From 6ef02e0bc845a85cc3111e8e6bc4037d47078c7e Mon Sep 17 00:00:00 2001 From: Digant Desai Date: Mon, 26 Oct 2026 22:40:24 -0700 Subject: [PATCH] Initial export of the PTN samples stack [ghstack-poisoned] --- .../native/examples/hf_llm/export_hf_llm.py | 83 +++++++++++++++++- .../examples/hf_llm/test_export_hf_llm.py | 85 +++++++++++++++++++ 2 files changed, 165 insertions(+), 3 deletions(-) diff --git a/backends/native/examples/hf_llm/export_hf_llm.py b/backends/native/examples/hf_llm/export_hf_llm.py index 2df9aa84e8b..7fcb6d93ba8 100644 --- a/backends/native/examples/hf_llm/export_hf_llm.py +++ b/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([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, "forward", "golden", samples) + print(f"wrote {len(prompts)} golden samples") 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 766d33f7cdc..9c26f4ac271 100644 --- a/backends/native/examples/hf_llm/test_export_hf_llm.py +++ b/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,81 @@ 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.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)