[Native] Load the first layers of a GGUF with --layers

backends/native/examples/hf_llm/export_hf_llm.py

@@ -124,11 +124,15 @@ def load_gguf(
     path: str,
     architecture: str | None = None,
     raw_iter: RawIter | None = None,
+    truncated: bool = False,
 ) -> None:
     """Replace every parameter with a GGUF file's weights; quantized ones stay
     ``ExportableGGUFTensor`` so the PTN carries the GGUF blocks. ``architecture``
-    defaults to the file's ``general.architecture``."""
-    key_map, convert = gguf_to_hf(model, architecture or gguf_architecture(path))
+    defaults to the file's ``general.architecture``. ``truncated`` loads a model
+    built with fewer layers from the file's first blocks."""
+    key_map, convert = gguf_to_hf(
+        model, architecture or gguf_architecture(path), truncated
+    )
     shapes = {n: p.shape for n, p in model.named_parameters(remove_duplicate=False)}
     expected = set(shapes)
     loaded = set()
@@ -419,7 +423,12 @@ def _parse_args() -> argparse.Namespace:
         action="store_true",
         help="Use only the model's config, with random weights.",
     )
-    parser.add_argument("--layers", type=int, default=0)
+    parser.add_argument(
+        "--layers",
+        type=int,
+        default=0,
+        help="Keep only the first N layers; with --gguf, load their weights.",
+    )
     parser.add_argument("--max_seq_length", type=int, default=128)
     parser.add_argument("--max_context_length", type=int, default=0)
     parser.add_argument("--output", default="")
@@ -450,10 +459,9 @@ def _parse_args() -> argparse.Namespace:
     args = parser.parse_args()
     if not 0 <= args.samples <= len(_PROMPTS) or (args.samples and not args.output):
         parser.error(f"--samples takes 0 to {len(_PROMPTS)} and needs --output")
-    if args.gguf and (args.layers or args.random_weights or args.quantize != "none"):
+    if args.gguf and (args.random_weights or args.quantize != "none"):
         parser.error(
-            "--gguf loads every weight; it excludes --layers, --random_weights "
-            "and --quantize"
+            "--gguf loads every weight; it excludes --random_weights and --quantize"
         )
     if args.offgraph_cache and (
         args.max_context_length or args.kv_cache_dtype != "fp32"
@@ -491,7 +499,7 @@ def main() -> None:
         offgraph_cache=args.offgraph_cache,
     )
     if args.gguf:
-        load_gguf(model, args.gguf)
+        load_gguf(model, args.gguf, truncated=bool(args.layers))
     if args.quantize == "8da4w":
         _quantize_8da4w(model, args.group_size, args.embedding_group_size)
     if args.offgraph_cache:

backends/native/examples/hf_llm/gguf_hf.py

@@ -16,6 +16,7 @@
 
 from __future__ import annotations
 
+import re
 from typing import Callable
 
 import gguf
@@ -107,11 +108,15 @@ def convert(fqn: str, value: torch.Tensor) -> torch.Tensor:
 }
 
 
-def gguf_to_hf(model: torch.nn.Module, architecture: str) -> tuple[KeyMap, Convert]:
+def gguf_to_hf(
+    model: torch.nn.Module, architecture: str, truncated: bool = False
+) -> tuple[KeyMap, Convert]:
     """Return the ``key_map`` and ``convert`` that load a GGUF file of
     ``architecture`` into ``model``, a Hugging Face model.
 
-    The key map raises for a GGUF tensor that names no parameter of ``model``.
+    The key map raises for a GGUF tensor that names no parameter of ``model``,
+    unless ``truncated``: a model built with fewer layers than the file holds
+    then skips the blocks past its last layer.
     """
     if architecture not in _ARCHITECTURES:
         raise ValueError(
@@ -137,6 +142,9 @@ def key_map(name: str) -> str | None:
             return names[name]
         if name in _DERIVED:
             return None
+        block = re.match(r"blk\.(\d+)\.", name)
+        if truncated and block and int(block[1]) >= config.num_hidden_layers:
+            return None
         raise ValueError(
             f"GGUF tensor {name!r} maps to no {config.model_type} parameter"
         )

backends/native/examples/hf_llm/test_export_hf_llm.py

@@ -200,6 +200,23 @@ def test_gguf_loads_into_a_model_built_without_weights(self) -> None:
                     got, want = loaded(ids).logits, reference(ids).logits
                 self.assertTrue(torch.equal(got, want))
 
+    def test_gguf_loads_the_first_layers_into_a_truncated_model(self) -> None:
+        ids = torch.tensor([[3, 4, 5, 6, 7, 8]])
+        torch.manual_seed(1)
+        reference = _tiny("llama")
+        pairs = _as_gguf_tensors(reference, "llama")
+        with tempfile.TemporaryDirectory() as path:
+            reference.config.save_pretrained(path)
+            loaded = _load(path, 1, False, empty_weights=True)
+
+        load_gguf(loaded, "model.gguf", "llama", lambda _: iter(pairs), truncated=True)
+
+        reference.model.layers = reference.model.layers[:1]
+        reference.config.num_hidden_layers = 1
+        with torch.no_grad():
+            got, want = loaded(ids).logits, reference(ids).logits
+        self.assertTrue(torch.equal(got, want))
+
     def test_gguf_derives_head_dim_when_config_has_none(self) -> None:
         torch.manual_seed(1)
         reference = _tiny()

backends/native/examples/hf_llm/tests/test_gguf_hf.py

@@ -85,6 +85,16 @@ def test_skips_derived_and_rejects_unknown_tensors(self) -> None:
             with self.assertRaisesRegex(ValueError, "maps to no llama parameter"):
                 key_map(name)
 
+    def test_truncated_model_skips_later_blocks(self) -> None:
+        key_map, _ = gguf_to_hf(_tiny(LlamaConfig), "llama", truncated=True)
+
+        self.assertIsNone(key_map("blk.2.attn_q.weight"))
+        self.assertEqual(
+            key_map("blk.1.attn_q.weight"), "model.layers.1.self_attn.q_proj.weight"
+        )
+        with self.assertRaisesRegex(ValueError, "maps to no llama parameter"):
+            key_map("blk.0.attn_qkv.weight")
+
     def test_checks_the_architecture_against_the_model(self) -> None:
         gguf_to_hf(_tiny(MistralConfig), "llama")
         with self.assertRaisesRegex(ValueError, "cannot load a 'qwen2' model"):