From dbcb2e1d1bc4015b854e96f22d9d60d8395b77c6 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Wed, 11 Nov 2026 14:04:59 -0800 Subject: [PATCH] Load the first layers of a GGUF with --layers --gguf excluded --layers because every GGUF tensor must name a parameter, and a truncated model has none for the later blocks. gguf_to_hf and load_gguf now take truncated=True to skip the blocks past the model's last layer, and the export passes it when --layers is set. That exports a few layers of a large GGUF checkpoint to test the lowering quickly. Without the flag, a tensor that names no parameter still raises. --- .../native/examples/hf_llm/export_hf_llm.py | 22 +++++++++++++------ backends/native/examples/hf_llm/gguf_hf.py | 12 ++++++++-- .../examples/hf_llm/test_export_hf_llm.py | 17 ++++++++++++++ .../examples/hf_llm/tests/test_gguf_hf.py | 10 +++++++++ 4 files changed, 52 insertions(+), 9 deletions(-) diff --git a/backends/native/examples/hf_llm/export_hf_llm.py b/backends/native/examples/hf_llm/export_hf_llm.py index 92879323cda..bda253bfec3 100644 --- a/backends/native/examples/hf_llm/export_hf_llm.py +++ b/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: diff --git a/backends/native/examples/hf_llm/gguf_hf.py b/backends/native/examples/hf_llm/gguf_hf.py index ab16d562069..e4bda9d72c7 100644 --- a/backends/native/examples/hf_llm/gguf_hf.py +++ b/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" ) diff --git a/backends/native/examples/hf_llm/test_export_hf_llm.py b/backends/native/examples/hf_llm/test_export_hf_llm.py index 35e8ef6e715..9c20b97ec6c 100644 --- a/backends/native/examples/hf_llm/test_export_hf_llm.py +++ b/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() diff --git a/backends/native/examples/hf_llm/tests/test_gguf_hf.py b/backends/native/examples/hf_llm/tests/test_gguf_hf.py index 96ea7b1edf3..b836020b782 100644 --- a/backends/native/examples/hf_llm/tests/test_gguf_hf.py +++ b/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"):