Skip to content

Commit 599f264

Browse files
[executorch][native] Store torchao q4 weights packed in PTN (pytorch#23330)
This PR was created by the merge bot to help merge the original PR into the main branch. ghstack PR number: pytorch#23225 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/105/base ghstack PR head: https://github.com/pytorch/executorch/tree/gh/digantdesai/105/head Merge bot PR base: https://github.com/pytorch/executorch/tree/gh/digantdesai/104/orig Merge bot PR head: https://github.com/pytorch/executorch/tree/gh/digantdesai/105/orig Differential Revision: [D121704315](https://our.internmc.facebook.com/intern/diff/D121704315/) @diff-train-skip-merge Co-authored-by: Digant Desai <digantdesai@meta.com>
1 parent 8ac8361 commit 599f264

2 files changed

Lines changed: 142 additions & 8 deletions

File tree

‎backends/native/serialization/graph_serialize.py‎

Lines changed: 124 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -253,6 +253,19 @@ def _packed_tensor_meta(codec: str, sizes: tuple[int, ...]) -> TensorMeta:
253253
)
254254

255255

256+
def _pack_signed_int4(weight: torch.Tensor) -> torch.Tensor:
257+
"""Pack signed int4 values as PTN AffineGroup qdata, even elements low."""
258+
if weight.dtype != torch.int8 or weight.ndim != 2:
259+
raise ValueError("native q4 weights must be rank-2 int8 tensors before packing")
260+
if weight.shape[1] % 2 != 0:
261+
raise ValueError("native q4 weight inner dimension must be even")
262+
if bool(torch.any((weight < -8) | (weight > 7)).item()):
263+
raise ValueError("native q4 weight values must be in [-8, 7]")
264+
# The uint8 cast wraps [-8, -1] to [248, 255], so + 8 wraps them to [0, 7].
265+
shifted = weight.to(torch.uint8) + 8
266+
return ((shifted[:, 1::2] << 4) | shifted[:, ::2]).contiguous()
267+
268+
256269
# ---------------------------------------------------------------------------
257270
# Argument dispatch.
258271
# ---------------------------------------------------------------------------
@@ -880,6 +893,13 @@ def _extract_constants_and_mutable_buffers(
880893
mutable_buffers: list[MutableBufferSpec] = []
881894
packed_quant = packed_quant or {}
882895

896+
fqn_by_name = {
897+
getattr(getattr(ispec, "arg", None), "name", None): getattr(
898+
ispec, "target", None
899+
)
900+
for ispec in getattr(graph_signature, "input_specs", []) or []
901+
}
902+
883903
for ispec in getattr(graph_signature, "input_specs", []) or []:
884904
if ispec.kind not in _INPUT_KIND_MAP:
885905
continue
@@ -899,11 +919,38 @@ def _extract_constants_and_mutable_buffers(
899919
continue
900920
tensor = tensor.contiguous()
901921
pq = packed_quant.get(name)
902-
meta = (
903-
_packed_tensor_meta(pq["codec"], pq["sizes"])
904-
if pq is not None
905-
else _tensor_meta(tensor)
906-
)
922+
# GGUF entries carry a codec; torchao q4 entries carry scale and
923+
# zero-point names.
924+
if pq is None:
925+
meta = _tensor_meta(tensor)
926+
elif "codec" in pq:
927+
meta = _packed_tensor_meta(pq["codec"], pq["sizes"])
928+
else:
929+
scale_data_key = fqn_by_name.get(pq["scale_name"])
930+
zero_point_data_key = fqn_by_name.get(pq["zero_point_name"])
931+
if not isinstance(scale_data_key, str) or not isinstance(
932+
zero_point_data_key, str
933+
):
934+
raise ValueError(
935+
"native q4 weight quantization parameters must be lifted "
936+
"constants"
937+
)
938+
tensor = _pack_signed_int4(tensor)
939+
meta = TensorMeta(
940+
dtype=ScalarType.BYTE,
941+
sizes=[_dim(s) for s in pq["sizes"]],
942+
quant=QuantSpec(
943+
scheme=AffineGroup(
944+
scale_data_key=scale_data_key,
945+
scale_dtype=pq["scale_dtype"],
946+
zero_point_data_key=zero_point_data_key,
947+
zero_point_dtype=pq["zero_point_dtype"],
948+
quant_min=-8,
949+
quant_max=7,
950+
group_size=pq["group_size"],
951+
)
952+
),
953+
)
907954
constant_refs.append(
908955
NamedTensorRef(
909956
name=name,
@@ -1007,6 +1054,65 @@ def _fold_gguf_dequant(graph_module: torch.fx.GraphModule) -> None:
10071054
node.meta["native_skip_serialize"] = True
10081055

10091056

1057+
def _mark_torchao_q4_weights(graph_module: torch.fx.GraphModule) -> None:
1058+
"""Mark portable torchao q4 weights for packed PTN storage.
1059+
1060+
The q/dq nodes remain in the serialized graph; only the weight storage
1061+
changes, to two four-bit values per byte.
1062+
"""
1063+
for node in graph_module.graph.nodes:
1064+
op = _resolve_op_overload(node.target) if node.op == "call_function" else None
1065+
if (
1066+
op is None
1067+
or op.name() != "torchao::dequantize_affine"
1068+
or len(node.args) < 7
1069+
# input_dtype, quant_min, quant_max
1070+
or node.args[4] != torch.int8
1071+
or node.args[5] != -8
1072+
or node.args[6] != 7
1073+
):
1074+
continue
1075+
weight, block_size, scale, zero_point = node.args[:4]
1076+
output = node.meta.get("val")
1077+
if (
1078+
not isinstance(weight, torch.fx.Node)
1079+
or weight.op != "placeholder"
1080+
or not isinstance(scale, torch.fx.Node)
1081+
or scale.op != "placeholder"
1082+
or not isinstance(zero_point, torch.fx.Node)
1083+
or zero_point.op != "placeholder"
1084+
or not isinstance(block_size, (list, tuple))
1085+
or len(block_size) != 2
1086+
or block_size[0] != 1
1087+
or not isinstance(block_size[1], int)
1088+
or block_size[1] <= 0
1089+
or not isinstance(output, torch.Tensor)
1090+
or output.ndim != 2
1091+
or output.shape[1] % block_size[1] != 0
1092+
):
1093+
continue
1094+
scale_value = scale.meta.get("val")
1095+
zero_point_value = zero_point.meta.get("val")
1096+
if not isinstance(scale_value, torch.Tensor) or not isinstance(
1097+
zero_point_value, torch.Tensor
1098+
):
1099+
continue
1100+
packed = {
1101+
"sizes": tuple(int(s) for s in output.shape),
1102+
"scale_name": scale.name,
1103+
"scale_dtype": _scalar_type(scale_value.dtype),
1104+
"zero_point_name": zero_point.name,
1105+
"zero_point_dtype": _scalar_type(zero_point_value.dtype),
1106+
"group_size": block_size[1],
1107+
}
1108+
existing = weight.meta.get("native_packed_quant")
1109+
if existing is not None and existing != packed:
1110+
raise ValueError(
1111+
f"constant {weight.name!r} has incompatible q4 interpretations"
1112+
)
1113+
weight.meta["native_packed_quant"] = packed
1114+
1115+
10101116
def _build_graph_body(
10111117
graph_module: torch.fx.GraphModule,
10121118
graph_signature: object | None,
@@ -1020,6 +1126,7 @@ def _build_graph_body(
10201126
dict[str, torch.Tensor],
10211127
]:
10221128
_fold_gguf_dequant(graph_module)
1129+
_mark_torchao_q4_weights(graph_module)
10231130
subgraph_map = _subgraph_map(graph_module)
10241131
nodes, val_by_name, output_names = _collect_fx_nodes(graph_module, subgraph_map)
10251132

@@ -1432,7 +1539,11 @@ def _validate_node(node: Node, defined: set[str], define: Any, check_ref: Any) -
14321539
_define_node_outputs(node, define)
14331540

14341541

1435-
def validate_graph(graph: Graph, defined_extra: set[str] | None = None) -> None:
1542+
def validate_graph(
1543+
graph: Graph,
1544+
defined_extra: set[str] | None = None,
1545+
meta_extra: set[str] | None = None,
1546+
) -> None:
14361547
"""Assert a (pure) graph body is structurally self-contained: every value
14371548
reference resolves and every input/output has tensor metadata. Recurses into
14381549
every ``GraphArg`` subgraph. Raises ``ValueError`` on the first inconsistency.
@@ -1472,7 +1583,7 @@ def check_ref(name: str, ctx: str) -> None:
14721583
for name in graph.outputs or []:
14731584
check_ref(name, "output")
14741585

1475-
_check_io_meta(graph, meta_names)
1586+
_check_io_meta(graph, meta_names | (meta_extra or set()))
14761587

14771588

14781589
def validate_method(
@@ -1491,7 +1602,12 @@ def validate_method(
14911602
"""
14921603
defined_extra: set[str] = {c.name for c in (method.constants or [])}
14931604
defined_extra.update(mb.name for mb in (method.mutable_buffers or []))
1494-
validate_graph(method.graph, defined_extra=defined_extra)
1605+
constant_meta = {c.name for c in (method.constants or [])}
1606+
validate_graph(
1607+
method.graph,
1608+
defined_extra=defined_extra,
1609+
meta_extra=constant_meta,
1610+
)
14951611

14961612
meta_names = {tv.name for tv in (method.graph.tensor_values or [])}
14971613
for mb in method.mutable_buffers or []:

‎backends/native/test/test_quant.py‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import torch.nn as nn
1111

1212
from executorch.backends.native.serialization import deserialize_program
13+
from executorch.backends.native.serialization.graph_serialize import _pack_signed_int4
1314
from executorch.backends.native.serialization.schema import PackedQuant, ScalarType
1415
from executorch.backends.native.test.utils import (
1516
_call_function_targets,
@@ -96,3 +97,20 @@ def test_gguf_embedding_serializes_as_packed_embedding(self):
9697
packed = self._packed_constants(program)
9798
self.assertEqual(len(packed), 1)
9899
self.assertEqual(packed[0].meta.quant.scheme.codec, "gguf:q4_k")
100+
101+
102+
class SignedInt4PackingTest(unittest.TestCase):
103+
def test_packs_even_value_into_low_nibble(self):
104+
weight = torch.tensor([[-8, -7, 6, 7]], dtype=torch.int8)
105+
106+
packed = _pack_signed_int4(weight)
107+
108+
torch.testing.assert_close(
109+
packed, torch.tensor([[0x10, 0xFE]], dtype=torch.uint8)
110+
)
111+
112+
def test_rejects_values_outside_signed_int4_range(self):
113+
for value in (-9, 8):
114+
with self.subTest(value=value):
115+
with self.assertRaisesRegex(ValueError, r"\[-8, 7\]"):
116+
_pack_signed_int4(torch.tensor([[value, 0]], dtype=torch.int8))

0 commit comments

Comments
 (0)