[ET][native] Bundle LLM prefill and decode samples from export_hf_llm

backends/native/examples/hf_llm/export_hf_llm.py

@@ -12,7 +12,8 @@
 ``(input_ids [1, T], cache_position [T])`` signature. Reports which operators
 the native partitioner leaves behind. ``--gguf`` takes the weights from a GGUF
 file of an architecture ``gguf_hf`` supports instead, keeping its quantized blocks
-packed in the PTN.
+packed in the PTN. ``--samples`` appends a ``golden`` sample set: eager logits
+for a prefill and a few greedy decode steps per prompt.
 """
 
 import argparse
@@ -34,7 +35,7 @@
     executorch_call_delegate,
     get_lowered_submodules,
 )
-from executorch.exir.native import to_native
+from executorch.exir.native import Call, Sample, SampleSet, to_native, write_samples
 from executorch.extension.llm.export.load import assign_one, iter_checkpoint, RawIter
 from torch.export import Dim, export
 from torchao.quantization import (
@@ -43,10 +44,12 @@
     quantize_,
 )
 from torchao.quantization.granularity import PerGroup
+from torchao.utils import TorchAOBaseTensor
 from transformers import (
     AttentionInterface,
     AutoConfig,
     AutoModelForCausalLM,
+    AutoTokenizer,
     PreTrainedModel,
 )
 from transformers.cache_utils import StaticLayer
@@ -218,6 +221,55 @@ def _quantize_8da4w(
     )
 
 
+_PROMPTS = (
+    "The capital of France is",
+    "Once upon a time, in a small village by the sea, there lived a",
+    "def fibonacci(n):",
+    "The three primary colors are red,",
+    "Water boils at a temperature of",
+)
+_DECODE_STEPS = 4
+# Advisory (atol, rtol) for engine logits against eager ones. Engines quantize
+# activations their own way, not torchao's, so quantized weights get a looser bound.
+_FLOAT_TOL = (1e-3, 1e-3)
+_QUANTIZED_TOL = (2.0, 5e-2)
+
+
+def _sample_prompts(
+    model: str, count: int, max_seq_length: int, cache_len: int
+) -> list[list[int]]:
+    tokenizer = AutoTokenizer.from_pretrained(model)
+    prompts = [tokenizer.encode(p) for p in _PROMPTS[:count]]
+    longest = max(map(len, prompts))
+    if longest > max_seq_length - 1 or longest + _DECODE_STEPS > cache_len:
+        raise ValueError("a sample prompt and its decode steps do not fit")
+    return prompts
+
+
+def golden_samples(
+    wrapper: LastTokenLM, prompts: list[list[int]], decode_steps: int
+) -> SampleSet:
+    """Per tokenized prompt, from an empty cache: a prefill call, then
+    `decode_steps` calls each fed the previous call's eager argmax."""
+    quantized = any(isinstance(p, TorchAOBaseTensor) for p in wrapper.parameters())
+    samples = []
+    with torch.no_grad():
+        for ids in prompts:
+            wrapper.cache.reset()
+            input_ids = torch.tensor([ids], dtype=torch.long)
+            positions = torch.arange(len(ids), dtype=torch.long)
+            calls = []
+            for _ in range(decode_steps + 1):
+                logits = wrapper(input_ids, positions)
+                calls.append(Call("forward", [input_ids, positions], [logits]))
+                input_ids = logits[:, -1].argmax(dim=-1, keepdim=True)
+                positions = positions[-1:] + 1
+            samples.append(Sample(calls))
+    wrapper.cache.reset()
+    tol = _QUANTIZED_TOL if quantized else _FLOAT_TOL
+    return SampleSet(samples, source="eager/cpu", tol=tol)
+
+
 def _histogram(gm: torch.fx.GraphModule) -> collections.Counter:
     return collections.Counter(
         str(n.target)
@@ -227,7 +279,7 @@ def _histogram(gm: torch.fx.GraphModule) -> collections.Counter:
     )
 
 
-def main() -> None:
+def _parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "--model", required=True, help="Hugging Face model directory or hub ID."
@@ -251,7 +303,17 @@ def main() -> None:
         help="GGUF file with the model's weights; --model supplies the config.",
     )
     parser.add_argument("--kv_cache_dtype", choices=["fp32", "fp16"], default="fp32")
+    parser.add_argument(
+        "--samples",
+        type=int,
+        default=0,
+        help=f"Append eager logits for the first N of {len(_PROMPTS)} built-in "
+        f"prompts, each a prefill and {_DECODE_STEPS} decode calls, as forward's "
+        "'golden' sample set; needs --output and a tokenizer in --model.",
+    )
     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"):
         parser.error(
             "--gguf loads every weight; it excludes --layers, --random_weights "
@@ -264,6 +326,17 @@ def main() -> None:
     cache_len = args.max_context_length or args.max_seq_length
     if cache_len < args.max_seq_length - 1:
         parser.error("the context must hold a full prompt of max_seq_length - 1 tokens")
+    return args
+
+
+def main() -> None:
+    args = _parse_args()
+    cache_len = args.max_context_length or args.max_seq_length
+    prompts = (
+        _sample_prompts(args.model, args.samples, args.max_seq_length, cache_len)
+        if args.samples
+        else []
+    )
     logging.basicConfig(level=logging.INFO)
     torch.manual_seed(0)
 
@@ -314,3 +387,7 @@ def main() -> None:
     if args.output:
         to_native(ep).save(args.output)
         print(f"saved {args.output}")
+    if prompts:
+        samples = golden_samples(wrapper, prompts, _DECODE_STEPS)
+        write_samples(args.output, "golden", samples)
+        print(f"wrote {len(prompts)} golden samples")

backends/native/examples/hf_llm/test_export_hf_llm.py

@@ -7,13 +7,23 @@
 import collections
 import copy
 import re
+import sys
 import tempfile
 import unittest
+from unittest import mock
 
 import gguf
 import torch
+from executorch.backends.native.examples.hf_llm import export_hf_llm
 from executorch.backends.native.examples.hf_llm.export_hf_llm import (
+    _DECODE_STEPS,
+    _FLOAT_TOL,
     _load,
+    _PROMPTS,
+    _quantize_8da4w,
+    _QUANTIZED_TOL,
+    _sample_prompts,
+    golden_samples,
     LastTokenLM,
     load_gguf,
 )
@@ -230,6 +240,82 @@ def test_fp16_kv_cache_exports_position_explicit_graph(self) -> None:
         self.assertEqual({t.dtype for t in caches.values()}, {torch.float16})
         self.assertEqual(_ops_except_casts(ep), _ops_except_casts(ep32))
 
+    def test_golden_samples_teacher_force_eager_argmax(self) -> None:
+        torch.manual_seed(7)
+        model = _tiny()
+        prompts = [[3, 4, 5, 6, 7, 8], [9, 10]]
+        wrapper = LastTokenLM(model, 32).eval()
+
+        samples = golden_samples(wrapper, prompts, 2)
+
+        ref = LastTokenLM(copy.deepcopy(model), 32).eval()
+        for ids, sample in zip(prompts, samples.samples, strict=True):
+            ref.cache.reset()
+            want = [
+                (torch.tensor([ids]), torch.arange(len(ids))),
+                (None, torch.tensor([len(ids)])),
+                (None, torch.tensor([len(ids) + 1])),
+            ]
+            self.assertEqual(len(sample.calls), len(want))
+            previous = None
+            for call, (input_ids, positions) in zip(sample.calls, want):
+                if input_ids is None:
+                    input_ids = previous.argmax(dim=-1)
+                self.assertEqual(call.method, "forward")
+                self.assertTrue(torch.equal(call.inputs[0], input_ids))
+                self.assertTrue(torch.equal(call.inputs[1], positions))
+                with torch.no_grad():
+                    previous = ref(input_ids, positions)
+                self.assertTrue(torch.equal(call.outputs[0], previous))
+        self.assertFalse(any(layer.keys.any() for layer in wrapper.cache.layers))
+
+    def test_samples_flag_needs_a_count_in_range_and_an_output(self) -> None:
+        for flags in (
+            ["--samples", str(len(_PROMPTS) + 1), "--output", "m.ptn"],
+            ["--samples", "-1", "--output", "m.ptn"],
+            ["--samples", "1"],
+        ):
+            argv = ["export_hf_llm", "--model", "unused", *flags]
+            with self.subTest(flags=flags), mock.patch.object(sys, "argv", argv):
+                with mock.patch.object(export_hf_llm, "_load") as load:
+                    with self.assertRaises(SystemExit):
+                        export_hf_llm.main()
+                load.assert_not_called()
+
+    def test_sample_prompts_must_fit_with_their_decode_steps(self) -> None:
+        tokenizer = mock.Mock()
+        tokenizer.encode.side_effect = lambda text: list(range(len(text.split())))
+        lengths = [len(p.split()) for p in _PROMPTS[:2]]
+        longest = max(lengths)
+
+        with mock.patch.object(
+            export_hf_llm.AutoTokenizer, "from_pretrained", return_value=tokenizer
+        ):
+            prompts = _sample_prompts("m", 2, longest + 1, longest + _DECODE_STEPS)
+            self.assertEqual(prompts, [list(range(n)) for n in lengths])
+            for max_seq_length, cache_len in (
+                (longest, 64),
+                (64, longest + _DECODE_STEPS - 1),
+            ):
+                with self.assertRaisesRegex(ValueError, "do not fit"):
+                    _sample_prompts("m", 2, max_seq_length, cache_len)
+
+    def test_golden_samples_tolerance_follows_weight_quantization(self) -> None:
+        source = _tiny()
+        gguf_float = _tiny()
+        pairs = _as_gguf_tensors(source, "llama")
+        load_gguf(gguf_float, "model.gguf", "llama", lambda _: iter(pairs))
+        quantized = _tiny()
+        _quantize_8da4w(quantized, group_size=32, embedding_group_size=32)
+
+        for model, tol in (
+            (source, _FLOAT_TOL),
+            (gguf_float, _FLOAT_TOL),
+            (quantized, _QUANTIZED_TOL),
+        ):
+            samples = golden_samples(LastTokenLM(model, 16).eval(), [[3, 4]], 0)
+            self.assertEqual(samples.tol, tol)
+
     def test_fp16_kv_cache_rejects_sliding_window_layers(self) -> None:
         with self.assertRaisesRegex(ValueError, "cannot change the KV dtype"):
             LastTokenLM(_tiny("gemma3"), 32, torch.float16)