[executorch][native-ggml] Reject weight-only quantized linears unless activation quant is allowed

backends/native/README.md

@@ -339,6 +339,9 @@ install/bin/run_hf_llm_ggml --device cpu --cpu_threads 4 \
 The GGML runner accepts `--device cpu|gpu`, `--device_index`, `--cpu_threads`,
 and `--max_cached_graphs` (default 1, 0 rebuilds each call). For GPU execution,
 enable the desired `GGML_*` backend when configuring ExecuTorch, as above.
+GGML quantizes the activations of its quantized matmuls, so it rejects
+weight-only quantized linears, including GGUF k-quants, unless
+`--allow_activation_quant_for_weight_only` accepts that loss of precision.
 
 With `EXECUTORCH_BUILD_TESTS=ON`, `native_llm_runner_utils_test` exercises the
 LLM generation utilities.

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

@@ -23,6 +23,11 @@ DEFINE_int32(
     max_cached_graphs,
     1,
     "Graphs kept per executable by input shapes; 0 rebuilds every call.");
+DEFINE_bool(
+    allow_activation_quant_for_weight_only,
+    false,
+    "Run weight-only quantized linears, such as GGUF k-quants, with ggml's "
+    "quantized activations.");
 
 int main(int argc, char** argv) {
   gflags::SetUsageMessage(
@@ -41,6 +46,8 @@ int main(int argc, char** argv) {
       throw std::invalid_argument("--max_cached_graphs must be non-negative");
     }
     options.max_cached_graphs = static_cast<size_t>(FLAGS_max_cached_graphs);
+    options.allow_activation_quant_for_weight_only =
+        FLAGS_allow_activation_quant_for_weight_only;
     return ptn::GgmlEngineHost::create(options);
   });
 }

backends/native/runtime/engine/ggml/GgmlEngine.h

@@ -40,6 +40,11 @@ struct GgmlEngineOptions {
   // Prints each newly built graph's node count per ggml op to stderr, which
   // shows whether a lowering rewrite fired.
   bool print_graph_ops = false;
+  // ggml backends quantize the activations of a quantized matmul themselves,
+  // so a linear with quantized weights, int4 or GGUF, whose input is not a
+  // dequantized activation is rejected unless the caller accepts that loss of
+  // precision.
+  bool allow_activation_quant_for_weight_only = false;
 };
 
 class GgmlEngineContext;

backends/native/runtime/engine/ggml/GgmlExecutor.cpp

@@ -542,7 +542,27 @@ class GgmlExecutable final : public EngineExecutable {
             "' may only feed linear or embedding weights, not '" + node.target +
             "'");
       }
+      if (op == "aten.linear.default" &&
+          !options_.allow_activation_quant_for_weight_only &&
+          !has_dequantized_input(node)) {
+        throw std::runtime_error(
+            "ggml: weight-only quantized linear '" + node.name +
+            "' needs allow_activation_quant_for_weight_only, since ggml "
+            "quantizes its activations");
+      }
+    }
+  }
+
+  bool has_dequantized_input(const Node& linear) const {
+    const NamedArgument* input = find_arg(linear, "input", 0);
+    if (input == nullptr || input->arg.kind() != ArgKind::Tensor) {
+      return false;
     }
+    const NodeId producer =
+        method_.graph.value(input->arg.as_tensor().id).producer_id;
+    return valid(producer) &&
+        op_key(method_.graph.node(producer).target) ==
+        "torchao.dequantize_affine.default";
   }
 
   // Converters read some dtypes through typed pointers, so the package's bytes

backends/native/test/e2e/test_cross_engine.py

@@ -35,7 +35,14 @@
 _MAX_SEQ = 16
 
 
-def _run(runner: str, package: Path, steps, outdir: Path, resize: bool = False):
+def _run(
+    runner: str,
+    package: Path,
+    steps,
+    outdir: Path,
+    resize: bool = False,
+    flags: tuple[str, ...] = (),
+):
     """Runs each step's inputs in order on one executable; returns each step's
     single float32 output."""
     input_steps, output_steps, shapes = [], [], []
@@ -55,6 +62,7 @@ def _run(runner: str, package: Path, steps, outdir: Path, resize: bool = False):
     ]
     if resize:
         command.append(f"--input_sizes={';'.join(shapes)}")
+    command.extend(flags)
     resources = files(__package__)
     with as_file(resources / runner) as binary, as_file(
         resources / "vulkan_driver"
@@ -390,11 +398,13 @@ def test_inplace_relu_of_dynamic_view(self):
         seq = Dim("seq", min=2, max=12)
         self._check_dynamic(_ViewRelu(), steps, ({1: seq},))
 
-    def _check_engines(self, package, steps, expected, outdir, atol_fraction, engines):
+    def _check_engines(
+        self, package, steps, expected, outdir, atol_fraction, engines, flags=()
+    ):
         tokens = {}
         for runner in engines:
             with self.subTest(engine=runner):
-                outputs = _run(runner, package, steps, outdir, resize=True)
+                outputs = _run(runner, package, steps, outdir, resize=True, flags=flags)
                 self.assertEqual(len(outputs), len(expected))
                 for step, (actual, reference) in enumerate(zip(outputs, expected)):
                     self.assertEqual(actual.shape, reference.shape)
@@ -433,6 +443,7 @@ def _check_llama(
         atol_fraction: float,
         engines=_ENGINES,
         max_seq: int = _MAX_SEQ,
+        flags: tuple[str, ...] = (),
     ):
         prompt = torch.randint(0, model.config.vocab_size, (1, 7))
         eager = LastTokenLM(copy.deepcopy(model), max_seq, kv_cache_dtype).eval()
@@ -457,7 +468,7 @@ def _check_llama(
             package = outdir / "tiny_llama.ptn"
             to_native(ep).save(str(package))
             tokens = self._check_engines(
-                package, steps, expected, outdir, atol_fraction, engines
+                package, steps, expected, outdir, atol_fraction, engines, flags
             )
         greedy = [e[0, -1].argmax().item() for e in expected]
         for runner, chosen in tokens.items():
@@ -472,16 +483,20 @@ def test_tiny_llama_gguf(self):
         model = _tiny_llama()
         pairs = _gguf_weights(model)
         load_gguf(model, "tiny.gguf", "llama", lambda _: iter(pairs))
-        # ggml requantizes activations to Q8_K for its K-quant matmuls, while
-        # native-vk decodes the K-quant weights exactly to fp32. Re-seeding
-        # gives both engines the same prompt.
-        for runner, atol_fraction in (
-            ("ggml_ptn_runner", 5e-2),
-            ("vulkan_ptn_runner", 1e-3),
+        # ggml requantizes activations to Q8_K for its K-quant matmuls, once
+        # allowed to, while native-vk decodes the K-quant weights exactly to
+        # fp32. Re-seeding gives both engines the same prompt.
+        for runner, atol_fraction, flags in (
+            ("ggml_ptn_runner", 5e-2, ("--allow_activation_quant_for_weight_only",)),
+            ("vulkan_ptn_runner", 1e-3, ()),
         ):
             torch.manual_seed(0)
             self._check_llama(
-                copy.deepcopy(model), torch.float32, atol_fraction, engines=(runner,)
+                copy.deepcopy(model),
+                torch.float32,
+                atol_fraction,
+                engines=(runner,),
+                flags=flags,
             )
 
     def test_tiny_llama_fp16_kv(self):

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

@@ -26,6 +26,10 @@ 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 {
 
@@ -44,6 +48,8 @@ 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;
   return GgmlEngineHost::create(options);
 }
 

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

@@ -81,6 +81,11 @@ flatbuffers::Offset<fbs::Argument> create_none_arg(
 
 constexpr size_t kQ4KBlockBytes = 144;
 
+// The safetensors header of the affine linear program's constants.
+constexpr const char* kQ4_0SafeTensorsHeader =
+    R"({"weight_scale":{"dtype":"F16","shape":[1,8],"data_offsets":[0,16]},)"
+    R"("weight":{"dtype":"U8","shape":[128],"data_offsets":[16,144]}})";
+
 // A [1, 256] x [256] linear with gguf:q4_k weights, or with symmetric int4
 // AffineGroup weights whose fp16 scales `weight_scale` cover groups of 32.
 // `cache` is an unused mutable buffer.
@@ -232,6 +237,10 @@ std::vector<uint8_t> make_linear_program(
       builder.GetBufferPointer() + builder.GetSize()};
 }
 
+// For tests of quantized linears whose activation numerics aren't under test.
+constexpr GgmlEngineOptions kAllowActivationQuant{
+    .allow_activation_quant_for_weight_only = true};
+
 std::unique_ptr<GgmlEngineContext> make_context(
     GgmlEngineHost& host,
     const std::vector<uint8_t>& program_bytes) {
@@ -473,7 +482,7 @@ std::vector<uint8_t> make_cache_update_program(
 
 // cppcheck-suppress syntaxError
 TEST(GgmlEngineTest, ExecutesPackedQ4KLinear) {
-  auto host = GgmlEngineHost::create();
+  auto host = GgmlEngineHost::create(kAllowActivationQuant);
   auto context = make_context(*host, make_linear_program());
   auto executable = context->compile("forward");
 
@@ -491,8 +500,9 @@ TEST(GgmlEngineTest, ExecutesPackedQ4KLinear) {
 TEST(GgmlEngineTest, RepeatedExecutionsReadFreshInputs) {
   for (const size_t max_cached_graphs : {size_t{0}, size_t{1}}) {
     SCOPED_TRACE(max_cached_graphs);
-    auto host = GgmlEngineHost::create(
-        GgmlEngineOptions{.max_cached_graphs = max_cached_graphs});
+    auto host = GgmlEngineHost::create(GgmlEngineOptions{
+        .max_cached_graphs = max_cached_graphs,
+        .allow_activation_quant_for_weight_only = true});
     auto context = make_context(*host, make_linear_program());
     auto executable = context->compile("forward");
     for (const float value : {1.0f, 2.0f, -1.0f}) {
@@ -580,7 +590,7 @@ TEST(GgmlEngineTest, RejectsUnsupportedOpAtCompile) {
 }
 
 TEST(GgmlEngineTest, RejectsSecondExecutableBindingMutableBuffer) {
-  auto host = GgmlEngineHost::create();
+  auto host = GgmlEngineHost::create(kAllowActivationQuant);
   auto context = make_context(
       *host,
       make_linear_program(
@@ -590,19 +600,43 @@ TEST(GgmlEngineTest, RejectsSecondExecutableBindingMutableBuffer) {
 }
 
 void expect_compile_error(
-    const std::vector<uint8_t>& program,
+    GgmlEngineContext& context,
     const std::string& message) {
-  auto host = GgmlEngineHost::create();
-  auto context = make_context(*host, program);
   try {
-    context->compile("forward");
+    context.compile("forward");
     FAIL() << "compile accepted the program";
   } catch (const std::runtime_error& error) {
     EXPECT_NE(std::string(error.what()).find(message), std::string::npos)
         << error.what();
   }
 }
 
+void expect_compile_error(
+    const std::vector<uint8_t>& program,
+    const std::string& message) {
+  auto host = GgmlEngineHost::create();
+  expect_compile_error(*make_context(*host, program), message);
+}
+
+TEST(GgmlEngineTest, RejectsWeightOnlyQuantizedLinearsByDefault) {
+  expect_compile_error(
+      make_linear_program(), "allow_activation_quant_for_weight_only");
+  const std::vector<uint8_t> program_bytes = make_linear_program(
+      "torch.ops.aten.linear.default", /*with_cache=*/false, /*affine=*/true);
+  auto host = GgmlEngineHost::create();
+  auto context = host->create_ggml_context(
+      std::make_shared<const Program>(
+          Program::load(program_bytes.data(), program_bytes.size())),
+      std::make_shared<const Package>(
+          Package::load(OwnedBytes::from_vector(testing::make_zip({
+              {kProgramEntry, program_bytes},
+              {kSafeTensorsEntry,
+               testing::make_safetensors(
+                   kQ4_0SafeTensorsHeader, /*scales + weights=*/16 + 128)},
+          })))));
+  expect_compile_error(*context, "allow_activation_quant_for_weight_only");
+}
+
 TEST(GgmlEngineTest, RejectsIntegerGelu) {
   TestProgram program;
   program.input("x", fbs::ScalarType::LONG, {2, 3});
@@ -773,7 +807,7 @@ TEST(GgmlEngineTest, WriteToCachedMaskReachesLaterAttention) {
 }
 
 TEST(GgmlEngineTest, ReleasesMutableBufferWithItsExecutable) {
-  auto host = GgmlEngineHost::create();
+  auto host = GgmlEngineHost::create(kAllowActivationQuant);
   auto context = make_context(
       *host,
       make_linear_program(
@@ -1067,7 +1101,7 @@ TEST(GgmlEngineTest, Q4KLinearOfAHalfInput) {
       {{"input", std::string("x")}, {"weight", std::string("weight")}},
       fbs::ScalarType::HALF,
       {1, 1});
-  auto host = GgmlEngineHost::create();
+  auto host = GgmlEngineHost::create(kAllowActivationQuant);
   auto context = make_context(*host, program.finish("y"));
   auto executable = context->compile("forward");
   // The IEEE half bit pattern of 1.0.
@@ -1399,7 +1433,9 @@ TEST(GgmlEngineTest, RejectsPackedWeightWithPartialBlocks) {
       {{"input", std::string("x")}, {"weight", std::string("weight")}},
       fbs::ScalarType::FLOAT,
       {1, 1});
-  expect_compile_error(program.finish("y"), "divisible by 256");
+  auto host = GgmlEngineHost::create(kAllowActivationQuant);
+  expect_compile_error(
+      *make_context(*host, program.finish("y")), "divisible by 256");
 }
 
 TEST(GgmlEngineTest, RejectsUnsupportedOpOnUnquantizedInput) {
@@ -1478,9 +1514,7 @@ TEST(GgmlEngineTest, ViewsQ4_0ScalesOfMappedPackage) {
     // alternate fp16 scales 0.5 and 1.0.
     constexpr size_t kScaleBytes = 16;
     constexpr size_t kWeightBytes = 128;
-    std::string header =
-        R"({"weight_scale":{"dtype":"F16","shape":[1,8],"data_offsets":[0,16]},)"
-        R"("weight":{"dtype":"U8","shape":[128],"data_offsets":[16,144]}})";
+    std::string header = kQ4_0SafeTensorsHeader;
     // The safetensors member comes first, so its payload starts right after
     // the local header, name, size prefix and JSON header. read_constant views
     // the scales in place only at an even offset.
@@ -1522,7 +1556,7 @@ TEST(GgmlEngineTest, ViewsQ4_0ScalesOfMappedPackage) {
     ASSERT_TRUE(view.has_value());
     ASSERT_EQ(reinterpret_cast<uintptr_t>(view->span().data()) % 2, misaligned);
 
-    auto host = GgmlEngineHost::create();
+    auto host = GgmlEngineHost::create(kAllowActivationQuant);
     auto context = host->create_ggml_context(
         std::make_shared<const Program>(
             Program::load(program_bytes.data(), program_bytes.size())),

backends/native/test/runtime/ggml/test_ggml_ops.py

@@ -305,6 +305,7 @@ def _run_ggml(
     max_cached_graphs: int = 1,
     graph_ops: collections.Counter | None = None,
     method: str = "forward",
+    allow_activation_quant_for_weight_only: bool = False,
 ):
     """Runs each step's inputs in order on one executable and returns each
     step's outputs as (bytes, shape) pairs. `graph_ops`, if given, collects
@@ -334,6 +335,8 @@ def _run_ggml(
         command.append(f"--input_sizes={';'.join(map(_shape_text, steps))}")
     if graph_ops is not None:
         command.append("--print_graph_ops")
+    if allow_activation_quant_for_weight_only:
+        command.append("--allow_activation_quant_for_weight_only")
     with as_file(files(__package__) / "ggml_ptn_runner") as runner:
         result = subprocess.run(
             [str(runner), *command],
@@ -387,15 +390,35 @@ def test_gguf_q6_k_linear(self):
         # block scale on GPUs. Test the weight layout tightly with exact
         # activations, then separately bound the randomized requantization loss.
         gpu = os.environ.get("GGML_PTN_DEVICE", "cpu") == "gpu"
-        self._check(model, (representable,), atol=5e-4, rtol=2e-4 if gpu else 1e-5)
-        self._check(model, (torch.randn(2, 256) / 8,), atol=0.35, rtol=0.02)
+        self._check(
+            model,
+            (representable,),
+            atol=5e-4,
+            rtol=2e-4 if gpu else 1e-5,
+            allow_activation_quant_for_weight_only=True,
+        )
+        self._check(
+            model,
+            (torch.randn(2, 256) / 8,),
+            atol=0.35,
+            rtol=0.02,
+            allow_activation_quant_for_weight_only=True,
+        )
 
     def test_gguf_q6_k_embedding(self):
         model = torch.nn.Embedding(8, 256)
         model.weight = torch.nn.Parameter(self._q6_k_weights(8), requires_grad=False)
         self._check(model, (torch.tensor([[0, 3, 6], [7, 1, 4]]),))
 
-    def _check(self, model, inputs, atol=1e-5, rtol=1e-5, graph_ops=None):
+    def _check(
+        self,
+        model,
+        inputs,
+        atol=1e-5,
+        rtol=1e-5,
+        graph_ops=None,
+        allow_activation_quant_for_weight_only=False,
+    ):
         model = model.eval()
         # Eager writes buffers in place; the package must keep them as constructed.
         with torch.no_grad():
@@ -406,7 +429,14 @@ def _check(self, model, inputs, atol=1e-5, rtol=1e-5, graph_ops=None):
             package = outdir / "model.ptn"
             to_native(torch.export.export(model, inputs)).save(str(package))
             (outputs,) = _run_ggml(
-                package, [inputs], outdir, len(expected), graph_ops=graph_ops
+                package,
+                [inputs],
+                outdir,
+                len(expected),
+                graph_ops=graph_ops,
+                allow_activation_quant_for_weight_only=(
+                    allow_activation_quant_for_weight_only
+                ),
             )
         self.assertEqual(len(outputs), len(expected))
         for (raw, shape), reference in zip(outputs, expected):