studio: group GGUF shards by variant in size-based fallback

The smallest-fitting-variant fallback now groups split GGUF shards
by their variant prefix and sums all shard sizes per variant.

For example, DeepSeek-V3.2 UD-Q4_K_XL has 9 shards totaling
379.8 GB. The previous code treated each shard as a separate
"variant" and would have incorrectly selected a single 50 GB shard
as fitting, ignoring the other 8 shards needed.

Tested with unsloth/DeepSeek-V3.2-GGUF (237 GGUF files, 27
variants from 150 GB to 1.25 TB). Correctly groups and sorts
all variants by total size.
This commit is contained in:
Daniel Han 2026-03-14 09:01:30 +00:00
commit 8ccb461570

View file

@ -287,10 +287,15 @@ class LlamaCppBackend:
free_bytes: int,
hf_token: Optional[str] = None,
) -> Optional[tuple[str, int]]:
"""Find the smallest single-file GGUF variant that fits in free_bytes.
"""Find the smallest GGUF variant (including all shards) that fits.
Returns (filename, size_bytes) or None if nothing fits.
Groups split shards by variant prefix and sums their sizes.
For example, UD-Q4_K_XL with 9 shards of 50 GB each = 450 GB total.
Returns (first_shard_filename, total_size_bytes) or None if nothing fits.
"""
import re
try:
from huggingface_hub import get_paths_info, list_repo_files
@ -300,16 +305,31 @@ class LlamaCppBackend:
return None
# Get sizes for all GGUF files
path_infos = list(get_paths_info(hf_repo, gguf_files, token = hf_token))
sized = [(p.path, p.size) for p in path_infos if p.size and p.size > 0]
if not sized:
return None
path_infos = list(
get_paths_info(hf_repo, gguf_files, token = hf_token)
)
size_map = {p.path: (p.size or 0) for p in path_infos}
# Sort by size ascending and pick the smallest that fits
sized.sort(key = lambda x: x[1])
for filename, size in sized:
if size <= free_bytes:
return filename, size
# Group files by variant: shards share a prefix before -NNNNN-of-NNNNN
shard_pat = re.compile(r"^(.*)-\d{5}-of-\d{5}\.gguf$")
variants: dict[str, list[str]] = {}
for f in gguf_files:
m = shard_pat.match(f)
key = m.group(1) if m else f
variants.setdefault(key, []).append(f)
# Sum shard sizes per variant, track the first shard (for download)
variant_sizes: list[tuple[str, int, list[str]]] = []
for key, shard_files in variants.items():
total = sum(size_map.get(f, 0) for f in shard_files)
first = sorted(shard_files)[0]
variant_sizes.append((first, total, shard_files))
# Sort by total size ascending and pick the smallest that fits
variant_sizes.sort(key = lambda x: x[1])
for first_file, total_size, _ in variant_sizes:
if total_size > 0 and total_size <= free_bytes:
return first_file, total_size
return None
except Exception: