Arm backend: add SmolLM2 VGF generation and evaluation workflow

add server-mode sampling in executorch runner

add default prompt generation and Wikitext perplexity example scripts

document the end-to-end FP32, linear8a8w, and linear16a8w workflow

Signed-off-by: Xingguo Li <xingguo.li@arm.com>

Change-Id: I14472f987c739ea80457acc543e89b0e5a5f4a4c

Author: xingguo01

examples/arm/smollm2_example_vgf/README.md

@@ -34,7 +34,7 @@ If you want the broader Arm backend setup flow, see
 ## 1. Tokenizer (one-time)
 ```bash
 mkdir -p data/tokenizers/smollm2
-hf download HuggingFaceTB/SmolLM2-135M-Instruct tokenizer.json \
+huggingface-cli download HuggingFaceTB/SmolLM2-135M-Instruct tokenizer.json \
   --local-dir data/tokenizers/smollm2
 ```
 The download lives at `data/tokenizers/smollm2/tokenizer.json`. Use this path in the export and sampling commands below.

examples/arm/smollm2_example_vgf/build_executor_runner_vkml.sh

@@ -24,8 +24,6 @@ output_folder=$(realpath "${output_folder}")
 cmake \
     -S "${et_root_dir}" \
     -B "${output_folder}" \
-    -Wall \
-    -Werror \
     -DCMAKE_BUILD_TYPE=RelWithDebInfo \
     -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \
     -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \

examples/arm/smollm2_example_vgf/eval_wikitext_perplexity.py

@@ -9,21 +9,27 @@
 
 import argparse
 import math
+import sys
 from pathlib import Path
 from typing import Iterable, List, Tuple
 
 import numpy as np
 
+sys.path.insert(0, str(Path(__file__).resolve().parent))
 from generate_sampled import (  # type: ignore[import-not-found]
     prepare_input,
     RunnerSession,
 )
-from pytorch_tokenizers import get_tokenizer  # type: ignore[import-untyped]
+from pytorch_tokenizers import (  # type: ignore[import-not-found, import-untyped]
+    get_tokenizer,
+)
 
 
 def _load_wikitext_lines(split: str) -> Iterable[str]:
     try:
-        from datasets import load_dataset  # type: ignore[import-untyped]
+        from datasets import (  # type: ignore[import-not-found, import-untyped]
+            load_dataset,
+        )
     except ImportError as exc:
         raise ImportError(
             "The 'datasets' package is required. Install it in the active environment "
@@ -103,13 +109,6 @@ def read_prompts(path: Path, limit: int) -> List[str]:
     return prompts[:limit]
 
 
-def token_nll(logits: np.ndarray, target_id: int) -> float:
-    max_logit = float(np.max(logits))
-    shifted = logits - max_logit
-    log_denom = max_logit + math.log(float(np.exp(shifted).sum()))
-    return log_denom - float(logits[target_id])
-
-
 def reshape_full_logits(*, logits: np.ndarray, window: int) -> np.ndarray:
     if window <= 0:
         raise ValueError("window must be > 0")
@@ -153,13 +152,18 @@ def eval_prompt_nll(
     logits = runner.run(window_tokens)
     logits_2d = reshape_full_logits(logits=logits, window=window)
 
-    total_nll = 0.0
-    for pos, target_id in enumerate(target_ids[-valid_len:]):
-        if target_id >= logits_2d.shape[1]:
-            raise RuntimeError(
-                f"Target token id {target_id} out of inferred vocab size {logits_2d.shape[1]}."
-            )
-        total_nll += token_nll(logits_2d[pos], target_id)
+    # prepare_input keeps the scored tokens at positions [0, valid_len).
+    rows = logits_2d[:valid_len]
+    targets = np.asarray(target_ids[-valid_len:], dtype=np.int64)
+    if targets.size and (targets.min() < 0 or targets.max() >= rows.shape[1]):
+        raise RuntimeError(
+            f"Target token id out of inferred vocab size {rows.shape[1]}."
+        )
+
+    max_logits = rows.max(axis=1)
+    shifted = rows - max_logits[:, None]
+    log_denom = max_logits + np.log(np.exp(shifted).sum(axis=1))
+    total_nll = float((log_denom - rows[np.arange(valid_len), targets]).sum())
 
     return total_nll, valid_len
 

examples/arm/smollm2_example_vgf/generate_sampled.py

@@ -10,10 +10,12 @@
 import subprocess  # nosec B404
 import tempfile
 from pathlib import Path
-from typing import List, Optional
+from typing import List, Optional, Tuple
 
 import numpy as np
-from pytorch_tokenizers import get_tokenizer  # type: ignore[import-untyped]
+from pytorch_tokenizers import (  # type: ignore[import-not-found, import-untyped]
+    get_tokenizer,
+)
 
 
 class RunnerSession:
@@ -130,14 +132,15 @@ def run(self, tokens: np.ndarray) -> np.ndarray:
 
         assert self._input_path is not None
         tokens.tofile(self._input_path)
+        output_path = self._output_path()
+        output_path.unlink(missing_ok=True)
 
         if not self._persistent:
             cmd = self._base_cmd()
             proc = subprocess.run(  # nosec B603
                 cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT
             )
             out = proc.stdout.decode(errors="replace")
-            output_path = self._output_path()
 
             # Prefer checking for the output file instead of relying on stdout.
             if output_path.exists() and output_path.stat().st_size > 0:
@@ -173,6 +176,8 @@ def run(self, tokens: np.ndarray) -> np.ndarray:
                 self._recent_stdout = self._recent_stdout[-400:]
             if "SERVER MODE DONE" in line:
                 break
+        if not output_path.exists() or output_path.stat().st_size == 0:
+            raise RuntimeError("executor_runner did not write logits output")
         return self._read_logits()
 
 
@@ -383,7 +388,7 @@ def select_last_token_logits(
     window: int,
     use_full_logits: bool,
     valid_len: int,
-) -> tuple[np.ndarray, Optional[int]]:
+) -> Tuple[np.ndarray, Optional[int]]:
     """Select the logits row used for sampling and infer vocab size.
 
     Args:

examples/portable/executor_runner/executor_runner.cpp

@@ -150,7 +150,6 @@ Error load_input_files(
       return Error::Internal;
     }
 
-    // Reserve memory for actual file contents.
     // Reserve memory for actual file contents.
     inputs_storage.emplace_back(static_cast<size_t>(file_size), '\0');
 
@@ -161,6 +160,7 @@ Error load_input_files(
 
     input_buffers.emplace_back(
         inputs_storage.back().data(), static_cast<size_t>(file_size));
+  }
 
   return Error::Ok;
 }
@@ -506,6 +506,9 @@ int main(int argc, char** argv) {
   }
 
   if (FLAGS_server_mode) {
+    ET_CHECK_MSG(
+        !FLAGS_output_file.empty(),
+        "--output_file must be set in --server_mode so outputs can be retrieved.");
     ET_LOG(Info, "Running in server mode; waiting for stdin triggers.");
     std::vector<EValue> outputs(method->outputs_size());
     ET_LOG(Info, "%zu outputs: ", outputs.size());
@@ -568,6 +571,10 @@ int main(int argc, char** argv) {
                 out_filename, 255, "%s-%d.bin", FLAGS_output_file.c_str(), i);
             ET_LOG(Info, "Writing output to file: %s", out_filename);
             FILE* out_file = fopen(out_filename, "wb");
+            ET_CHECK_MSG(
+                out_file != nullptr,
+                "Failed to open output file: %s",
+                out_filename);
             fwrite(tensor.const_data_ptr<char>(), 1, tensor.nbytes(), out_file);
             fclose(out_file);
           }
@@ -662,6 +669,10 @@ int main(int argc, char** argv) {
         snprintf(out_filename, 255, "%s-%d.bin", FLAGS_output_file.c_str(), i);
         ET_LOG(Info, "Writing output to file: %s", out_filename);
         FILE* out_file = fopen(out_filename, "wb");
+        ET_CHECK_MSG(
+            out_file != nullptr,
+            "Failed to open output file: %s",
+            out_filename);
         fwrite(tensor.const_data_ptr<char>(), 1, tensor.nbytes(), out_file);
         fclose(out_file);
       }