From ede5d6cd531853e670b731f47e472c9d37654ec1 Mon Sep 17 00:00:00 2001 From: Digant Desai Date: Mon, 26 Oct 2026 22:40:29 -0700 Subject: [PATCH] Initial export of the PTN samples stack [ghstack-poisoned] --- backends/native/examples/hf_llm/BUCK | 15 ++ backends/native/examples/hf_llm/runtime/BUCK | 2 + .../examples/hf_llm/runtime/LlmRunner.cpp | 167 +++++++++++++++++- .../hf_llm/runtime/LlmRunnerUtils.cpp | 26 +++ .../examples/hf_llm/runtime/LlmRunnerUtils.h | 15 ++ .../hf_llm/runtime/LlmRunnerUtilsTest.cpp | 42 +++++ .../examples/hf_llm/test_verify_samples.py | 158 +++++++++++++++++ 7 files changed, 423 insertions(+), 2 deletions(-) create mode 100644 backends/native/examples/hf_llm/test_verify_samples.py diff --git a/backends/native/examples/hf_llm/BUCK b/backends/native/examples/hf_llm/BUCK index 37bdeb93fed..535e80081a6 100644 --- a/backends/native/examples/hf_llm/BUCK +++ b/backends/native/examples/hf_llm/BUCK @@ -39,6 +39,21 @@ fbcode_target( ], ) +fbcode_target( + _kind = runtime.python_test, + name = "test_verify_samples", + srcs = ["test_verify_samples.py"], + env = { + "RUN_HF_LLM": "$(location //executorch/backends/native/examples/hf_llm/runtime:run_hf_llm_ggml_cpu)", + }, + deps = [ + ":export_hf_llm_lib", + "//caffe2:torch", + "//executorch/exir/native:native", + "fbsource//third-party/pypi/transformers:transformers", + ], +) + fbcode_target( _kind = runtime.python_binary, name = "convert_meta_llama", diff --git a/backends/native/examples/hf_llm/runtime/BUCK b/backends/native/examples/hf_llm/runtime/BUCK index ac110c40504..8bfa0ca1041 100644 --- a/backends/native/examples/hf_llm/runtime/BUCK +++ b/backends/native/examples/hf_llm/runtime/BUCK @@ -30,6 +30,8 @@ runtime.cxx_library( "//executorch/backends/native/examples/runner_utils:runner_utils", "//executorch/backends/native/runtime:runtime", "//executorch/backends/native/runtime/deserialize:package", + "//executorch/backends/native/runtime/deserialize:samples", + "//executorch/backends/native/runtime/deserialize:zip_reader", "//executorch/backends/native/runtime/graph:scalar_type", "//executorch/examples/models/llama/tokenizer:tiktoken", ], diff --git a/backends/native/examples/hf_llm/runtime/LlmRunner.cpp b/backends/native/examples/hf_llm/runtime/LlmRunner.cpp index 0c3a5903d54..81ce7a3cde7 100644 --- a/backends/native/examples/hf_llm/runtime/LlmRunner.cpp +++ b/backends/native/examples/hf_llm/runtime/LlmRunner.cpp @@ -8,6 +8,7 @@ #include #include #include +#include #include #include #include @@ -24,6 +25,8 @@ #include #include #include +#include +#include #include #include #include @@ -54,6 +57,11 @@ DEFINE_bool( warmup_prefill, false, "Replay the identical prompt once before timing prefill (benchmark only)."); +DEFINE_string( + verify_samples, + "", + "Replay this sample set of --method instead of --prompt, comparing each " + "call's logits and forcing the sample's next token; exits 1 on mismatch."); namespace { using Clock = std::chrono::steady_clock; @@ -86,6 +94,27 @@ size_t kv_cache_length(const ptn::Method& method) { return length; } +// The values of an [n] or [1, n] Int or Long sample tensor. +std::vector integer_values(const ptn::SampleTensor& tensor) { + const bool row = tensor.sizes.size() == 1 || + (tensor.sizes.size() == 2 && tensor.sizes[0] == 1); + if ((tensor.dtype != ptn::kLong && tensor.dtype != ptn::kInt) || !row) { + throw std::runtime_error( + "sample token and position inputs must be [n] or [1, n] Int or Long"); + } + std::vector values(static_cast(tensor.sizes.back())); + for (size_t i = 0; i < values.size(); ++i) { + if (tensor.dtype == ptn::kLong) { + std::memcpy(&values[i], tensor.bytes.data() + 8 * i, 8); + } else { + int32_t value = 0; + std::memcpy(&value, tensor.bytes.data() + 4 * i, 4); + values[i] = value; + } + } + return values; +} + // Chooses the token after a logits output whose last dimension is the // vocabulary. using NextToken = std::function< @@ -317,11 +346,11 @@ int run_llm(const std::function()>& make_host) { method.graph.output_ids.size(), user_outputs); const std::unique_ptr host = make_host(); - const std::unique_ptr context = + std::unique_ptr context = host->create_context(program, package); const Clock::time_point load_end = Clock::now(); const Clock::time_point compile_start = Clock::now(); - const std::unique_ptr executable = + std::unique_ptr executable = context->compile(method_name); const Clock::time_point compile_end = Clock::now(); @@ -375,6 +404,140 @@ int run_llm(const std::function()>& make_host) { "logits must be Float or Half, got " + std::string(ptn::scalar_type_name(output_dtype))); } + if (!FLAGS_verify_samples.empty()) { + if (FLAGS_warmup_prefill || !FLAGS_dump_first_logits.empty()) { + throw std::invalid_argument( + "verify_samples cannot be combined with warmup_prefill or " + "dump_first_logits"); + } + const ptn::SampleSet set = ptn::read_sample_set( + ptn::ZipReader::open(FLAGS_package), + method_name, + FLAGS_verify_samples); + size_t failed_calls = 0; + for (size_t i = 0; i < set.samples.size(); ++i) { + const std::vector& calls = set.samples[i].calls; + // Mutable state lives in the context, so a fresh one is the program's + // initial state. + if (i > 0) { + executable.reset(); + context.reset(); + context = host->create_context(program, package); + executable = context->compile(method_name); + } + std::vector prompt; + std::vector forced; + for (size_t j = 0; j < calls.size(); ++j) { + if (calls[j].inputs.size() != 2 || !calls[j].outputs || + calls[j].outputs->size() <= logits_output || + (*calls[j].outputs)[logits_output].dtype != ptn::kFloat || + (*calls[j].outputs)[logits_output].sizes.empty() || + calls[j].inputs[0].sizes.size() != 2 || + calls[j].inputs[1].sizes.size() != 1) { + throw std::runtime_error( + "sample calls need [1, n] tokens, [n] positions and Float logits"); + } + const std::vector ids = integer_values(calls[j].inputs[0]); + const int64_t vocabulary_size = + (*calls[j].outputs)[logits_output].sizes.back(); + if (std::any_of(ids.begin(), ids.end(), [&](int64_t id) { + return id < 0 || id >= vocabulary_size; + })) { + throw std::runtime_error( + "sample " + std::to_string(i) + " call " + std::to_string(j) + + " has a token outside the vocabulary"); + } + // The runner sends what generate() would: the prompt from 0, then one + // token per call. + const size_t count = j == 0 ? ids.size() : 1; + const int64_t start = + static_cast(j == 0 ? 0 : prompt.size() + j - 1); + std::vector positions(per_token_positions ? count : 1); + std::iota(positions.begin(), positions.end(), start); + if (ids.empty() || (j > 0 && ids.size() != 1) || + integer_values(calls[j].inputs[1]) != positions) { + throw std::runtime_error( + "sample " + std::to_string(i) + " call " + std::to_string(j) + + " is not a prefill or the next decode step"); + } + if (ids.size() > static_cast(token_shape[1])) { + throw std::runtime_error( + "sample " + std::to_string(i) + + " prefill exceeds the method's " + + std::to_string(token_shape[1]) + " token input"); + } + if (j == 0) { + prompt.assign(ids.begin(), ids.end()); + } else { + forced.push_back(static_cast(ids[0])); + } + } + // generate() runs one call per generated token. + if (ptn::generation_limit( + prompt.size(), + calls.size(), + static_cast(FLAGS_max_seq_length), + kv_cache_length(method), + static_cast(token_shape[1])) < calls.size()) { + throw std::runtime_error( + "sample " + std::to_string(i) + " does not fit the method"); + } + size_t call = 0; + std::string report; + const NextToken check_then_force = [&](const std::vector& logits, + size_t vocabulary_size) { + const ptn::SampleTensor& want = (*calls[call].outputs)[logits_output]; + if (want.bytes.size() != logits.size() * sizeof(float)) { + throw std::runtime_error("sample logits have a different size"); + } + std::vector want_logits(logits.size()); + std::memcpy(want_logits.data(), want.bytes.data(), want.bytes.size()); + const ptn::Tolerance tol = calls[call].tol.value_or(ptn::Tolerance{}); + const ptn::Closeness result = ptn::compare_close( + logits.data(), + want_logits.data(), + logits.size(), + tol.atol, + tol.rtol); + failed_calls += result.close ? 0 : 1; + char line[160]; + std::snprintf( + line, + sizeof(line), + "sample %zu call %zu: max_abs_error=%.6g nan=%zu (atol %g, rtol %g) " + "%s\n", + i, + call, + result.max_abs_error, + result.nan_count, + tol.atol, + tol.rtol, + result.close ? "ok" : "FAIL"); + report += line; + ++call; + return call < calls.size() + ? forced[call - 1] + : ptn::greedy_argmax_last_row(logits, vocabulary_size); + }; + generate( + *executable, + logits_output, + per_token_positions, + *tokenizer, + prompt, + calls.size(), + {}, + check_then_force); + std::printf("%s", report.c_str()); + } + std::printf( + "verify: %s/%s, %zu calls failed\n", + method_name.c_str(), + FLAGS_verify_samples.c_str(), + failed_calls); + return failed_calls == 0 ? 0 : 1; + } + const size_t limit = ptn::generation_limit( prompt_tokens.size(), static_cast(FLAGS_max_new_tokens), diff --git a/backends/native/examples/hf_llm/runtime/LlmRunnerUtils.cpp b/backends/native/examples/hf_llm/runtime/LlmRunnerUtils.cpp index 6b3e096c9d4..642f0fe6c25 100644 --- a/backends/native/examples/hf_llm/runtime/LlmRunnerUtils.cpp +++ b/backends/native/examples/hf_llm/runtime/LlmRunnerUtils.cpp @@ -175,4 +175,30 @@ size_t generation_limit( return std::min(max_new_tokens, max_total_tokens - prompt_length); } +// cppcheck-suppress unusedFunction +Closeness compare_close( + const float* got, + const float* want, + size_t n, + double atol, + double rtol) { + Closeness result; + for (size_t i = 0; i < n; ++i) { + if (got[i] == want[i]) { + continue; + } + const double error = + std::abs(static_cast(got[i]) - static_cast(want[i])); + if (std::isnan(error)) { + ++result.nan_count; + result.close = false; + continue; + } + result.close &= std::isfinite(error) && + error <= atol + rtol * std::abs(static_cast(want[i])); + result.max_abs_error = std::max(result.max_abs_error, error); + } + return result; +} + } // namespace ptn diff --git a/backends/native/examples/hf_llm/runtime/LlmRunnerUtils.h b/backends/native/examples/hf_llm/runtime/LlmRunnerUtils.h index 6e686d698f5..aaf182e379b 100644 --- a/backends/native/examples/hf_llm/runtime/LlmRunnerUtils.h +++ b/backends/native/examples/hf_llm/runtime/LlmRunnerUtils.h @@ -42,4 +42,19 @@ size_t generation_limit( size_t cache_length, size_t token_upper_bound); +struct Closeness { + double max_abs_error = 0.0; + size_t nan_count = 0; + bool close = true; +}; +// Largest non-NaN |got - want|, the number of elements whose difference is NaN, +// and whether every element is within atol + rtol * |want|. Equal values, +// infinities included, are close; a NaN never is. +Closeness compare_close( + const float* got, + const float* want, + size_t n, + double atol, + double rtol); + } // namespace ptn diff --git a/backends/native/examples/hf_llm/runtime/LlmRunnerUtilsTest.cpp b/backends/native/examples/hf_llm/runtime/LlmRunnerUtilsTest.cpp index c4b56aaff5f..7c69af5848a 100644 --- a/backends/native/examples/hf_llm/runtime/LlmRunnerUtilsTest.cpp +++ b/backends/native/examples/hf_llm/runtime/LlmRunnerUtilsTest.cpp @@ -140,6 +140,48 @@ TEST(LlmRunnerUtilsTest, TopLogitIndicesRankNanFirst) { (std::vector{2, 104, 103, 102, 101})); EXPECT_EQ(top_logit_indices(row.data(), 2, 5), (std::vector{0, 1})); } + +TEST(LlmRunnerUtilsTest, CompareCloseAppliesAbsoluteAndRelativeBounds) { + const std::vector want = {1.0f, -100.0f, 0.0f}; + const std::vector got = {1.5f, -101.0f, 0.0f}; + + const Closeness loose = compare_close(got.data(), want.data(), 3, 0.5, 0.01); + EXPECT_TRUE(loose.close); + EXPECT_DOUBLE_EQ(loose.max_abs_error, 1.0); + EXPECT_FALSE(compare_close(got.data(), want.data(), 3, 0.4, 0.01).close); + EXPECT_FALSE(compare_close(got.data(), want.data(), 3, 0.5, 0.004).close); +} + +TEST(LlmRunnerUtilsTest, CompareCloseScalesRtolByWant) { + const std::vector small = {100.0f}; + const std::vector large = {102.0f}; + + EXPECT_FALSE(compare_close(large.data(), small.data(), 1, 0.0, 0.0199).close); + EXPECT_TRUE(compare_close(small.data(), large.data(), 1, 0.0, 0.0199).close); +} + +TEST(LlmRunnerUtilsTest, CompareCloseRejectsNonFiniteMismatches) { + const float nan = std::numeric_limits::quiet_NaN(); + const float inf = std::numeric_limits::infinity(); + const std::vector want = {0.0f, inf, 2.0f}; + + const Closeness result = compare_close( + std::vector{nan, inf, 2.5f}.data(), want.data(), 3, 1.0, 1.0); + EXPECT_FALSE(result.close); + EXPECT_DOUBLE_EQ(result.max_abs_error, 0.5); + EXPECT_EQ(result.nan_count, 1); + const Closeness both_nan = compare_close( + std::vector{nan}.data(), std::vector{nan}.data(), 1, 1, 1); + EXPECT_FALSE(both_nan.close); + EXPECT_EQ(both_nan.nan_count, 1); + EXPECT_TRUE(compare_close(want.data() + 1, want.data() + 1, 2, 0, 0).close); + EXPECT_FALSE( + compare_close(std::vector{-inf}.data(), want.data() + 1, 1, 1, 1) + .close); + EXPECT_FALSE( + compare_close(std::vector{0.0f}.data(), want.data() + 1, 1, 1, 1) + .close); +} // cppcheck-suppress-end syntaxError } // namespace diff --git a/backends/native/examples/hf_llm/test_verify_samples.py b/backends/native/examples/hf_llm/test_verify_samples.py new file mode 100644 index 00000000000..876a4b0fa3d --- /dev/null +++ b/backends/native/examples/hf_llm/test_verify_samples.py @@ -0,0 +1,158 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import base64 +import copy +import os +import subprocess +import tempfile +import unittest + +import torch +from executorch.backends.native.examples.hf_llm.export_hf_llm import ( + golden_samples, + LastTokenLM, +) +from executorch.exir.native import Call, Sample, SampleSet, to_native, write_samples +from transformers import AutoModelForCausalLM, LlamaConfig + +_VOCAB = 128 +_CACHE = 40 +_MAX_TOKENS = 31 + + +def _tiny_lm() -> LastTokenLM: + config = LlamaConfig( + vocab_size=_VOCAB, + hidden_size=64, + intermediate_size=128, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + max_position_embeddings=64, + ) + model = AutoModelForCausalLM.from_config(config).eval().requires_grad_(False) + model.config._attn_implementation = "sdpa" + model.generation_config.use_cache = True + model.generation_config.cache_implementation = "static" + return LastTokenLM(model, _CACHE).eval() + + +def _forced(wrapper: LastTokenLM, prompt: list[int], forced: list[int]) -> Sample: + """A prefill of `prompt`, then one decode call per `forced` token.""" + wrapper.cache.reset() + input_ids = torch.tensor([prompt]) + positions = torch.arange(len(prompt)) + calls = [] + with torch.no_grad(): + for token in [None, *forced]: + if token is not None: + input_ids = torch.tensor([[token]]) + positions = positions[-1:] + 1 + calls.append(Call([input_ids, positions], [wrapper(input_ids, positions)])) + wrapper.cache.reset() + return Sample(calls) + + +class VerifySamplesTest(unittest.TestCase): + """Writes sample sets with the Python writer and replays them through the + C++ runner, which reads them with the C++ reader.""" + + @classmethod + def setUpClass(cls) -> None: + cls._tmp = tempfile.TemporaryDirectory() + # Byte-level TikToken: one rank per byte, so token IDs below 256 decode. + cls.tokenizer = os.path.join(cls._tmp.name, "tokenizer.model") + with open(cls.tokenizer, "w") as f: + for i in range(256): + f.write(f"{base64.b64encode(bytes([i])).decode()} {i}\n") + + torch.manual_seed(0) + cls.wrapper = _tiny_lm() + seq = torch.export.Dim("seq", min=1, max=_MAX_TOKENS) + with torch.no_grad(): + ep = torch.export.export( + cls.wrapper, + (torch.zeros((1, 6), dtype=torch.long), torch.arange(6)), + dynamic_shapes={"input_ids": {1: seq}, "cache_position": {0: seq}}, + strict=True, + ) + cls.path = os.path.join(cls._tmp.name, "m.ptn") + to_native(ep).save(cls.path) + cls.golden = golden_samples(cls.wrapper, [[3, 4, 5, 6, 7, 8], [9, 10]], 3) + write_samples(cls.path, "forward", "golden", cls.golden) + + @classmethod + def tearDownClass(cls) -> None: + cls._tmp.cleanup() + + def _verify(self, name: str, samples: SampleSet) -> subprocess.CompletedProcess: + write_samples(self.path, "forward", name, samples) + return subprocess.run( + [ + os.environ["RUN_HF_LLM"], + f"--package={self.path}", + f"--tokenizer={self.tokenizer}", + f"--verify_samples={name}", + ], + capture_output=True, + text=True, + ) + + def _modified_golden(self, edit) -> SampleSet: + samples = copy.deepcopy(self.golden) + edit(samples.samples[0].calls) + return samples + + def test_replays_golden_samples(self) -> None: + result = self._verify("replay", self.golden) + + self.assertEqual(result.returncode, 0, result.stdout + result.stderr) + self.assertEqual(result.stdout.count(" ok\n"), 8) + self.assertIn("forward/replay, 0 calls failed", result.stdout) + + def test_forces_recorded_tokens_over_its_own_choice(self) -> None: + forced = [17, 42, 99] + sample = _forced(self.wrapper, [3, 4, 5], forced) + argmax = [c.outputs[0][0, -1].argmax().item() for c in sample.calls[:-1]] + self.assertNotEqual(argmax, forced) + + result = self._verify("forced", SampleSet([sample], tol=(1e-3, 1e-3))) + + self.assertEqual(result.returncode, 0, result.stdout + result.stderr) + self.assertEqual(result.stdout.count(" ok\n"), 4) + + def test_fails_a_call_whose_logits_differ(self) -> None: + def perturb(calls: list[Call]) -> None: + calls[1].outputs = [calls[1].outputs[0] + 1.0] + + result = self._verify("wrong", self._modified_golden(perturb)) + + self.assertEqual(result.returncode, 1) + self.assertRegex(result.stdout, r"sample 0 call 1: max_abs_error=1 .* FAIL") + self.assertIn("forward/wrong, 1 calls failed", result.stdout) + + def test_rejects_samples_it_cannot_replay(self) -> None: + def out_of_vocabulary(calls: list[Call]) -> None: + calls[1].inputs[0] = torch.tensor([[_VOCAB]]) + + def wrong_position(calls: list[Call]) -> None: + calls[1].inputs[1] = torch.tensor([99]) + + cases = [ + (self._modified_golden(out_of_vocabulary), "outside the vocabulary"), + (self._modified_golden(wrong_position), "not a prefill or the next decode"), + ( + golden_samples(self.wrapper, [list(range(_MAX_TOKENS + 1))], 0), + "prefill exceeds", + ), + ] + for i, (samples, message) in enumerate(cases): + with self.subTest(message=message): + result = self._verify(f"invalid{i}", samples) + self.assertEqual(result.returncode, 2) + self.assertIn(message, result.stderr)