diff --git a/CHANGELOG.rst b/CHANGELOG.rst index c45d8025386..41de07cdf5d 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -64,6 +64,7 @@ Changelog - Fix Megatron Top-K/Top-P distillation gradients being multiplied by the tensor-parallel size when tensor parallelism is enabled; reported loss values are unchanged. - Fix the Puzzletron memory objective pricing the KV cache with ``hidden_size // num_attention_heads`` instead of the model's ``head_dim``, which misestimated it for models such as Mistral-Small, GPT-OSS and Gemma-3 and for every candidate with a pruned hidden size. Recompute subblock statistics produced with earlier versions. - Fix ONNX INT4 quantization ending the Python process when the INT4 packing extension fails to build, for example without a C++ compiler. It now falls back to the Python implementation. +- Fix ``--dq_only`` ONNX export producing a model that ONNX Runtime cannot load when the weights are quantized to an unsigned type. Such exports now load and match the source model. - Fix ``fsdp2_aware_weight_update`` masking an error raised during its unshard (e.g. a CUDA OOM) with a spurious ``UnboundLocalError``. 0.48.0 (2026-10-xx) diff --git a/modelopt/onnx/quantization/qdq_utils.py b/modelopt/onnx/quantization/qdq_utils.py index 89259c2bdce..619bc25dc74 100644 --- a/modelopt/onnx/quantization/qdq_utils.py +++ b/modelopt/onnx/quantization/qdq_utils.py @@ -60,6 +60,13 @@ onnx_bit_dtype_signed_map = {4: "INT4", 8: "INT8"} onnx_bit_dtype_unsigned_map = {4: "UINT4", 8: "UINT8"} +# Integer weight types qdq_to_dq can emit, keyed by the zero point's ONNX type. The +# DequantizeLinear left in the graph keeps its zero point, so the weight must match it. +qdq_weight_np_dtype_map = { + onnx.TensorProto.INT8: np.int8, + onnx.TensorProto.UINT8: np.uint8, +} + np_dtype_map = { "Float": np.float32, "Half": np.float16, @@ -657,7 +664,7 @@ def _convert_weight( ValueError: If scale shape doesn't match weight shape for the operation Note: - - INT8 weights are clipped to [-128, 127] + - Integer weights are clipped to the range of the zero point's type - FP8 weights use float8e4m3fn format """ weight_shape = weight_array.shape @@ -690,16 +697,29 @@ def _convert_weight( scale_array = scale_array.reshape(*reshape_dims) zp_array = zp_array.reshape(*reshape_dims) - # Convert to INT8/FP8 + # Convert to integer/FP8 if zp.data_type == onnx_dtype_map["Float8"]: scaled = np.asarray(weight_array / scale_array) + zp_array else: + weight_dtype = _qdq_weight_dtype(zp) + bounds = np.iinfo(weight_dtype) scaled = np.asarray((weight_array / scale_array).round()) - np.clip(scaled + zp_array, -128, 127, out=scaled) + np.clip(scaled + zp_array, bounds.min, bounds.max, out=scaled) return scaled +def _qdq_weight_dtype(zp: onnx.TensorProto) -> type[np.signedinteger | np.unsignedinteger]: + """Returns the numpy type a converted weight must use for this zero point.""" + weight_dtype = qdq_weight_np_dtype_map.get(zp.data_type) + if weight_dtype is None: + raise ValueError( + f"Unsupported zero point type for real weight quantization: " + f"{onnx.TensorProto.DataType.Name(zp.data_type)}" + ) + return weight_dtype + + def _cast_fp8(array: np.ndarray) -> np.ndarray: """Cast a numpy array to FLOAT8E4M3FN using PyTorch.""" array_f32_t = torch.from_numpy(array) @@ -791,8 +811,9 @@ def qdq_to_dq(onnx_model: onnx.ModelProto) -> onnx.ModelProto: new_weight = _create_fp8_tensor(scaled, weight_name) logger.debug(f"Converted {weight_name} to FP8") else: - new_weight = onnx.numpy_helper.from_array(scaled.astype("int8"), weight_name) - logger.debug(f"Converted {weight_name} to INT8") + weight_dtype = _qdq_weight_dtype(zp) + new_weight = onnx.numpy_helper.from_array(scaled.astype(weight_dtype), weight_name) + logger.debug(f"Converted {weight_name} to {weight_dtype.__name__.upper()}") if len(tensor_consumers[weight_name]) > 1 or weight_name in graph_outputs: # Keep the float value for other consumers and reuse the removed Q output name. new_weight.name = node.output[0] diff --git a/tests/unit/onnx/quantization/test_qdq_utils.py b/tests/unit/onnx/quantization/test_qdq_utils.py index 1b4670b6457..f932c38f7bf 100644 --- a/tests/unit/onnx/quantization/test_qdq_utils.py +++ b/tests/unit/onnx/quantization/test_qdq_utils.py @@ -1176,6 +1176,67 @@ def test_constant_node_scale_path_still_patched(self): assert (scale_arr > 0).all() +def create_test_model_with_matmul_qdq(zero_point: np.ndarray, scale: np.ndarray, weight): + """Create weight -> Q -> DQ -> MatMul, per-channel on the weight's output axis.""" + graph = helper.make_graph( + nodes=[ + helper.make_node( + "QuantizeLinear", ["weight", "w_scale", "w_zp"], ["w_q"], name="w_q", axis=1 + ), + helper.make_node( + "DequantizeLinear", ["w_q", "w_scale", "w_zp"], ["w_dq"], name="w_dq", axis=1 + ), + helper.make_node("MatMul", ["input", "w_dq"], ["output"], name="matmul"), + ], + name="matmul_qdq", + inputs=[helper.make_tensor_value_info("input", TensorProto.FLOAT, [4, weight.shape[0]])], + outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, [4, weight.shape[1]])], + initializer=[ + numpy_helper.from_array(weight, "weight"), + numpy_helper.from_array(scale, "w_scale"), + numpy_helper.from_array(zero_point, "w_zp"), + ], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 21)]) + model.ir_version = 10 + onnx.checker.check_model(model) + return model + + +class TestQdqToDqWeightType: + """The converted weight must keep the type of the zero point the DQ node still uses.""" + + @pytest.mark.parametrize("unsigned", [False, True]) + def test_converted_weight_matches_zero_point_type(self, unsigned): + weight = np.random.RandomState(3).randn(16, 32).astype(np.float32) + if unsigned: + # Asymmetric UINT8: the zero point is non-zero and the weight spans 0..255. + low, high = weight.min(axis=0), weight.max(axis=0) + scale = ((high - low) / 255.0).astype(np.float32) + zero_point = np.round(-low / scale).astype(np.uint8) + expected_type = TensorProto.UINT8 + else: + scale = (np.abs(weight).max(axis=0) / 127.0).astype(np.float32) + zero_point = np.zeros(32, dtype=np.int8) + expected_type = TensorProto.INT8 + + model = create_test_model_with_matmul_qdq(zero_point, scale, weight) + inputs = {"input": np.random.RandomState(5).randn(4, 16).astype(np.float32)} + + def run(m): + return ort.InferenceSession( + m.SerializeToString(), providers=["CPUExecutionProvider"] + ).run(None, inputs)[0] + + expected = run(model) + converted = qdq_to_dq(onnx.ModelProto.FromString(model.SerializeToString())) + + stored = next(t for t in converted.graph.initializer if t.name == "weight") + assert stored.data_type == expected_type + # onnx.checker does not catch a weight/zero-point type mismatch, so load it. + assert np.array_equal(run(converted), expected) + + def create_test_model_with_dq_transpose_matmul( out_features: int, in_features: int,