diff --git a/python/freetoken/models/gguf/reader.py b/python/freetoken/models/gguf/reader.py index b950d929..449871fb 100644 --- a/python/freetoken/models/gguf/reader.py +++ b/python/freetoken/models/gguf/reader.py @@ -6,12 +6,42 @@ reversed), the ggml quant type, and a zero-copy ``uint8`` view of the packed block bytes laid out as ``[rows, row_bytes]`` (rows = product of all but the fastest ggml dim; row_bytes spans whole quant blocks of the fastest dim). + +Multi-shard GGUF support (llama.cpp split convention, per ground truth in spec): + +Filenames follow ``-%05d-of-%05d.gguf``, both numbers 1-based. The split layout is: + + - Shard 1 (``-00001-of-000NN``) holds the FULL KV metadata: ``general.architecture``, + all ``.*`` config keys, all ``tokenizer.*`` keys. It also carries tensor count + and a ``split.no = 0`` marker (0-based, even though filenames are 1-based). + - Shards 2..N (``-00002-of-000NN`` to ``-000NN-of-000NN``) carry exactly 3 metadata keys: + ``split.no`` (1..N-1), ``split.count`` (always NN), and ``split.tensors.count`` (the + TOTAL tensor count across all shards, not per-shard). They list no architecture keys. + - Tensor distribution: e.g. Hy3 IQ1_M (1298 total) splits as shard 1 with 694 tensors, + shard 2 with 604 tensors. ``split.tensors.count`` is always 1298. + +Examples: + + - A bare ``.gguf`` file (no shard marker) -> no change, single-file path throughout. + - ``model-00001-of-00002.gguf`` passed to any function -> caller gets shard 1 metadata + and tensors from both shards 1..2 in order. + - A directory containing ``model-00001-of-00002.gguf`` -> caller passes the dir, + is_gguf_path resolves it, downstream gets the first-shard path. + +Validation on open: if shard 1 declares ``split.count = N``, all N shards must exist +(1..N), no gaps. Also assert summed tensor count across all shards equals +``split.tensors.count``. Both keys are read from shard 1 only. + +Shard readers are cached per path (one per shard file), so opening ``-00002-of-00002`` +after ``-00001-of-00002`` will reuse the first-shard reader (no double-load of shard 1). """ from __future__ import annotations import functools +import glob import os +import re import struct from dataclasses import dataclass from typing import Any, Iterator @@ -20,11 +50,127 @@ import torch +def gguf_shards(path: str) -> list[str]: + r"""Return the ordered list of shard paths given any shard's path (or a plain .gguf). + + Matches the llama.cpp pattern ``(?P.+)-(\d{5})-of-(\d{5})\.gguf$`` on the + basename. A non-shard path returns ``[path]``. For a shard path, globs sibling shards, + sorts by index, and validates the set is complete 1..N with none missing. + + Raises a clear error naming the missing indices if any are absent (truncated downloads + are the common failure case and must not load silently). + """ + # A directory: find the first shard inside it and continue from there. Users routinely + # pass the folder a split model was downloaded into rather than a specific shard. + if os.path.isdir(path): + first = sorted(glob.glob(os.path.join(path, "*-00001-of-?????.gguf"))) + if len(first) > 1: + raise ValueError( + f"{path}: contains {len(first)} different split models " + f"({[os.path.basename(f) for f in first]}); point at one shard instead" + ) + if not first: + return [path] + path = first[0] + + basename = os.path.basename(path) + match = re.match(r"(?P.+)-(\d{5})-of-(\d{5})\.gguf$", basename) + if not match: + # Not a shard file; return as single-file path. + return [path] + + base, shard_idx_str, total_shards_str = match.group("base"), match.group(2), match.group(3) + total_shards = int(total_shards_str) + shard_dir = os.path.dirname(path) + + # Glob all sibling shards + pattern = os.path.join(shard_dir, f"{base}-?????-of-{total_shards_str}.gguf") + found_shards = sorted(glob.glob(pattern)) + + # Parse indices and validate completeness + shard_indices = set() + shard_map = {} # index -> path + for shard_path in found_shards: + shard_basename = os.path.basename(shard_path) + shard_match = re.match(rf"{re.escape(base)}-(\d{{5}})-of-{total_shards_str}\.gguf$", shard_basename) + if shard_match: + idx = int(shard_match.group(1)) + shard_indices.add(idx) + shard_map[idx] = shard_path + + # Verify complete range 1..N + expected = set(range(1, total_shards + 1)) + if shard_indices != expected: + missing = sorted(expected - shard_indices) + raise ValueError( + f"Incomplete shard set for {base}: expected shards 1..{total_shards}, " + f"missing {missing}. (Truncated download?)" + ) + + # Return in order 1..N + return [shard_map[i] for i in range(1, total_shards + 1)] + + +def resolve_gguf_path(model_path: str) -> str | None: + """Resolve a path to the first shard (shard 1) of a GGUF file. + + Accepts: + - A single ``.gguf`` file -> returns it as-is. + - A shard file (e.g., ``-00002-of-00002.gguf``) -> returns shard 1 path. + - A directory containing exactly one shard 1 file -> returns that file path. + + Returns ``None`` if the path is none of the above. + """ + if not isinstance(model_path, str): + return None + + # Case 1: A single .gguf file (not a shard) + if os.path.isfile(model_path) and model_path.endswith(".gguf"): + basename = os.path.basename(model_path) + if not re.match(r".+-\d{5}-of-\d{5}\.gguf$", basename): + # Plain .gguf, not a shard + return model_path + + # Case 2: A shard file or a directory + if os.path.isfile(model_path) and model_path.endswith(".gguf"): + # It's a shard file; get shard 1 + shards = gguf_shards(model_path) + return shards[0] if shards else None + + if os.path.isdir(model_path): + # Look for exactly one shard-1 file in the directory + pattern = os.path.join(model_path, "*-00001-of-?????.gguf") + candidates = glob.glob(pattern) + if len(candidates) == 1: + return candidates[0] + + return None + + def is_gguf_path(model_path: str) -> bool: - """A single ``.gguf`` file (the only GGUF layout FreeToken loads directly).""" - return isinstance(model_path, str) and os.path.isfile(model_path) and model_path.endswith( - ".gguf" - ) + """A ``.gguf`` file or directory, supporting single files and multi-shard layouts. + + Accepts: + - A single ``.gguf`` file. + - Any shard of a multi-shard ``.gguf`` (e.g., shard 2 of 5). + - A directory containing exactly one shard 1 file. + + Returns ``True`` only if one of these conditions holds. + """ + if not isinstance(model_path, str): + return False + + # Case 1: A .gguf file (single or shard) + if os.path.isfile(model_path) and model_path.endswith(".gguf"): + return True + + # Case 2: A directory with a shard 1 file + if os.path.isdir(model_path): + pattern = os.path.join(model_path, "*-00001-of-?????.gguf") + candidates = glob.glob(pattern) + return len(candidates) == 1 + + return False # Canonical name of the metadata-only GGUF that ``convert_checkpoint`` drops into an FTW @@ -42,18 +188,24 @@ def is_gguf_path(model_path: str) -> bool: def gguf_config_source(model_path: str) -> str | None: """The ``.gguf`` file to source config/tokenizer/metadata from, or ``None``. - A bare ``.gguf`` file resolves to itself; an FTW dir carrying a - :data:`FTW_METADATA_GGUF` resolves to that embedded metadata file. This is the single - seam config/tokenizer dispatch uses to decide "this checkpoint is GGUF-config-sourced" - -- a real file and a converted-FTW dir both land on a genuine ``.gguf`` path the reader - can parse, so no downstream code learns about the FTW wrapper. + A bare ``.gguf`` file or any shard resolves to the first shard; an FTW dir carrying + a :data:`FTW_METADATA_GGUF` resolves to that embedded metadata file. A directory + containing shards resolves to shard 1. This is the single seam config/tokenizer + dispatch uses to decide "this checkpoint is GGUF-config-sourced" -- a real file, a + shard file, a shard directory, and a converted-FTW dir all land on a genuine ``.gguf`` + path the reader can parse, so no downstream code learns about the layout. """ - if is_gguf_path(model_path): - return model_path + # Case 1: Check for FTW metadata file first (highest priority) if isinstance(model_path, str) and os.path.isdir(model_path): cand = os.path.join(model_path, FTW_METADATA_GGUF) if os.path.isfile(cand): return cand + + # Case 2: Try to resolve to a GGUF (single, shard, or directory) + resolved = resolve_gguf_path(model_path) + if resolved is not None: + return resolved + return None @@ -121,61 +273,147 @@ def _field_value(reader, name: str) -> Any: @functools.cache def _reader(model_path: str): + """Get or create a GGUFReader for the given path, with shard validation. + + For single-shard files, this is a pass-through. For shard 1 of a multi-shard set, + this validates that: + 1. All shards 1..N are present and complete (no missing indices). + 2. The summed tensor count across all shards matches split.tensors.count (if present). + """ import gguf - return gguf.GGUFReader(model_path) + reader = gguf.GGUFReader(model_path) + + # Check if this is shard 1 of a multi-shard set + split_count = _field_value(reader, "split.count") + split_no = _field_value(reader, "split.no") + + if split_count is not None and split_no == 0: + # This is shard 1 of a multi-shard set; validate completeness + try: + shards = gguf_shards(model_path) + if len(shards) != split_count: + raise ValueError( + f"GGUF shard validation: {model_path} declares split.count={split_count}, " + f"but found {len(shards)} shards" + ) + + # Validate tensor count sum if split.tensors.count is declared + split_tensors_count = _field_value(reader, "split.tensors.count") + if split_tensors_count is not None: + total_tensor_count = 0 + for shard_path in shards: + shard_reader = gguf.GGUFReader(shard_path) + total_tensor_count += len(shard_reader.tensors) + + if total_tensor_count != split_tensors_count: + raise ValueError( + f"GGUF shard validation: {model_path} declares " + f"split.tensors.count={split_tensors_count}, but summed tensor count " + f"across all shards is {total_tensor_count}" + ) + except ValueError: + raise + + return reader @functools.cache def load_gguf_metadata(model_path: str) -> dict[str, Any]: - """All GGUF KV metadata as ``{field_name: python_value}`` (arrays -> lists).""" - reader = _reader(model_path) + """All GGUF KV metadata as ``{field_name: python_value}`` (arrays -> lists). + + Metadata is read from shard 1 only. If the caller passes any other shard, this + function resolves to shard 1 first. + """ + shard1_path = resolve_gguf_path(model_path) + if shard1_path is None: + raise ValueError(f"Cannot resolve GGUF path: {model_path}") + + reader = _reader(shard1_path) return {name: field.contents() for name, field in reader.fields.items()} def gguf_architecture(model_path: str) -> str: - arch = _field_value(_reader(model_path), "general.architecture") + """The model architecture string (e.g., "qwen3moe", "qwen35moe"). + + Architecture is read from shard 1 only. If the caller passes any other shard, + this function resolves to shard 1 first. + """ + shard1_path = resolve_gguf_path(model_path) + if shard1_path is None: + raise ValueError(f"Cannot resolve GGUF path: {model_path}") + + arch = _field_value(_reader(shard1_path), "general.architecture") if arch is None: - raise ValueError(f"GGUF file {model_path} has no general.architecture") + raise ValueError(f"GGUF file {shard1_path} has no general.architecture") return str(arch) def iter_gguf_tensors(model_path: str) -> Iterator[GgufTensor]: - """Yield every tensor with its torch shape, ggml type, and packed block bytes.""" + """Yield every tensor with its torch shape, ggml type, and packed block bytes. + + For multi-shard files, yields tensors from shard 1, then shard 2, ..., in order. + Single-shard files take exactly the same code path (gguf_shards returns [path]). + """ import gguf - reader = _reader(model_path) - for t in reader.tensors: - ne = [int(s) for s in t.shape] # ggml order, fastest dim first - torch_shape = tuple(reversed(ne)) - block, type_size = gguf.GGML_QUANT_SIZES[t.tensor_type] - n_fast = ne[0] - if n_fast % block != 0: - raise ValueError( - f"{t.name}: fastest dim {n_fast} not a multiple of block {block} " - f"for {t.tensor_type.name}" + shard1_path = resolve_gguf_path(model_path) + if shard1_path is None: + raise ValueError(f"Cannot resolve GGUF path: {model_path}") + + # Get all shard paths in order + shards = gguf_shards(shard1_path) + + # Iterate over each shard and yield tensors + for shard_path in shards: + reader = _reader(shard_path) + for t in reader.tensors: + ne = [int(s) for s in t.shape] # ggml order, fastest dim first + torch_shape = tuple(reversed(ne)) + block, type_size = gguf.GGML_QUANT_SIZES[t.tensor_type] + n_fast = ne[0] + if n_fast % block != 0: + raise ValueError( + f"{t.name}: fastest dim {n_fast} not a multiple of block {block} " + f"for {t.tensor_type.name}" + ) + row_bytes = n_fast // block * type_size + rows = int(np.prod(ne[1:])) if len(ne) > 1 else 1 + # gguf-py returns quantized tensors as raw uint8 but F32/F16 as typed arrays; + # normalize everything to a flat byte view before shaping into [rows, row_bytes]. + flat = np.ascontiguousarray(t.data).reshape(-1).view(np.uint8) + raw = flat.reshape(rows, row_bytes) + yield GgufTensor( + name=t.name, + shape=torch_shape, + ggml_type=int(t.tensor_type), + rows=rows, + row_bytes=row_bytes, + _raw=raw, ) - row_bytes = n_fast // block * type_size - rows = int(np.prod(ne[1:])) if len(ne) > 1 else 1 - # gguf-py returns quantized tensors as raw uint8 but F32/F16 as typed arrays; - # normalize everything to a flat byte view before shaping into [rows, row_bytes]. - flat = np.ascontiguousarray(t.data).reshape(-1).view(np.uint8) - raw = flat.reshape(rows, row_bytes) - yield GgufTensor( - name=t.name, - shape=torch_shape, - ggml_type=int(t.tensor_type), - rows=rows, - row_bytes=row_bytes, - _raw=raw, - ) def gguf_tensor_names(model_path: str) -> set[str]: - return {t.name for t in _reader(model_path).tensors} + """The union of tensor names across all shards. + + For multi-shard files, returns the union of tensor names from shard 1, shard 2, etc. + Single-shard files take exactly the same code path. + """ + shard1_path = resolve_gguf_path(model_path) + if shard1_path is None: + raise ValueError(f"Cannot resolve GGUF path: {model_path}") + + shards = gguf_shards(shard1_path) + names = set() + for shard_path in shards: + reader = _reader(shard_path) + names.update(t.name for t in reader.tensors) + return names __all__ = [ + "gguf_shards", + "resolve_gguf_path", "is_gguf_path", "FTW_METADATA_GGUF", "OUTPUT_WEIGHT_PRESENT_KV", diff --git a/tests/models/test_gguf_shards.py b/tests/models/test_gguf_shards.py new file mode 100644 index 00000000..1bef402f --- /dev/null +++ b/tests/models/test_gguf_shards.py @@ -0,0 +1,200 @@ +"""Multi-shard GGUF reading: discovery, shard-1 metadata, tensor aggregation. + +Large GGUF checkpoints ship split (``-00001-of-000NN``), and llama.cpp's convention has +three properties that are easy to get wrong and nearly invisible when you do: + +* ``split.no`` is 0-BASED while the filenames are 1-BASED. Shard ``-00002-of-00003`` + carries ``split.no = 1``. +* ``split.tensors.count`` is the TOTAL across every shard, not this shard's count. +* Only shard 1 carries the real metadata. Later shards hold exactly three ``split.*`` keys + and no ``general.architecture`` at all, so anything that reads arch or tokenizer from an + arbitrary shard gets nothing. + +The fixtures below write GGUF bytes directly rather than going through ``gguf.GGUFWriter``: +the format is small, and hand-writing it keeps these tests independent of that writer's +API (which takes a required ``arch`` positional and has moved around between releases). +Layout per the spec: magic, uint32 version, uint64 tensor_count, uint64 kv_count, the KV +pairs, the tensor infos, then padding to ``general.alignment`` and the tensor data. +""" + +from __future__ import annotations + +import struct +from pathlib import Path + +import pytest + +from freetoken.models.gguf.reader import ( + gguf_architecture, + gguf_shards, + gguf_tensor_names, + is_gguf_path, + iter_gguf_tensors, + load_gguf_metadata, +) + +# GGUF value type tags +_UINT32, _UINT64, _STRING = 4, 10, 8 +_F32_TENSOR_TYPE = 0 +_ALIGN = 32 + + +def _u32(v: int) -> bytes: + return struct.pack(" bytes: + return struct.pack(" bytes: + raw = s.encode("utf-8") + return _u64(len(raw)) + raw + + +def _kv(key: str, tag: int, value) -> bytes: + out = _string(key) + _u32(tag) + if tag == _STRING: + return out + _string(value) + if tag == _UINT32: + return out + _u32(value) + if tag == _UINT64: + return out + _u64(value) + raise AssertionError(f"unhandled tag {tag}") + + +def _write_gguf(path: Path, kvs: list[bytes], tensors: list[tuple[str, int]]) -> None: + """Write a GGUF with ``tensors`` as [(name, n_elements)], each F32. + + Tensor data is written contiguously after the aligned header; the values themselves are + irrelevant here since these tests only exercise discovery, metadata and the tensor + table. + """ + head = b"GGUF" + _u32(3) + _u64(len(tensors)) + _u64(len(kvs)) + head += b"".join(kvs) + offset = 0 + infos = b"" + for name, n in tensors: + infos += _string(name) + _u32(1) + _u64(n) + _u32(_F32_TENSOR_TYPE) + _u64(offset) + nbytes = n * 4 + offset += (nbytes + _ALIGN - 1) // _ALIGN * _ALIGN + body = head + infos + pad = (-len(body)) % _ALIGN + body += b"\0" * pad + body += b"\0" * offset + path.write_bytes(body) + + +def _full_kvs(arch: str = "qwen3moe", *, extra: list[bytes] | None = None) -> list[bytes]: + """Shard 1's KV block: the real metadata.""" + kvs = [ + _kv("general.architecture", _STRING, arch), + _kv("general.alignment", _UINT32, _ALIGN), + _kv(f"{arch}.block_count", _UINT32, 4), + _kv(f"{arch}.embedding_length", _UINT32, 128), + ] + return kvs + (extra or []) + + +def _split_kvs(no: int, count: int, total_tensors: int) -> list[bytes]: + """A non-first shard's KV block: exactly the three split keys, no architecture.""" + return [ + _kv("split.no", _UINT32, no), + _kv("split.count", _UINT32, count), + _kv("split.tensors.count", _UINT32, total_tensors), + ] + + +def _make_split(tmp_path: Path, base: str, per_shard: list[list[str]], *, + declared_count: int | None = None) -> list[Path]: + """Write a split set; returns the shard paths in order.""" + n = len(per_shard) + declared = declared_count if declared_count is not None else n + total = sum(len(names) for names in per_shard) + paths = [] + for i, names in enumerate(per_shard): + p = tmp_path / f"{base}-{i + 1:05d}-of-{n:05d}.gguf" + kvs = (_full_kvs() + _split_kvs(0, declared, total)) if i == 0 \ + else _split_kvs(i, declared, total) + _write_gguf(p, kvs, [(nm, 8) for nm in names]) + paths.append(p) + return paths + + +class TestSingleFileUnchanged: + def test_single_file_unchanged(self, tmp_path: Path): + """A plain one-file GGUF must behave exactly as before the shard work.""" + p = tmp_path / "single.gguf" + _write_gguf(p, _full_kvs(), [("token_embd.weight", 8), ("output.weight", 8)]) + assert is_gguf_path(str(p)) + assert gguf_shards(str(p)) == [str(p)] + assert gguf_architecture(str(p)) == "qwen3moe" + assert load_gguf_metadata(str(p))["qwen3moe.block_count"] == 4 + assert gguf_tensor_names(str(p)) == {"token_embd.weight", "output.weight"} + assert len(list(iter_gguf_tensors(str(p)))) == 2 + + +class TestShardDiscovery: + def test_discovery_from_any_shard_or_directory(self, tmp_path: Path): + paths = _make_split(tmp_path, "m", [["a.weight"], ["b.weight"], ["c.weight"]]) + want = [str(p) for p in paths] + for handed in (paths[0], paths[1], paths[2]): + assert gguf_shards(str(handed)) == want, f"from {handed.name}" + # a user pointing at the folder must work too + assert gguf_shards(str(tmp_path)) == want + + def test_every_shard_is_a_gguf_path(self, tmp_path: Path): + paths = _make_split(tmp_path, "m", [["a.weight"], ["b.weight"]]) + for p in paths: + assert is_gguf_path(str(p)) + + +class TestMissingShard: + def test_missing_shard_raises_naming_the_index(self, tmp_path: Path): + """A truncated download must fail loudly, never load as a partial model.""" + paths = _make_split(tmp_path, "m", [["a.weight"], ["b.weight"], ["c.weight"]]) + paths[1].unlink() # drop the middle shard + with pytest.raises(Exception) as e: + gguf_shards(str(paths[0])) + assert "2" in str(e.value), f"error should name the missing index: {e.value}" + + +class TestMetadataFromShardOne: + def test_metadata_resolves_to_shard_one(self, tmp_path: Path): + """Later shards carry no architecture, so reads must resolve back to shard 1.""" + paths = _make_split(tmp_path, "m", [["a.weight"], ["b.weight"], ["c.weight"]]) + for p in paths: + assert gguf_architecture(str(p)) == "qwen3moe", f"from {p.name}" + assert load_gguf_metadata(str(p))["qwen3moe.block_count"] == 4, f"from {p.name}" + + def test_later_shards_really_lack_arch(self, tmp_path: Path): + """Guards the fixture itself: if shard 2 carried arch, the test above proves nothing.""" + paths = _make_split(tmp_path, "m", [["a.weight"], ["b.weight"]]) + import gguf as gguf_pkg + + r = gguf_pkg.GGUFReader(str(paths[1])) + assert "general.architecture" not in r.fields + assert "split.no" in r.fields + + +class TestTensorAggregation: + def test_tensors_aggregate_across_shards(self, tmp_path: Path): + per = [["a.weight", "b.weight"], ["c.weight"], ["d.weight", "e.weight"]] + paths = _make_split(tmp_path, "m", per) + expected = {n for names in per for n in names} + for p in paths: + assert gguf_tensor_names(str(p)) == expected, f"from {p.name}" + names = [t.name for t in iter_gguf_tensors(str(paths[0]))] + assert names == [n for names_ in per for n in names_], "shard order must be preserved" + assert len(names) == load_gguf_metadata(str(paths[0]))["split.tensors.count"] + + +class TestDeclaredCountMismatch: + def test_declared_count_mismatch_raises(self, tmp_path: Path): + """shard 1 says 3 shards but only 2 exist on disk.""" + tmp = tmp_path / "sub" + tmp.mkdir() + _make_split(tmp, "m", [["a.weight"], ["b.weight"]], declared_count=3) + first = tmp / "m-00001-of-00002.gguf" + with pytest.raises(Exception): + list(iter_gguf_tensors(str(first)))