Support exclusive raw-logit output dumps (#24505)

.ci/scripts/test_native_llm_runner.py

@@ -67,6 +67,7 @@ def main() -> None:
         prompt = "abc"
         unreachable_eos_id = config.vocab_size
         tokens = [config.bos_token_id, *prompt.encode()]
+        expected_logits = model(torch.tensor([tokens]), use_cache=False).logits[0, -1]
         expected = []
         # Recompute the full prefix without a cache for an independent reference.
         for _ in range(max_new_tokens):
@@ -75,6 +76,7 @@ def main() -> None:
             expected.append(token)
             tokens.append(token)
 
+        logits_dump = root / "first_logits.bin"
         result = subprocess.run(
             [
                 str(args.runner.resolve()),
@@ -87,6 +89,7 @@ def main() -> None:
                 f"--max_new_tokens={max_new_tokens}",
                 f"--max_seq_length={max_seq_length}",
                 "--print_token_ids",
+                f"--dump_first_logits={logits_dump}",
             ],
             capture_output=True,
             text=True,
@@ -102,7 +105,14 @@ def main() -> None:
         actual = [int(token) for token in match[1].split(",")]
         if actual != expected:
             raise AssertionError(f"Generated tokens {actual} != eager {expected}")
-        print(f"GGML HF LLM runner matched all {max_new_tokens} eager token IDs")
+        actual_logits = torch.frombuffer(
+            bytearray(logits_dump.read_bytes()), dtype=torch.float32
+        )
+        torch.testing.assert_close(actual_logits, expected_logits, atol=1e-3, rtol=1e-3)
+        print(
+            f"GGML HF LLM runner matched all {max_new_tokens} eager token IDs "
+            f"and {config.vocab_size} initial logits"
+        )
 
 
 if __name__ == "__main__":

backends/native/README.md

@@ -217,7 +217,8 @@ runs `native_cpu_engine_test` and builds `run_classification_cpu`. The GGML job
 also runs `native_llm_runner_utils_test` and executes `run_hf_llm_ggml` on a tiny
 seeded Hugging Face Llama with a generated byte tokenizer. Its eight generated
 token IDs must match eager PyTorch, covering prefill and KV-cache decoding without
-downloading a model. The GGML job requires the Python export dependencies
+downloading a model, and its `--dump_first_logits` output must match the eager
+prefill logits within 1e-3. The GGML job requires the Python export dependencies
 installed by the `Test Native Engines` workflow. The Vulkan job builds
 `run_classification` and `run_hf_llm_vulkan`; its tests do not require a GPU.
 

backends/native/examples/hf_llm/runtime/LlmRunner.cpp

@@ -202,20 +202,9 @@ Generation generate(
     executable.get_output(
         logits_output, logits.data(), logits.size(), ptn::kFloat);
     if (position == 0 && !FLAGS_dump_first_logits.empty()) {
-      FILE* output = std::fopen(FLAGS_dump_first_logits.c_str(), "wbx");
-      if (output == nullptr) {
-        throw std::runtime_error(
-            "cannot create new logits dump: " + FLAGS_dump_first_logits);
-      }
-      const float* row = logits.data() + logits_numel - vocabulary_size;
-      const size_t written =
-          std::fwrite(row, sizeof(float), vocabulary_size, output);
-      const int close_status = std::fclose(output);
-      if (written != vocabulary_size || close_status != 0) {
-        std::remove(FLAGS_dump_first_logits.c_str());
-        throw std::runtime_error(
-            "failed to write complete initial logits dump");
-      }
+      ptn::runner_utils::write_floats(
+          FLAGS_dump_first_logits,
+          {logits.data() + logits_numel - vocabulary_size, vocabulary_size});
     }
     if (FLAGS_print_top_logits) {
       const float* row = logits.data() + logits_numel - vocabulary_size;

backends/native/examples/runner_utils/runner_utils.cpp

@@ -35,6 +35,20 @@ std::vector<float> read_floats(const std::string& path) {
   return out;
 }
 
+void write_floats(const std::string& path, std::span<const float> values) {
+  FILE* file = std::fopen(path.c_str(), "wbx");
+  if (!file) {
+    throw std::runtime_error("cannot create " + path);
+  }
+  const size_t written =
+      std::fwrite(values.data(), sizeof(float), values.size(), file);
+  const int closed = std::fclose(file);
+  if (written != values.size() || closed != 0) {
+    std::remove(path.c_str());
+    throw std::runtime_error("cannot write " + path);
+  }
+}
+
 std::vector<std::string> read_lines(const std::string& path) {
   std::vector<std::string> lines;
   std::ifstream file(path);
@@ -73,9 +87,9 @@ std::vector<int> top_k(const std::vector<float>& values, int k) {
   return order;
 }
 
-bool compare_floats(
-    const std::vector<float>& actual,
-    const std::vector<float>& expected,
+bool can_compare_floats(
+    size_t actual_size,
+    size_t expected_size,
     double rtol,
     double atol) {
   if (!std::isfinite(rtol) || !std::isfinite(atol) || rtol < 0 || atol < 0) {
@@ -87,12 +101,23 @@ bool compare_floats(
         atol);
     return false;
   }
-  if (actual.size() != expected.size()) {
+  if (actual_size != expected_size) {
     std::fprintf(
         stderr,
         "error: reference has %zu values, model produced %zu\n",
-        expected.size(),
-        actual.size());
+        expected_size,
+        actual_size);
+    return false;
+  }
+  return true;
+}
+
+bool compare_floats(
+    const std::vector<float>& actual,
+    const std::vector<float>& expected,
+    double rtol,
+    double atol) {
+  if (!can_compare_floats(actual.size(), expected.size(), rtol, atol)) {
     return false;
   }
   size_t failures = 0;

backends/native/examples/runner_utils/runner_utils.h

@@ -7,18 +7,25 @@
 #pragma once
 
 #include <cstdint>
+#include <span>
 #include <string>
 #include <vector>
 
 namespace ptn::runner_utils {
 
 std::vector<float> read_floats(const std::string& path);
+void write_floats(const std::string& path, std::span<const float> values);
 std::vector<std::string> read_lines(const std::string& path);
 
 int64_t numel_of(const std::vector<int64_t>& sizes);
 std::string shape_str(const std::vector<int64_t>& sizes);
 std::vector<int> top_k(const std::vector<float>& values, int k);
 
+bool can_compare_floats(
+    size_t actual_size,
+    size_t expected_size,
+    double rtol,
+    double atol);
 bool compare_floats(
     const std::vector<float>& actual,
     const std::vector<float>& expected,

backends/native/examples/runner_utils/test_runner_utils.cpp

@@ -6,15 +6,120 @@
 
 #include <executorch/backends/native/examples/runner_utils/runner_utils.h>
 
+#include <bit>
+#include <csignal>
+#include <cstdint>
+#include <filesystem>
 #include <limits>
+#include <random>
+#include <stdexcept>
+#include <string>
 #include <vector>
 
 #include <gtest/gtest.h>
 
+#if !defined(_WIN32)
+#include <sys/resource.h>
+#endif
+
 namespace ptn::runner_utils {
 namespace {
 
+class RunnerOutputTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    directory_ = std::filesystem::path(::testing::TempDir()) /
+        ("runner_output_" + std::to_string(std::random_device{}()));
+    owns_directory_ = std::filesystem::create_directory(directory_);
+    ASSERT_TRUE(owns_directory_);
+    path_ = (directory_ / "output.bin").string();
+  }
+
+  void TearDown() override {
+    if (owns_directory_) {
+      std::filesystem::remove_all(directory_);
+    }
+  }
+
+  std::filesystem::path directory_;
+  std::string path_;
+  bool owns_directory_ = false;
+};
+
 // cppcheck-suppress-begin syntaxError
+TEST_F(RunnerOutputTest, WriteFloatsPreservesAllBits) {
+  const std::vector<uint32_t> bits{
+      0x00000000, 0x80000000, 0x3f800000, 0xc0400000, 0x7f800000, 0x7fc12345};
+  std::vector<float> values;
+  values.reserve(bits.size());
+  for (const uint32_t value : bits) {
+    values.push_back(std::bit_cast<float>(value));
+  }
+
+  write_floats(path_, values);
+
+  EXPECT_EQ(std::filesystem::file_size(path_), bits.size() * sizeof(float));
+  std::vector<uint32_t> actual;
+  for (const float value : read_floats(path_)) {
+    actual.push_back(std::bit_cast<uint32_t>(value));
+  }
+  EXPECT_EQ(actual, bits);
+}
+
+TEST_F(RunnerOutputTest, WriteFloatsDoesNotOverwriteExistingFile) {
+  const std::vector<float> original{1.0f, -3.0f};
+  write_floats(path_, original);
+
+  EXPECT_THROW(
+      write_floats(path_, std::vector<float>{4.0f}), std::runtime_error);
+
+  EXPECT_EQ(read_floats(path_), original);
+}
+
+TEST_F(RunnerOutputTest, WriteFloatsCreatesEmptyOutput) {
+  write_floats(path_, {});
+
+  EXPECT_EQ(std::filesystem::file_size(path_), 0);
+  EXPECT_TRUE(read_floats(path_).empty());
+}
+
+TEST_F(RunnerOutputTest, WriteFloatsReportsInvalidOutputPath) {
+  const std::string missing = (directory_ / "missing" / "output.bin").string();
+
+  EXPECT_THROW(
+      write_floats(missing, std::vector<float>{1.0f}), std::runtime_error);
+}
+
+#if !defined(_WIN32)
+TEST_F(RunnerOutputTest, WriteFloatsRemovesIncompleteOutput) {
+  const std::vector<float> values(1024, 1.0f);
+  rlimit original{};
+  ASSERT_EQ(getrlimit(RLIMIT_FSIZE, &original), 0);
+  rlimit limited = original;
+  limited.rlim_cur = sizeof(float);
+  const auto previous_handler = std::signal(SIGXFSZ, SIG_IGN);
+  ASSERT_EQ(setrlimit(RLIMIT_FSIZE, &limited), 0);
+
+  bool threw = false;
+  try {
+    write_floats(path_, values);
+  } catch (const std::runtime_error&) {
+    threw = true;
+  }
+  setrlimit(RLIMIT_FSIZE, &original);
+  std::signal(SIGXFSZ, previous_handler);
+
+  EXPECT_TRUE(threw);
+  EXPECT_FALSE(std::filesystem::exists(path_));
+}
+#endif
+
+TEST(RunnerUtilsTest, CanCompareFloatsChecksSizesAndTolerances) {
+  EXPECT_TRUE(can_compare_floats(2, 2, 1e-3, 0.0));
+  EXPECT_FALSE(can_compare_floats(1, 2, 1e-3, 1e-3));
+  EXPECT_FALSE(can_compare_floats(1, 1, -1e-3, 1e-3));
+}
+
 TEST(RunnerUtilsTest, CompareFloatsUsesAbsoluteAndRelativeTolerance) {
   EXPECT_TRUE(compare_floats({1.0f, 100.0f}, {1.0005f, 100.05f}, 1e-3, 1e-3));
   EXPECT_FALSE(compare_floats({1.0f, 100.0f}, {1.1f, 100.0f}, 1e-3, 1e-3));

backends/native/examples/vision/ClassificationRunner.cpp

@@ -28,6 +28,7 @@ DEFINE_string(
     "Reference logits to compare against, raw float32.");
 DEFINE_string(categories, "", "Newline-separated class labels, for display.");
 DEFINE_string(method, "", "Method to run. Defaults to the first one.");
+DEFINE_string(dump_output, "", "Write raw float32 logits to a new file.");
 DEFINE_double(rtol, 1e-3, "Relative tolerance for the logits comparison.");
 DEFINE_double(atol, 1e-3, "Absolute tolerance for the logits comparison.");
 namespace {
@@ -57,6 +58,7 @@ void print_top5(
 
 namespace ptn {
 
+// cppcheck-suppress unusedFunction
 int run_classification(
     const std::function<std::unique_ptr<EngineHost>()>& make_host,
     const std::function<ConstantBytes(const EngineContext&)>& constant_bytes) {
@@ -156,6 +158,18 @@ int run_classification(
     if (!FLAGS_categories.empty()) {
       labels = ptn::runner_utils::read_lines(FLAGS_categories);
     }
+    const std::vector<float> reference = FLAGS_expected.empty()
+        ? std::vector<float>{}
+        : ptn::runner_utils::read_floats(FLAGS_expected);
+    if (!FLAGS_expected.empty() &&
+        !ptn::runner_utils::can_compare_floats(
+            logits.size(), reference.size(), FLAGS_rtol, FLAGS_atol)) {
+      return 2;
+    }
+    if (!FLAGS_dump_output.empty()) {
+      ptn::runner_utils::write_floats(FLAGS_dump_output, logits);
+    }
+
     print_top5(logits, labels);
     const int argmax = ptn::runner_utils::top_k(logits, 1).front();
     std::printf(
@@ -165,8 +179,6 @@ int run_classification(
       std::printf("result:    ran (no reference given)\n");
       return 0;
     }
-    const std::vector<float> reference =
-        ptn::runner_utils::read_floats(FLAGS_expected);
     const bool pass = ptn::runner_utils::compare_floats(
         logits, reference, FLAGS_rtol, FLAGS_atol);
     std::printf("result:    %s\n", pass ? "PASS" : "FAIL");