diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 68927078e83..f3666f1bee4 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -54,6 +54,7 @@ Changelog - Avoid querying CUDA/Blackwell capability when ``NVFP4QTensor.quantize`` uses its CPU path or has the optional TensorRT-LLM fast path disabled. - Fix NVFP4 ONNX export to quantize FP4 weights with the published FP8 block scales, matching eager ModelOpt packed weights. Block scales below ``2**-9`` are now clamped to that minimum, and non-finite or negative scales raise an error. +- Fix FP8 ONNX export of BF16 models during real-weight compression. - Fix Megatron-Bridge Quantization Aware Distillation of a vision-language model silently discarding the ModelOpt state, so the distilled checkpoint restored no quantizers and exported as an unquantized model. Re-run QAD to regenerate any affected checkpoint. - Fix Megatron-Core HuggingFace export silently omitting fused (grouped GEMM) MoE experts for architectures without an ``experts.linear_fc1`` rule (e.g. ``Qwen3MoeForCausalLM``), which produced a valid-looking checkpoint containing no expert weights. The exporter now raises instead of writing that checkpoint; the scripts also avoid the situation by selecting ``SequentialMLP`` for those architectures. - Fix GatedDeltaNet (Qwen3.5) quantizer exclusions on Megatron-Core: the recipe patterns name the HuggingFace ``linear_attn`` module, so the ``conv1d`` was calibrated and the alpha / beta gate projections were exported in FP8. ``conv1d`` now has a ``self_attention`` alias in the default disabled-quantizer units, and the alpha / beta projections are exported in BF16 (they share Megatron's fused ``in_proj`` quantizer and cannot be disabled by name). diff --git a/modelopt/onnx/export/fp8_exporter.py b/modelopt/onnx/export/fp8_exporter.py index 427a7791f3b..b23019a2c03 100644 --- a/modelopt/onnx/export/fp8_exporter.py +++ b/modelopt/onnx/export/fp8_exporter.py @@ -17,6 +17,7 @@ import time +import ml_dtypes import numpy as np import onnx import onnx_graphsurgeon as gs @@ -33,6 +34,13 @@ _FP8_E4M3_SOFTMAX_SCALE = 1.0 / _FP8_E4M3_MAX +def _torch_from_numpy(array: np.ndarray) -> torch.Tensor: + """Convert a NumPy array to a PyTorch tensor while preserving BF16 values.""" + if array.dtype == ml_dtypes.bfloat16: + return torch.from_numpy(array.view(np.int16)).view(torch.bfloat16) + return torch.from_numpy(array) + + class FP8QuantExporter(ONNXQuantExporter): """Exporter for FP8 quantization.""" @@ -78,8 +86,11 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: weights = node.inputs[0] scale = node.inputs[1] - torch_weights = torch.from_numpy(weights.values) - torch_scale = torch.from_numpy(scale.values) + torch_weights = _torch_from_numpy(weights.values) + torch_scale = _torch_from_numpy(scale.values) + if torch.bfloat16 in (torch_weights.dtype, torch_scale.dtype): + torch_weights = torch_weights.float() + torch_scale = torch_scale.float() quantizer_name = scale.name.rsplit("/", 1)[0] dq_op = node.outputs[0].outputs[0] if dq_op.op != "TRT_FP8DequantizeLinear": @@ -194,14 +205,28 @@ def _quantize_conv_weights_to_fp8(graph: gs.Graph) -> int: if any(out.op == "DequantizeLinear" for out in weight_input.outputs): continue - torch_weights = torch.from_numpy(weight_input.values.copy()) + torch_weights = _torch_from_numpy(weight_input.values.copy()) amax = torch_weights.abs().max().float() if amax == 0: continue scale_val = (amax / _FP8_E4M3_MAX).item() + scale_data = np.array(scale_val, dtype=weight_input.values.dtype) + # Round up so normalizing by the serialized scale stays within the FP8 range. + if scale_data < scale_val: + np.nextafter( + scale_data, + np.array(np.inf, dtype=scale_data.dtype), + out=scale_data, + ) + torch_scale = _torch_from_numpy(scale_data) + if torch.bfloat16 in (torch_weights.dtype, torch_scale.dtype): + torch_weights = torch_weights.float() + torch_scale = torch_scale.float() # Quantize weights to FP8 (WAR: numpy doesn't support fp8) - fp8_data = (torch_weights / scale_val).to(torch.float8_e4m3fn).view(torch.uint8).numpy() + fp8_data = ( + (torch_weights / torch_scale).to(torch.float8_e4m3fn).view(torch.uint8).numpy() + ) fp8_tensor = onnx.TensorProto() fp8_tensor.data_type = onnx.TensorProto.FLOAT8E4M3FN fp8_tensor.dims.extend(fp8_data.shape) @@ -210,13 +235,16 @@ def _quantize_conv_weights_to_fp8(graph: gs.Graph) -> int: node.name + "/weight_quantizer/fp8_weights", LazyValues(fp8_tensor) ) - # Scale in FP16 — DQ output type matches scale dtype, must match activation type scale_constant = gs.Constant( node.name + "/weight_quantizer/scale", - np.array(scale_val, dtype=np.float16), + scale_data, ) - dq_output = gs.Variable(node.name + "/weight_quantizer/dq_output") + dq_output = gs.Variable( + node.name + "/weight_quantizer/dq_output", + scale_data.dtype, + weight_input.values.shape, + ) dq_node = gs.Node( op="DequantizeLinear", name=node.name + "/weight_quantizer/DequantizeLinear", diff --git a/modelopt/onnx/quantization/gs_patching.py b/modelopt/onnx/quantization/gs_patching.py index a0eea84951e..503a91f9721 100644 --- a/modelopt/onnx/quantization/gs_patching.py +++ b/modelopt/onnx/quantization/gs_patching.py @@ -101,9 +101,15 @@ def _export_value_info_proto(tensor: gs.Variable, do_type_check: bool) -> onnx.V ) if tensor.dtype is not None: - dtype = getattr( - tensor, "explicit_dtype", onnx.helper.np_dtype_to_tensor_dtype(np.dtype(tensor.dtype)) - ) + dtype = getattr(tensor, "explicit_dtype", None) + if dtype is None: + dtype = tensor.dtype + if isinstance(dtype, (int, np.integer)): + dtype = int(dtype) + if dtype not in onnx.TensorProto.DataType.values(): + raise ValueError(f"Unknown ONNX tensor dtype for {tensor.name}: {dtype}") + else: + dtype = onnx.helper.np_dtype_to_tensor_dtype(np.dtype(dtype)) onnx_tensor = onnx.helper.make_tensor_value_info(tensor.name, dtype, tensor.shape) else: onnx_tensor = onnx.helper.make_empty_tensor_value_info(tensor.name) diff --git a/modelopt/torch/_deploy/utils/torch_onnx.py b/modelopt/torch/_deploy/utils/torch_onnx.py index 01fb754bbae..6d828225986 100644 --- a/modelopt/torch/_deploy/utils/torch_onnx.py +++ b/modelopt/torch/_deploy/utils/torch_onnx.py @@ -24,6 +24,7 @@ import shutil import tempfile from contextlib import nullcontext +from itertools import chain from typing import Any import onnx @@ -527,6 +528,13 @@ def get_onnx_bytes_and_metadata( if isinstance(model, (DataParallel, DistributedDataParallel)): model = model.module + source_floating_dtypes = { + tensor.dtype + for tensor in chain(model.parameters(), model.buffers()) + if tensor.is_floating_point() + } + source_floating_dtype_names = ", ".join(sorted(map(str, source_floating_dtypes))) or "none" + # Standardize model args and also tensorize them so they also appear in the onnx graph! # Floats/ints are tensorized when they are provided, but not tensorized when they are not # provided which is somewhat inconsistent (we always tensorize them!) @@ -634,14 +642,28 @@ def get_onnx_bytes_and_metadata( if dq_only: onnx_opt_graph = qdq_to_dq(onnx_opt_graph) - if weights_dtype in ["fp16", "bf16"]: - if ( - is_int4_quantized(model) - or is_mxfp8_quantized(model) - or is_fp8_quantized(model) - or is_int8_quantized(model) - ): - assert weights_dtype == "fp16", "BF16 + MXFP8/INT4 mixed precision is not supported yet" + uses_fp8 = is_fp8_quantized(model) + uses_other_unsupported_quantizer = ( + is_int4_quantized(model) or is_mxfp8_quantized(model) or is_int8_quantized(model) + ) + if weights_dtype == "fp16" and uses_fp8 and torch.bfloat16 in source_floating_dtypes: + raise AssertionError( + "Converting a BF16 FP8 ONNX graph to FP16 is not supported yet " + f"(source floating dtypes: {source_floating_dtype_names})" + ) + + is_bf16_fp8_noop = ( + weights_dtype == "bf16" + and source_floating_dtypes == {torch.bfloat16} + and uses_fp8 + and not uses_other_unsupported_quantizer + ) + if weights_dtype in ["fp16", "bf16"] and not is_bf16_fp8_noop: + if uses_other_unsupported_quantizer or uses_fp8: + assert weights_dtype == "fp16", ( + "Converting a quantized ONNX graph to BF16 is not supported yet " + f"(source floating dtypes: {source_floating_dtype_names})" + ) onnx_opt_graph = convert_float_to_float16( onnx_opt_graph, keep_io_types=False, diff --git a/tests/unit/onnx/quantization/test_qdq_utils.py b/tests/unit/onnx/quantization/test_qdq_utils.py index 4b1e69ec538..f370b15211e 100644 --- a/tests/unit/onnx/quantization/test_qdq_utils.py +++ b/tests/unit/onnx/quantization/test_qdq_utils.py @@ -15,14 +15,22 @@ import warnings +import ml_dtypes import numpy as np +import onnx import onnx_graphsurgeon as gs import onnxruntime as ort import pytest from onnx import TensorProto, helper, numpy_helper -from modelopt.onnx.export import INT4QuantExporter, MXFP8QuantExporter, NVFP4QuantExporter +from modelopt.onnx.export import ( + FP8QuantExporter, + INT4QuantExporter, + MXFP8QuantExporter, + NVFP4QuantExporter, +) from modelopt.onnx.export.nvfp4_exporter import _cast_fp4 +from modelopt.onnx.quantization.gs_patching import _export_value_info_proto from modelopt.onnx.quantization.qdq_utils import ( _cast_fp8, apply_column_major_transformation, @@ -484,6 +492,59 @@ def test_cast_fp4(self, input_array, expected_array): assert np.all(result == expected_array) +def test_graphsurgeon_value_info_accepts_onnx_dtype(): + tensor = gs.Variable("input", dtype=TensorProto.BFLOAT16, shape=[1]) + value_info = _export_value_info_proto(tensor, do_type_check=True) + + assert value_info.type.tensor_type.elem_type == TensorProto.BFLOAT16 + + +class TestFP8QuantExporter: + """Test suite for FP8QuantExporter.""" + + def test_bf16_weights_and_scale_are_compressed(self): + weight_data = np.array([0.001312255859375], dtype=ml_dtypes.bfloat16) + scale_data = np.array(0.00099945068359375, dtype=ml_dtypes.bfloat16) + weight = gs.Constant("weight", weight_data) + scale = gs.Constant("linear/weight_quantizer/scale", scale_data) + quantized = gs.Variable("quantized", dtype=np.uint8, shape=weight_data.shape) + dequantized = gs.Variable("dequantized", dtype=ml_dtypes.bfloat16, shape=weight_data.shape) + graph = gs.Graph( + nodes=[ + gs.Node( + op="TRT_FP8QuantizeLinear", + inputs=[weight, scale], + outputs=[quantized], + ), + gs.Node( + op="TRT_FP8DequantizeLinear", + inputs=[quantized, scale], + outputs=[dequantized], + ), + ], + outputs=[dequantized], + opset=23, + ) + + converted_model = FP8QuantExporter.compress_weights(gs.export_onnx(graph)) + + onnx.checker.check_model(converted_model) + assert [node.op_type for node in converted_model.graph.node] == ["DequantizeLinear"] + fp8_weight = next( + initializer + for initializer in converted_model.graph.initializer + if initializer.name == "linear/weight_quantizer/fp8_weights" + ) + assert fp8_weight.data_type == TensorProto.FLOAT8E4M3FN + assert fp8_weight.raw_data == b"\x3b" + output_scale = next( + initializer + for initializer in converted_model.graph.initializer + if initializer.name == scale.name + ) + assert output_scale.data_type == TensorProto.BFLOAT16 + + class TestMXFP8QuantExporter: """Test suite for MXFP8QuantExporter.""" diff --git a/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py b/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py index 399085ad699..24135c8f7f8 100644 --- a/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py +++ b/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import copy import json from contextlib import nullcontext @@ -24,6 +25,7 @@ import torch.nn as nn from _test_utils.torch.deploy.lib_test_models import BaseDeployModel, get_deploy_models +import modelopt.torch.quantization as mtq from modelopt.onnx.utils import get_batch_size_from_bytes, validate_batch_size from modelopt.torch._deploy.utils import ( OnnxBytes, @@ -57,6 +59,32 @@ } +def _export_fp8_model(source_dtype, weights_dtype, conv=False): + if conv: + model = nn.Sequential(nn.Conv2d(1, 1, 1, bias=False)) + sample_input = torch.ones(1, 1, 2, 2, dtype=source_dtype) + else: + model = nn.Sequential(nn.Linear(4, 4, bias=False)) + sample_input = torch.arange(4, dtype=source_dtype).reshape(1, 4) + model = model.eval().to(source_dtype) + if conv: + with torch.no_grad(): + model[0].weight.fill_(1e-38 if source_dtype == torch.bfloat16 else 1.0) + model = mtq.quantize( + model, + mtq.FP8_DEFAULT_CFG, + forward_loop=lambda quantized_model: quantized_model(sample_input), + ) + onnx_bytes, _ = get_onnx_bytes_and_metadata( + model, + (sample_input,), + weights_dtype=weights_dtype, + onnx_opset=23, + ) + onnx_bytes_obj = OnnxBytes.from_bytes(onnx_bytes) + return onnx.load_model_from_string(onnx_bytes_obj.get_onnx_model_file_bytes()) + + @pytest.mark.parametrize( "model", deploy_benchmark_dynamo.values(), ids=deploy_benchmark_dynamo.keys() ) @@ -157,6 +185,139 @@ def test_onnx_export_and_inputs(model: BaseDeployModel): ) +@pytest.mark.parametrize( + ("source_dtype", "weights_dtype", "expected_onnx_dtype", "conv"), + [ + (torch.bfloat16, "bf16", onnx.TensorProto.BFLOAT16, False), + (torch.bfloat16, "bf16", onnx.TensorProto.BFLOAT16, True), + (torch.float32, "fp16", onnx.TensorProto.FLOAT16, False), + (torch.float32, "fp16", onnx.TensorProto.FLOAT16, True), + ], +) +def test_fp8_export_with_supported_weights_dtype( + source_dtype, weights_dtype, expected_onnx_dtype, conv +): + exported_model = _export_fp8_model(source_dtype, weights_dtype, conv) + + onnx.checker.check_model(exported_model, full_check=True) + assert not any( + node.op_type in {"TRT_FP8QuantizeLinear", "TRT_FP8DequantizeLinear"} + for node in exported_model.graph.node + ) + initializer_by_name = { + initializer.name: initializer for initializer in exported_model.graph.initializer + } + fp8_weight_dq_nodes = [ + node + for node in exported_model.graph.node + if node.op_type == "DequantizeLinear" + and node.input[0] in initializer_by_name + and initializer_by_name[node.input[0]].data_type == onnx.TensorProto.FLOAT8E4M3FN + ] + assert fp8_weight_dq_nodes + assert all( + initializer_by_name[node.input[1]].data_type == expected_onnx_dtype + for node in fp8_weight_dq_nodes + ) + if conv: + assert all( + set(initializer_by_name[node.input[0]].raw_data).isdisjoint({0x7F, 0xFF}) + for node in fp8_weight_dq_nodes + ) + graph_io = [*exported_model.graph.input, *exported_model.graph.output] + assert all(value.type.tensor_type.elem_type == expected_onnx_dtype for value in graph_io) + + +@pytest.mark.parametrize( + ("source_dtype", "weights_dtype", "error"), + [ + ( + torch.float32, + "bf16", + r"Converting a quantized ONNX graph to BF16.*source floating dtypes: torch.float32", + ), + ( + torch.bfloat16, + "fp16", + r"Converting a BF16 FP8 ONNX graph to FP16.*source floating dtypes: torch.bfloat16", + ), + ], +) +def test_fp8_export_rejects_unsupported_dtype_conversion(source_dtype, weights_dtype, error): + with pytest.raises(AssertionError, match=error): + _export_fp8_model(source_dtype, weights_dtype) + + +def test_fp8_bf16_noop_rejects_incompatible_mixed_format(): + model = nn.Sequential(*(nn.Linear(128, 128, bias=False) for _ in range(2))) + model = model.eval().to(torch.bfloat16) + sample_input = torch.ones(1, 128, dtype=torch.bfloat16) + config = copy.deepcopy(mtq.FP8_DEFAULT_CFG) + config["quant_cfg"].extend( + [ + { + "quantizer_name": "1.weight_quantizer", + "cfg": {"num_bits": 4, "block_sizes": {-1: 128, "type": "static"}}, + }, + {"quantizer_name": "1.input_quantizer", "enable": False}, + ] + ) + model = mtq.quantize(model, config, forward_loop=lambda model: model(sample_input)) + with pytest.raises( + AssertionError, + match="Converting a quantized ONNX graph to BF16 is not supported", + ): + get_onnx_bytes_and_metadata(model, (sample_input,), weights_dtype="bf16", onnx_opset=23) + + +def test_fp8_bf16_noop_rejects_mixed_source_dtypes(): + class MixedDtypeModel(nn.Module): + def __init__(self): + super().__init__() + self.bf16_layer = nn.Linear(4, 4, bias=False).to(torch.bfloat16) + self.fp32_layer = nn.Linear(4, 4, bias=False) + + def forward(self, inputs): + hidden = self.bf16_layer(inputs) + return self.fp32_layer(hidden.float()) + + model = MixedDtypeModel().eval() + sample_input = torch.ones(1, 4, dtype=torch.bfloat16) + model = mtq.quantize( + model, + mtq.FP8_DEFAULT_CFG, + forward_loop=lambda quantized_model: quantized_model(sample_input), + ) + with pytest.raises( + AssertionError, match="Converting a quantized ONNX graph to BF16 is not supported" + ): + get_onnx_bytes_and_metadata(model, (sample_input,), weights_dtype="bf16", onnx_opset=23) + + +def test_fp8_bf16_noop_rejects_mixed_buffer_dtype(): + class MixedBufferModel(nn.Module): + def __init__(self): + super().__init__() + self.bf16_layer = nn.Linear(4, 4, bias=False).to(torch.bfloat16) + self.register_buffer("fp32_offset", torch.ones(4)) + + def forward(self, inputs): + return self.bf16_layer(inputs).float() + self.fp32_offset + + model = MixedBufferModel().eval() + sample_input = torch.ones(1, 4, dtype=torch.bfloat16) + model = mtq.quantize( + model, + mtq.FP8_DEFAULT_CFG, + forward_loop=lambda quantized_model: quantized_model(sample_input), + ) + with pytest.raises( + AssertionError, + match=r"source floating dtypes: torch.bfloat16, torch.float32", + ): + get_onnx_bytes_and_metadata(model, (sample_input,), weights_dtype="bf16", onnx_opset=23) + + class SingleArgModel(nn.Module): def forward(self, x: torch.Tensor): return torch.add(x, x) - x