@@ -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+
10101116def _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
14781589def 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 []:
0 commit comments