[executorch][native] Fold torchao q4 weight dequantizes into linear (pytorch#23337)

This PR was created by the merge bot to help merge the original PR into
the main branch.
ghstack PR number: pytorch#23232 by
@digantdesai
^ Please use this as the source of truth for the PR details, comments,
and reviews
ghstack PR base:
https://github.com/pytorch/executorch/tree/gh/digantdesai/112/base
ghstack PR head:
https://github.com/pytorch/executorch/tree/gh/digantdesai/112/head
Merge bot PR base:
https://github.com/pytorch/executorch/tree/gh/digantdesai/111/orig
Merge bot PR head:
https://github.com/pytorch/executorch/tree/gh/digantdesai/112/orig
Differential Revision:
[D121704321](https://our.internmc.facebook.com/intern/diff/D121704321/)
@diff-train-skip-merge

Co-authored-by: Digant Desai <digantdesai@meta.com>

Author: pytorchbot

backends/native/serialization/graph_serialize.py

@@ -390,8 +390,8 @@ def _node_to_arg_value(
     v: torch.fx.Node,
     subgraph_map: "dict[str, torch.fx.GraphModule] | None",
 ) -> ArgumentValue:
-    # A GGUF weight consumed through dequantize_gguf serializes as the raw packed
-    # constant (see _fold_gguf_dequant), so redirect the reference to it.
+    # A packed weight consumed through a folded dequantize serializes as the packed
+    # constant itself (see _fold_gguf_dequant, _mark_torchao_q4_weights).
     redirect = v.meta.get("native_serialize_as")
     if redirect is not None:
         return TensorArg(name=redirect)
@@ -800,8 +800,8 @@ def _collect_fx_nodes(
             ):
                 continue
 
-        # Folded GGUF dequantize nodes are not emitted; the consuming op references
-        # the packed weight directly (see _fold_gguf_dequant).
+        # Folded dequantize nodes are not emitted; the consuming op references the
+        # packed weight directly.
         if fx_node.meta.get("native_skip_serialize"):
             continue
 
@@ -1054,12 +1054,35 @@ def _fold_gguf_dequant(graph_module: torch.fx.GraphModule) -> None:
         node.meta["native_skip_serialize"] = True
 
 
-def _mark_torchao_q4_weights(graph_module: torch.fx.GraphModule) -> None:
+def _is_linear_weight_use(dequantize: torch.fx.Node, user: torch.fx.Node) -> bool:
+    op = _resolve_op_overload(user.target) if user.op == "call_function" else None
+    return (
+        op is not None
+        and op.name() == "aten::linear"
+        and len(user.args) >= 2
+        and user.args[1] is dequantize
+        and user.args[0] is not dequantize
+        and dequantize not in user.args[2:]
+        and dequantize not in user.kwargs.values()
+    )
+
+
+def _mark_torchao_q4_weights(
+    graph_module: torch.fx.GraphModule, lifted_constants: set[str] | None
+) -> None:
     """Mark portable torchao q4 weights for packed PTN storage.
 
-    The q/dq nodes remain in the serialized graph; only the weight storage
-    changes, to two four-bit values per byte.
+    A weight dequantize read only as the weight of `aten.linear`, whose output
+    dtype matches its scales and whose operands are serialized state inputs, is
+    folded: the linear reads the packed `AffineGroup` weight directly and the
+    dequantize is not serialized. A weight is packed only if every one of its
+    readers is such a dequantize with the same parameters. Any other dequantize
+    stays in the graph over the plain int8 weight, with torchao semantics.
+    Activation q/dq nodes are never folded.
     """
+    if lifted_constants is None:
+        return
+    folds: dict[torch.fx.Node, list[tuple[torch.fx.Node, dict[str, object]]]] = {}
     for node in graph_module.graph.nodes:
         op = _resolve_op_overload(node.target) if node.op == "call_function" else None
         if (
@@ -1105,12 +1128,35 @@ def _mark_torchao_q4_weights(graph_module: torch.fx.GraphModule) -> None:
             "zero_point_dtype": _scalar_type(zero_point_value.dtype),
             "group_size": block_size[1],
         }
-        existing = weight.meta.get("native_packed_quant")
-        if existing is not None and existing != packed:
-            raise ValueError(
-                f"constant {weight.name!r} has incompatible q4 interpretations"
-            )
+        if (
+            {weight.name, scale.name, zero_point.name} <= lifted_constants
+            and output.dtype == scale_value.dtype
+            and node.users
+            and all(_is_linear_weight_use(node, user) for user in node.users)
+        ):
+            folds.setdefault(weight, []).append((node, packed))
+
+    for weight, weight_folds in folds.items():
+        packed = weight_folds[0][1]
+        if len(weight.users) != len(weight_folds) or any(
+            p != packed for _, p in weight_folds
+        ):
+            continue
         weight.meta["native_packed_quant"] = packed
+        for node, _ in weight_folds:
+            node.meta["native_serialize_as"] = weight.name
+            node.meta["native_skip_serialize"] = True
+
+
+def _lifted_constant_names(graph_signature: object | None) -> set[str] | None:
+    if graph_signature is None:
+        return None
+    return {
+        ispec.arg.name
+        for ispec in getattr(graph_signature, "input_specs", []) or []
+        if ispec.kind in _INPUT_KIND_MAP
+        and isinstance(getattr(ispec, "arg", None), TensorArgument)
+    }
 
 
 def _build_graph_body(
@@ -1126,7 +1172,7 @@ def _build_graph_body(
     dict[str, torch.Tensor],
 ]:
     _fold_gguf_dequant(graph_module)
-    _mark_torchao_q4_weights(graph_module)
+    _mark_torchao_q4_weights(graph_module, _lifted_constant_names(graph_signature))
     subgraph_map = _subgraph_map(graph_module)
     nodes, val_by_name, output_names = _collect_fx_nodes(graph_module, subgraph_map)
 

backends/native/test/BUCK

@@ -59,7 +59,9 @@ fbcode_target(
         "//executorch/backends/native:lib",
         "//executorch/backends/native/test:utils",
         "//executorch/exir:lib",
+        "//executorch/exir/native:native",
         "//executorch/extension/llm/export:gguf",
+        "//pytorch/ao:torchao",
     ],
 )
 

backends/native/test/test_quant.py

@@ -11,13 +11,25 @@
 
 from executorch.backends.native.serialization import deserialize_program
 from executorch.backends.native.serialization.graph_serialize import _pack_signed_int4
-from executorch.backends.native.serialization.schema import PackedQuant, ScalarType
+from executorch.backends.native.serialization.schema import (
+    AffineGroup,
+    PackedQuant,
+    ScalarType,
+    TensorArg,
+)
 from executorch.backends.native.test.utils import (
     _call_function_targets,
     _get_delegate_blob,
     _lower,
 )
+from executorch.exir.native import to_native
 from executorch.extension.llm.export.gguf import ExportableGGUFTensor
+from torchao.quantization import (
+    Int8DynamicActivationIntxWeightConfig,
+    IntxWeightOnlyConfig,
+    quantize_,
+)
+from torchao.quantization.granularity import PerGroup
 
 
 class FuseGGUFPassTest(unittest.TestCase):
@@ -114,3 +126,132 @@ def test_rejects_values_outside_signed_int4_range(self):
             with self.subTest(value=value):
                 with self.assertRaisesRegex(ValueError, r"\[-8, 7\]"):
                     _pack_signed_int4(torch.tensor([[value, 0]], dtype=torch.int8))
+
+
+def _linear_weight(method):
+    [linear] = [n for n in method.graph.nodes if n.target and "aten.linear" in n.target]
+    [weight] = [a.arg.value for a in linear.inputs if a.name == "weight"]
+    assert isinstance(weight, TensorArg)
+    return weight.name
+
+
+class TorchaoQ4LinearTest(unittest.TestCase):
+    N, K, GROUP_SIZE = 16, 64, 32
+
+    def _serialize(self, config):
+        model = nn.Sequential(nn.Linear(self.K, self.N, bias=False))
+        quantize_(model, config)
+        manager = to_native(torch.export.export(model, (torch.randn(2, self.K),)))
+        return deserialize_program(manager._ptg).methods[0], manager._constants
+
+    def _assert_linear_reads_packed_weight(self, method, constants):
+        [weight] = [c for c in method.constants if c.meta.quant is not None]
+        scheme = weight.meta.quant.scheme
+        self.assertIsInstance(scheme, AffineGroup)
+        self.assertEqual(_linear_weight(method), weight.name)
+        self.assertEqual(weight.meta.dtype, ScalarType.BYTE)
+        self.assertEqual([d.max for d in weight.meta.sizes], [self.N, self.K])
+        self.assertEqual(
+            (scheme.quant_min, scheme.quant_max, scheme.group_size),
+            (-8, 7, self.GROUP_SIZE),
+        )
+        self.assertEqual(scheme.scale_dtype, ScalarType.FLOAT)
+        self.assertEqual(constants[weight.data_key].dtype, torch.uint8)
+        self.assertEqual(constants[weight.data_key].numel(), self.N * self.K // 2)
+        self.assertIn(scheme.scale_data_key, constants)
+        self.assertIn(scheme.zero_point_data_key, constants)
+
+    def test_8da4w_linear_reads_packed_weight(self):
+        method, constants = self._serialize(
+            Int8DynamicActivationIntxWeightConfig(
+                weight_dtype=torch.int4, weight_granularity=PerGroup(self.GROUP_SIZE)
+            )
+        )
+        targets = _call_function_targets(method.graph)
+        # Only the activation keeps its choose_qparams -> quantize -> dequantize.
+        self.assertEqual(sum("dequantize_affine" in t for t in targets), 1)
+        self.assertEqual(sum("quantize_affine" in t for t in targets), 2)
+        self.assertTrue(any("choose_qparams_affine" in t for t in targets))
+        self._assert_linear_reads_packed_weight(method, constants)
+
+    def test_4w_linear_reads_packed_weight(self):
+        method, constants = self._serialize(
+            IntxWeightOnlyConfig(
+                weight_dtype=torch.int4, granularity=PerGroup(self.GROUP_SIZE)
+            )
+        )
+        targets = _call_function_targets(method.graph)
+        self.assertFalse(any("quantize_affine" in t for t in targets))
+        self._assert_linear_reads_packed_weight(method, constants)
+
+
+class _DequantizedLinear(nn.Module):
+    """linear(x, dequantize_affine(weight)) over an int4 [8, 16] weight in groups
+    of 8, optionally with the weight as an input or the dequantize read twice."""
+
+    def __init__(self, output_dtype=torch.float32, weight_input=False, reuse=False):
+        super().__init__()
+        self.output_dtype = output_dtype
+        self.weight_input = weight_input
+        self.reuse = reuse
+        self.register_buffer("weight", torch.randint(-8, 8, (8, 16), dtype=torch.int8))
+        self.register_buffer("scale", torch.rand(8, 2) + 0.5)
+        self.register_buffer("zero_point", torch.zeros(8, 2, dtype=torch.int8))
+
+    def example_inputs(self):
+        x = torch.randn(2, 16, dtype=self.output_dtype)
+        return (x, self.weight.clone()) if self.weight_input else (x,)
+
+    def forward(self, x, weight=None):
+        weight = self.weight if weight is None else weight
+        dequantized = torch.ops.torchao.dequantize_affine(
+            weight,
+            [1, 8],
+            self.scale,
+            self.zero_point,
+            torch.int8,
+            -8,
+            7,
+            output_dtype=self.output_dtype,
+        )
+        out = torch.nn.functional.linear(x, dequantized)
+        return out + dequantized.sum() if self.reuse else out
+
+
+class FoldTorchaoQ4DequantizeTest(unittest.TestCase):
+    def _method(self, model):
+        return deserialize_program(
+            _get_delegate_blob(_lower(model, model.example_inputs()))
+        ).methods[0]
+
+    def _weight_dequantizes(self, model):
+        targets = _call_function_targets(self._method(model).graph)
+        return sum("dequantize_affine" in t for t in targets)
+
+    def _packed_constants(self, model):
+        return [c for c in self._method(model).constants if c.meta.quant is not None]
+
+    def test_folds_constant_weight_read_by_linear(self):
+        self.assertEqual(self._weight_dequantizes(_DequantizedLinear()), 0)
+        self.assertEqual(len(self._packed_constants(_DequantizedLinear())), 1)
+
+    def test_keeps_unfolded_weight_unpacked(self):
+        for model in (
+            _DequantizedLinear(reuse=True),
+            _DequantizedLinear(output_dtype=torch.float16),
+        ):
+            self.assertEqual(self._packed_constants(model), [])
+
+    def test_keeps_dequantize_of_runtime_weight(self):
+        self.assertEqual(
+            self._weight_dequantizes(_DequantizedLinear(weight_input=True)), 1
+        )
+
+    def test_keeps_dequantize_read_by_another_op(self):
+        self.assertEqual(self._weight_dequantizes(_DequantizedLinear(reuse=True)), 1)
+
+    def test_keeps_dequantize_to_another_dtype_than_scales(self):
+        self.assertEqual(
+            self._weight_dequantizes(_DequantizedLinear(output_dtype=torch.float16)),
+            1,
+        )