[executorch][native] Record embedding op samples

.ci/scripts/test-native.sh

@@ -43,7 +43,7 @@ case "${component}" in
     )
     # TODO: run the other sample tests once they pass.
     targets+=(run_hf_llm_ggml run_whisper_ggml_cpu native_ggml_engine_test native_llm_runner_utils_test run_classification_ggml_cpu native_output_check_test native_ggml_sample_test)
-    test_regex='^native_(ggml_engine|llm_runner_utils|output_check)_test$|^native_ggml_sample_test_act_quant$'
+    test_regex='^native_(ggml_engine|llm_runner_utils|output_check)_test$|^native_ggml_sample_test_(act_quant|no_linear)$'
     ;;
   vulkan)
     cmake_args+=(

.github/workflows/test-native.yml

@@ -18,6 +18,7 @@ on:
       - backends/cpu/runtime/**
       - backends/native/**
       - backends/cadence/utils/FACTO
+      - backends/test/suite/operators/test_embedding.py
       - backends/test/suite/operators/test_linear.py
       - backends/vulkan/CMakeLists.txt
       - backends/vulkan/cmake/**

backends/native/test/CMakeLists.txt

@@ -62,6 +62,7 @@ add_custom_command(
     executorch.backends.native.test.samples.generate --out ${_native_op_samples}
   COMMAND ${CMAKE_COMMAND} -E touch ${_native_op_samples}.stamp
   DEPENDS flatc ${_native_op_samples_srcs}
+          ${EXECUTORCH_ROOT}/backends/test/suite/operators/test_embedding.py
           ${EXECUTORCH_ROOT}/backends/test/suite/operators/test_linear.py
   VERBATIM
 )
@@ -160,9 +161,21 @@ if(EXECUTORCH_BUILD_NATIVE_GGML)
   set_tests_properties(native_ggml_sample_test PROPERTIES DISABLED TRUE)
   # native_ggml_sample_test_act_quant lets ggml quantize the activations of
   # weight-only quantized linears, judged against golden-int8_act.
-  add_test(NAME native_ggml_sample_test_act_quant
-           COMMAND native_ggml_sample_test --packages=${_native_op_samples}
-                   --device=cpu --allow_activation_quant_for_weight_only
+  # native_ggml_sample_test_no_linear judges the other ops without it.
+  #
+  # TODO: Drop the int4wo f32 embedding filters once ggml keeps their float
+  # scales exact (next diff).
+  add_test(
+    NAME native_ggml_sample_test_act_quant
+    COMMAND
+      native_ggml_sample_test --packages=${_native_op_samples} --device=cpu
+      --allow_activation_quant_for_weight_only
+      --gtest_filter=-embedding_int4wo_*.*_f32
+  )
+  add_test(
+    NAME native_ggml_sample_test_no_linear
+    COMMAND native_ggml_sample_test --packages=${_native_op_samples}
+            --device=cpu --gtest_filter=-linear_*:embedding_int4wo_*.*_f32
   )
   target_compile_options(
     native_ggml_sample_test

backends/native/test/runtime/ggml/BUCK

@@ -97,7 +97,8 @@ _VARIANTS = [
 ]
 
 # ggml_sample_test_act_quant lets ggml quantize the activations of weight-only
-# quantized linears, judged against golden-int8_act.
+# quantized linears, judged against golden-int8_act. ggml_sample_test_no_linear
+# judges the other ops without it.
 [
     fbcode_target(
         _kind = runtime.cxx_test,
@@ -109,15 +110,18 @@ _VARIANTS = [
         env = v.get("env", {}) | v.get("sample_env", {}),
         remote_execution = v.get("remote_execution"),
         # TODO: run the others in CI once they pass.
-        labels = [] if v["suffix"] + name == "_act_quant" else ci.labels(ci.map(ci.skip_test())),
+        labels = [] if v["suffix"] + name in ("_act_quant", "_no_linear") else ci.labels(ci.map(ci.skip_test())),
         test_main = "//executorch/backends/native/test/samples:sample_test_main",
         deps = [":ggml_engine_host" + v["suffix"]],
         **v.get("compat", {})
     )
     for v in _VARIANTS
+    # TODO: drop the int4wo f32 embedding filters once ggml keeps their float
+    # scales exact (next diff).
     for name, flags in [
         ("", []),
-        ("_act_quant", ["--allow_activation_quant_for_weight_only"]),
+        ("_act_quant", ["--allow_activation_quant_for_weight_only", "--gtest_filter=-embedding_int4wo_*.*_f32"]),
+        ("_no_linear", ["--gtest_filter=-linear_*:embedding_int4wo_*.*_f32"]),
     ]
 ]
 

backends/native/test/samples/BUCK

@@ -8,6 +8,7 @@ fbcode_target(
     name = "record",
     srcs = [
         "cases.py",
+        "embedding.py",
         "linear.py",
         "ops.py",
         "record.py",

backends/native/test/samples/cases.py

@@ -20,6 +20,8 @@
     from facto.inputgen.specs.model import Spec
 
 Inputs = tuple[torch.Tensor, ...]
+# Draws a tensor of the given shape from a generator, as torch.randn does.
+Sampler = Callable[..., torch.Tensor]
 
 # Caps a suite case's input and output elements over its samples, keeping at
 # least one sample: a large output alone has enough elements for statistics.
@@ -30,7 +32,8 @@
 class Case:
     name: str
     model: torch.nn.Module
-    # Positional float32 inputs of each sample; the count may vary by case.
+    # Positional inputs of each sample, floating ones float32; the count may
+    # vary by case.
     inputs: list[Inputs]
     # `{dim: Dim}` of each positional input. A dynamic case's inputs are the
     # calls of one sample, which run in order on one executable.
@@ -41,20 +44,27 @@ def outputs(result) -> list[torch.Tensor]:
     return list(result) if isinstance(result, (tuple, list)) else [result]
 
 
-def draws(count: int, *shapes: tuple[int, ...]) -> list[Inputs]:
-    """`count` samples of standard-normal inputs of the given shapes, from a fixed
-    seed, so equal calls return equal samples."""
+def draws(
+    count: int, *shapes: tuple[int, ...], sampler: Sampler = torch.randn
+) -> list[Inputs]:
+    """`count` samples of inputs of the given shapes from `sampler`, standard
+    normal by default, and a fixed seed, so equal calls return equal samples."""
     generator = torch.Generator().manual_seed(0)
     return [
-        tuple(torch.randn(shape, generator=generator) for shape in shapes)
+        tuple(sampler(shape, generator=generator) for shape in shapes)
         for _ in range(count)
     ]
 
 
-def chain(name: str, model: torch.nn.Module, *shapes: tuple[int | Dim, ...]) -> Case:
-    """A dynamic case with standard-normal inputs of `shapes`, from a fixed seed.
-    Each Dim starts mid-range, which export traces, then visits both bounds and
-    returns, every other Dim in the opposite order."""
+def chain(
+    name: str,
+    model: torch.nn.Module,
+    *shapes: tuple[int | Dim, ...],
+    sampler: Sampler = torch.randn,
+) -> Case:
+    """A dynamic case with inputs of `shapes` from `sampler`, standard normal by
+    default, and a fixed seed. Each Dim starts mid-range, which export traces,
+    then visits both bounds and returns, every other Dim in the opposite order."""
     dims = list(dict.fromkeys(d for s in shapes for d in s if not isinstance(d, int)))
     walks = {}
     for k, d in enumerate(dims):
@@ -63,7 +73,7 @@ def chain(name: str, model: torch.nn.Module, *shapes: tuple[int | Dim, ...]) ->
     generator = torch.Generator().manual_seed(0)
     inputs = [
         tuple(
-            torch.randn(
+            sampler(
                 [walks[d][step] if d in walks else d for d in s], generator=generator
             )
             for s in shapes
@@ -76,6 +86,17 @@ def chain(name: str, model: torch.nn.Module, *shapes: tuple[int | Dim, ...]) ->
     return Case(name, model, inputs, dynamic_shapes)
 
 
+class Parallel(torch.nn.Module):
+    """Modules on one input, each giving its own output."""
+
+    def __init__(self, *modules: torch.nn.Module) -> None:
+        super().__init__()
+        self.branches = torch.nn.ModuleList(modules)
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, ...]:
+        return tuple(m(x) for m in self.branches)
+
+
 class _SuiteRunner:
     """Stands in for the backend test suite's runner, keeping what a test lowers."""
 

backends/native/test/samples/embedding.py

@@ -0,0 +1,93 @@
+# 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.
+
+"""Embedding cases: the backend test suite's, LLM vocabularies and token
+shapes, index edge cases, and dynamic shapes."""
+
+from functools import partial
+
+import torch
+from executorch.backends.native.test.samples.cases import (
+    Case,
+    chain,
+    draws,
+    Parallel,
+    Sampler,
+    suite_cases,
+)
+from executorch.backends.native.test.samples.schemes import GGUF_SUPER_SCALES
+from executorch.backends.test.suite.operators import test_embedding
+from torch.export import Dim
+
+
+def _index_sampler(vocab: int) -> Sampler:
+    """A sampler of uniform int64 indices into a `vocab`-row table."""
+    return partial(torch.randint, vocab)
+
+
+def cases() -> list[Case]:
+    suite = suite_cases(test_embedding, samples=4)
+    torch.manual_seed(0)
+    tokens, batch = Dim("tokens", min=1, max=64), Dim("batch", min=1, max=4)
+    return [
+        *suite,
+        Case(
+            "decode",
+            torch.nn.Embedding(4096, 256),
+            draws(8, (1,), sampler=_index_sampler(4096)),
+        ),
+        Case(
+            "prefill",
+            torch.nn.Embedding(4096, 256),
+            draws(2, (64,), sampler=_index_sampler(4096)),
+        ),
+        Case(
+            "wide",
+            torch.nn.Embedding(512, 2048),
+            draws(4, (4,), sampler=_index_sampler(512)),
+        ),
+        Case(
+            "leading_batch",
+            torch.nn.Embedding(1000, 256),
+            draws(4, (2, 7), sampler=_index_sampler(1000)),
+        ),
+        Case("scalar_index", torch.nn.Embedding(64, 256), [(torch.tensor(5),)]),
+        Case(
+            "edge_rows",
+            torch.nn.Embedding(512, 256),
+            [(torch.tensor([0, 511, 0, 511, 7, 7]),)],
+        ),
+        Case(
+            "int32_indices",
+            torch.nn.Embedding(512, 256),
+            [(x.int(),) for (x,) in draws(2, (16,), sampler=_index_sampler(512))],
+        ),
+        # A zero row 0, and padding_idx in the graph, in the float and torchao
+        # packages.
+        Case(
+            "padding_idx",
+            torch.nn.Embedding(512, 256, padding_idx=0),
+            [(torch.tensor([0, 3, 0, 9]),)],
+        ),
+        # One output per family of GGUF super-block scales.
+        Case(
+            "scale_families",
+            Parallel(*(torch.nn.Embedding(64, 256) for _ in GGUF_SUPER_SCALES["q4_k"])),
+            draws(2, (8,), sampler=_index_sampler(64)),
+        ),
+        chain(
+            "dynamic_tokens",
+            torch.nn.Embedding(4096, 256),
+            (tokens,),
+            sampler=_index_sampler(4096),
+        ),
+        chain(
+            "dynamic_batch_tokens",
+            torch.nn.Embedding(1000, 256),
+            (batch, tokens),
+            sampler=_index_sampler(1000),
+        ),
+    ]

backends/native/test/samples/linear.py

@@ -13,6 +13,7 @@
     chain,
     draws,
     facto_cases,
+    Parallel,
     suite_cases,
 )
 from executorch.backends.native.test.samples.schemes import GGUF_SUPER_SCALES
@@ -81,17 +82,6 @@ def _facto_case(name: str, x, weight, bias) -> Case:
     return Case(name, model, [(x,)])
 
 
-class _Parallel(torch.nn.Module):
-    """nn.Linears on one input, each giving its own output."""
-
-    def __init__(self, *linears: torch.nn.Linear) -> None:
-        super().__init__()
-        self.linears = torch.nn.ModuleList(linears)
-
-    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, ...]:
-        return tuple(m(x) for m in self.linears)
-
-
 class _Pair(torch.nn.Module):
     """Two modules, each on its own input."""
 
@@ -202,7 +192,7 @@ def cases() -> list[Case]:
         # One output per family of GGUF super-block scales.
         Case(
             "scale_families",
-            _Parallel(*(torch.nn.Linear(256, 32) for _ in GGUF_SUPER_SCALES["q4_k"])),
+            Parallel(*(torch.nn.Linear(256, 32) for _ in GGUF_SUPER_SCALES["q4_k"])),
             draws(2, (4, 256)),
         ),
         # FACTO varies one argument at a time, so only every fifth tuple is kept,

backends/native/test/samples/ops.py

@@ -9,7 +9,7 @@
 from collections.abc import Callable
 from dataclasses import dataclass
 
-from executorch.backends.native.test.samples import linear
+from executorch.backends.native.test.samples import embedding, linear
 from executorch.backends.native.test.samples.cases import Case
 from executorch.backends.native.test.samples.record import DTYPES
 
@@ -21,4 +21,4 @@ class Op:
 
 
 def ops() -> dict[str, Op]:
-    return {"linear": Op(linear.cases)}
+    return {"embedding": Op(embedding.cases), "linear": Op(linear.cases)}

backends/native/test/samples/record.py

@@ -10,11 +10,12 @@
 and two sample sets over the same inputs: `golden`, the eager outputs, and
 `precise`, the same math in float64 with quantized weights decoded in float32,
 whatever the method's dtype, and activations unquantized, floating outputs
-stored as float32. Weight-only schemes add `golden-int8_act`, the eager outputs
-with the same weights and int8 per-token dynamic activations. A dynamic case
-is exported with its dynamic shapes and recorded as one sample whose calls
-change them. Nothing in it depends on an engine or a device: consumers judge an
-engine by its distance from `precise` relative to eager's.
+stored as float32. Weight-only schemes over a case with an nn.Linear add
+`golden-int8_act`, the eager outputs with the same weights and int8 per-token
+dynamic activations. A dynamic case is exported with its dynamic shapes and
+recorded as one sample whose calls change them. Nothing in it depends on an
+engine or a device: consumers judge an engine by its distance from `precise`
+relative to eager's.
 """
 
 import copy
@@ -61,6 +62,9 @@ def record(
 ) -> None:
     """Writes no package, removing any at `path`, when the scheme fits no case."""
     programs, golden, precise, golden_act = {}, [], [], []
+    act_quant = any(
+        scheme.adds_int8_act(c.model) for c in cases if scheme.fits(c.model)
+    )
     for case in cases:
         if not scheme.fits(case.model):
             logger.info(f"{path.name}: {case.name} does not fit the scheme")
@@ -92,7 +96,7 @@ def record(
                         t.float() if t.is_floating_point() else t for t in outputs(ref)
                     ]
                     precise_calls.append(Call(name, list(x), ref))
-                    if scheme.weight_only:
+                    if act_quant:
                         with int8_activations(method):
                             act_calls.append(Call(name, list(x), outputs(method(*x))))
             golden += _samples(golden_calls, chained)
@@ -114,7 +118,7 @@ def record(
             source,
         ),
     )
-    if scheme.weight_only:
+    if act_quant:
         write_samples(
             str(path),
             "golden-int8_act",

backends/native/test/samples/schemes.py

@@ -39,10 +39,15 @@ class Scheme:
 
     @property
     def weight_only(self) -> bool:
-        """Whether the scheme quantizes only weights; `record` then adds
-        `golden-int8_act`."""
+        """Whether the scheme quantizes only weights."""
         return False
 
+    def adds_int8_act(self, model: torch.nn.Module) -> bool:
+        """Whether `record` adds `golden-int8_act` for `model`: the scheme
+        quantizes only weights and `model` has inputs int8_activations
+        quantizes."""
+        return self.weight_only and bool(_targets(model, _DQA_QSCHEME_MODULES))
+
     def fits(self, model: torch.nn.Module) -> bool:
         return True