GGUF: read multi-shard checkpoints

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 ``<base>-%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 ``<arch>.*`` 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<base>.+)-(\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<base>.+)-(\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",