[executorch][native] Judge activation-quantized engines against an int8-activation golden

backends/native/test/CMakeLists.txt

@@ -155,6 +155,12 @@ if(EXECUTORCH_BUILD_NATIVE_GGML)
   )
   # TODO: Enable once every case passes.
   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
+  )
   target_compile_options(
     native_ggml_sample_test
     PRIVATE $<$<CXX_COMPILER_ID:MSVC>:/EHsc>

backends/native/test/runtime/PtnRunner.h

@@ -12,8 +12,15 @@
 
 namespace ptn::testing {
 
+// Numerics a driver lets an engine trade away. Drivers then judge the engine as
+// having made the trade, so an engine that never makes it rejects the option.
+struct HostOptions {
+  // Quantize the activations of linears whose weights alone are quantized.
+  bool allow_activation_quant_for_weight_only = false;
+};
+
 // Defined once per binary by an engine's host library; drivers such as
 // ptn_runner_main.cpp run every engine through the host it returns.
-std::unique_ptr<EngineHost> make_engine_host();
+std::unique_ptr<EngineHost> make_engine_host(const HostOptions& options);
 
 } // namespace ptn::testing

backends/native/test/runtime/ggml/BUCK

@@ -91,14 +91,16 @@ _VARIANTS = [
     for v in _VARIANTS
 ]
 
+# ggml_sample_test_act_quant lets ggml quantize the activations of weight-only
+# quantized linears, judged against golden-int8_act.
 [
     fbcode_target(
         _kind = runtime.cxx_test,
-        name = "ggml_sample_test" + v["suffix"],
+        name = "ggml_sample_test" + name + v["suffix"],
         args = [
             "--packages=$(location //executorch/backends/native/test/samples:op_samples)",
             "--device=" + v["device"],
-        ],
+        ] + flags,
         env = v.get("env", {}),
         remote_execution = v.get("remote_execution"),
         # TODO: run in CI once they pass.
@@ -108,6 +110,10 @@ _VARIANTS = [
         **v.get("compat", {})
     )
     for v in _VARIANTS
+    for name, flags in [
+        ("", []),
+        ("_act_quant", ["--allow_activation_quant_for_weight_only"]),
+    ]
 ]
 
 [

backends/native/test/runtime/ggml/ggml_ptn_runner.cpp

@@ -26,15 +26,11 @@ DEFINE_bool(
     print_graph_ops,
     false,
     "Print each built graph's node count per ggml op to stderr.");
-DEFINE_bool(
-    allow_activation_quant_for_weight_only,
-    false,
-    "Run weight-only quantized linears with ggml's quantized activations.");
 
 namespace ptn::testing {
 
 // cppcheck-suppress unusedFunction
-std::unique_ptr<EngineHost> make_engine_host() {
+std::unique_ptr<EngineHost> make_engine_host(const HostOptions& host_options) {
   GgmlEngineOptions options;
   if (FLAGS_device == "gpu") {
     options.device = GgmlDevice::Gpu;
@@ -49,7 +45,7 @@ std::unique_ptr<EngineHost> make_engine_host() {
   options.max_cached_graphs = static_cast<size_t>(FLAGS_max_cached_graphs);
   options.print_graph_ops = FLAGS_print_graph_ops;
   options.allow_activation_quant_for_weight_only =
-      FLAGS_allow_activation_quant_for_weight_only;
+      host_options.allow_activation_quant_for_weight_only;
   return GgmlEngineHost::create(options);
 }
 

backends/native/test/runtime/ptn_runner_main.cpp

@@ -65,6 +65,10 @@ DEFINE_int32(
     iterations,
     1,
     "Timed executions; above one, prints median and min latency.");
+DEFINE_bool(
+    allow_activation_quant_for_weight_only,
+    false,
+    "Let the engine quantize the activations of weight-only quantized linears.");
 
 namespace {
 
@@ -201,7 +205,9 @@ void run() {
   const ptn::ByteSpan program_bytes = package->program_bytes();
   const auto program = std::make_shared<const ptn::Program>(
       ptn::Program::load(program_bytes.data(), program_bytes.size()));
-  const auto host = ptn::testing::make_engine_host();
+  const auto host = ptn::testing::make_engine_host(
+      {.allow_activation_quant_for_weight_only =
+           FLAGS_allow_activation_quant_for_weight_only});
   const auto context = host->create_context(program, package);
   const auto executable = context->compile(FLAGS_method);
 

backends/native/test/runtime/vulkan/vulkan_ptn_runner.cpp

@@ -6,13 +6,20 @@
 
 // The Vulkan engine host for PTN test drivers.
 
+#include <stdexcept>
+
 #include <executorch/backends/native/runtime/vulkan/VulkanEngine.h>
 #include <executorch/backends/native/test/runtime/PtnRunner.h>
 
 namespace ptn::testing {
 
 // cppcheck-suppress unusedFunction
-std::unique_ptr<EngineHost> make_engine_host() {
+std::unique_ptr<EngineHost> make_engine_host(const HostOptions& options) {
+  if (options.allow_activation_quant_for_weight_only) {
+    throw std::invalid_argument(
+        "vulkan: never quantizes activations, so it rejects "
+        "allow_activation_quant_for_weight_only");
+  }
   return VulkanEngineHost::create();
 }
 

backends/native/test/samples/record.py

@@ -9,9 +9,10 @@
 A package holds one method per case and dtype under one quantization scheme,
 and two sample sets over the same inputs: `golden`, the eager outputs, and
 `precise`, the same math in float64 with weights dequantized and activations
-unquantized, floating outputs stored as float32. Nothing in it depends on an
-engine or a device: consumers judge an engine by its distance from `precise`
-relative to eager's.
+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. 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
@@ -50,7 +51,8 @@ def record(
     dtypes: Iterable[str] = tuple(DTYPES),
 ) -> None:
     """Writes no package, removing any at `path`, when the scheme fits no case."""
-    programs, golden, precise = {}, [], []
+    act_scheme = scheme.act_quant()
+    programs, golden, precise, golden_act = {}, [], [], []
     for case in cases:
         if not scheme.fits(case.model):
             logger.info(f"{path.name}: {case.name} does not fit the scheme")
@@ -63,6 +65,9 @@ def record(
             model = copy.deepcopy(case.model).to(dtype).eval()
             scheme.quantize(model)
             method = _Method({name: model})
+            if act_scheme is not None:
+                act_model = copy.deepcopy(case.model).to(dtype).eval()
+                act_scheme.quantize(act_model)
             state = {
                 k: scheme.decode(v).double() for k, v in method.state_dict().items()
             }
@@ -76,6 +81,10 @@ def record(
                         t.float() if t.is_floating_point() else t for t in outputs(ref)
                     ]
                     precise.append(Sample([Call(name, list(x), ref)]))
+                    if act_scheme is not None:
+                        golden_act.append(
+                            Sample([Call(name, list(x), outputs(act_model(*x)))])
+                        )
     if not programs:
         logger.info(f"{path.name}: no case fits the scheme, not written")
         path.unlink(missing_ok=True)
@@ -90,3 +99,9 @@ def record(
             precise, "float64, weights dequantized, activations unquantized", source
         ),
     )
+    if act_scheme is not None:
+        write_samples(
+            str(path),
+            "golden-int8_act",
+            SampleSet(golden_act, "eager, int8 per-token dynamic activations", source),
+        )

backends/native/test/samples/sample_test_main.cpp

@@ -11,7 +11,8 @@
 //
 // Usage:
 //   <test> --packages=a.ptn,dir [--rmse_ratio --max_ratio --rtol --sigmas]
-//          [engine host flags] [gtest flags]
+//          [--allow_activation_quant_for_weight_only] [engine host flags]
+//          [gtest flags]
 
 #include <algorithm>
 #include <cstdio>
@@ -56,21 +57,29 @@ DEFINE_double(
     sigmas,
     ptn::testing::Policy{}.sigmas,
     "Standard errors the fitted biases may deviate by.");
+DEFINE_bool(
+    allow_activation_quant_for_weight_only,
+    false,
+    "Let the engine quantize the activations of weight-only quantized linears, "
+    "judging it against golden-int8_act samples where a package has them.");
 
 namespace ptn::testing {
 namespace {
 
 struct Bundle {
   std::shared_ptr<const Package> package;
   std::shared_ptr<const Program> program;
+  std::string golden_set;
   SampleSet golden;
   SampleSet precise;
 };
 
 // Created on first use, so listing tests touches no device.
 EngineHost& host() {
   static const std::unique_ptr<EngineHost> instance = [] {
-    std::unique_ptr<EngineHost> made = make_engine_host();
+    std::unique_ptr<EngineHost> made = make_engine_host(
+        {.allow_activation_quant_for_weight_only =
+             FLAGS_allow_activation_quant_for_weight_only});
     std::printf(
         "engine %s on %s\n", made->name().c_str(), made->device_name().c_str());
     return made;
@@ -111,7 +120,8 @@ class SampleTest : public ::testing::Test {
         .rtol = FLAGS_rtol,
         .sigmas = FLAGS_sigmas};
     for (const auto& [key, check] : checks) {
-      const std::string name = key.first + " out " + std::to_string(key.second);
+      const std::string name = key.first + " out " +
+          std::to_string(key.second) + " vs " + bundle_->golden_set;
       std::printf("%s: %s\n", name.c_str(), check.summary().c_str());
       for (const std::string& failure : check.failures(policy)) {
         ADD_FAILURE() << name << ": " << failure;
@@ -129,7 +139,8 @@ class SampleTest : public ::testing::Test {
       const SampleCall& call = golden.calls[c];
       ASSERT_EQ(call.method, precise.calls[c].method);
       ASSERT_TRUE(call.inputs == precise.calls[c].inputs)
-          << call.method << " has golden and precise inputs that differ";
+          << call.method << " has " << bundle_->golden_set
+          << " and precise inputs that differ";
       ASSERT_TRUE(call.outputs && precise.calls[c].outputs)
           << call.method << " has no recorded outputs";
       const std::vector<SampleTensor>& want = *call.outputs;
@@ -208,7 +219,12 @@ void register_package(const std::filesystem::path& path) {
   bundle->program = std::make_shared<const Program>(
       Program::load(program.data(), program.size()));
   const ZipReader zip = ZipReader::open(path.string());
-  bundle->golden = read_sample_set(zip, "golden");
+  // Eager outputs with int8 per-token dynamic activations, which packages with
+  // weight-only quantized linears also carry.
+  const bool act_quant_golden = FLAGS_allow_activation_quant_for_weight_only &&
+      zip.member_size("samples/golden-int8_act/index.json").has_value();
+  bundle->golden_set = act_quant_golden ? "golden-int8_act" : "golden";
+  bundle->golden = read_sample_set(zip, bundle->golden_set);
   bundle->precise = read_sample_set(zip, "precise");
 
   std::set<std::string> methods;

backends/native/test/samples/schemes.py

@@ -33,6 +33,11 @@ def decode(self, value: torch.Tensor) -> torch.Tensor:
         """A state tensor of a quantized model as a plain tensor."""
         return value
 
+    def act_quant(self) -> "Scheme | None":
+        """This weight-only scheme's weights with int8 dynamic activations, or
+        None when the scheme is not weight-only."""
+        return None
+
 
 def _fits(model: torch.nn.Module, group: int | None) -> bool:
     """Whether `model` has an nn.Linear and every one's in_features split into
@@ -68,6 +73,18 @@ def quantize(self, model: torch.nn.Module) -> None:
     def decode(self, value: torch.Tensor) -> torch.Tensor:
         return value.dequantize() if isinstance(value, TorchAOBaseTensor) else value
 
+    def act_quant(self) -> Scheme | None:
+        if not isinstance(self.config, IntxWeightOnlyConfig):
+            return None
+        return Torchao(
+            Int8DynamicActivationIntxWeightConfig(
+                weight_dtype=self.config.weight_dtype,
+                weight_granularity=self.config.granularity,
+                weight_mapping_type=self.config.mapping_type,
+                weight_scale_dtype=self.config.scale_dtype,
+            )
+        )
+
 
 def schemes() -> dict[str, Scheme]:
     return {

backends/native/test/samples/test_record.py

@@ -77,6 +77,25 @@ def test_record_golden_is_eager_and_precise_is_float64(self) -> None:
             atol=0,
         )
 
+    def test_record_adds_the_8da4w_golden_to_weight_only_packages(self) -> None:
+        case = Case("linear", torch.nn.Linear(64, 8), draws(2, (3, 64)))
+        sets = {}
+        with tempfile.TemporaryDirectory() as tmp:
+            for name in ("float", "int4wo_g32", "8da4w_g32"):
+                path = Path(tmp) / f"{name}.ptn"
+                record([case], schemes()[name], path, dtypes=["f32"])
+                sets[name] = read_samples(str(path))
+        self.assertNotIn("golden-int8_act", sets["float"])
+        self.assertNotIn("golden-int8_act", sets["8da4w_g32"])
+        for got, want in zip(
+            sets["int4wo_g32"]["golden-int8_act"].samples,
+            sets["8da4w_g32"]["golden"].samples,
+            strict=True,
+        ):
+            torch.testing.assert_close(
+                got.calls[0].outputs, want.calls[0].outputs, rtol=0, atol=0
+            )
+
     def test_record_keeps_integer_outputs_precise(self) -> None:
         class Argmax(torch.nn.Module):
             def forward(self, x: torch.Tensor) -> torch.Tensor: