[6508436] Fix BF16 FP8 ONNX export

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).

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",

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)

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,

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."""