[6721556] Fix mixed dtypes in NVFP4 FP16 ONNX export

CHANGELOG.rst

@@ -25,6 +25,7 @@ Changelog
 
 **Bug Fixes**
 
+- Fix NVFP4 ONNX exports that failed ONNX or TensorRT parsing due to mixed-precision ``MatMul``, ``Gemm``, and elementwise inputs. Export now preserves FP32 graph boundaries around low-precision NVFP4 compute and rejects unsupported FP32-to-BF16 and BF16-to-FP16 conversions; use FP32-to-FP16 or native-BF16 export instead.
 - Fix ONNX INT8 entropy calibration failing or producing invalid quantization parameters for FP16 activations.
 - Fix ``--use_fsdp2`` HuggingFace checkpoint export gathering the whole model onto rank 0, which made export the dominant phase of a PTQ run and could exhaust host memory on large models. The model is now split into per-decoder-layer units dealt round-robin across ranks; each rank gathers every unit but keeps, packs, and writes only the ones it owns, so a rank buffers roughly ``model / world_size`` instead of the whole checkpoint, and rank 0 writes the combined index. Export configurations that cannot be split this way now raise instead of producing a mismatched checkpoint: FSDP2 combined with another DTensor parallelism (for example FSDP2 + tensor parallel on a 2-D mesh; HSDP is supported), models whose decoder layers cannot be discovered, a decoder layer object reused across layers, and a module that holds the decoder layers while owning parameters of its own.
 - Speed up ``mtq.quantize`` on FSDP2-sharded fused-MoE models. Promoting static-block weight quantizers gathered each expert's slice of the fused weight across ranks even though only quantizer state is read, adding a collective per expert to calibration.

modelopt/onnx/export/nvfp4_exporter.py

@@ -331,7 +331,7 @@ def post_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto:
         initializer_indices = {
             initializer.name: idx for idx, initializer in enumerate(graph.initializer)
         }
-        value_info_map = {vi.name: vi for vi in graph.value_info}
+        value_info_map = {vi.name: vi for vi in [*graph.value_info, *graph.output]}
         graph_inputs = {inp.name for inp in graph.input}
         cast_output_cache: dict[tuple[str, str], str] = {}
 
@@ -351,11 +351,15 @@ def _cast_input_dtypes(node: onnx.NodeProto, precision_dtype: str):
                 assert maybe_matmul.op_type == "MatMul"
                 node = maybe_matmul
 
-            # Create Cast nodes for each input of the target node except bias
-            for i, input_name in enumerate(node.input[:2]):
+            precision_onnx_dtype = onnx_dtype_map[precision_dtype]
+            cast_output_suffix = "bf16" if precision_dtype == "BFloat16" else "f16"
+
+            compute_inputs = node.input[:3] if node.op_type == "Gemm" else node.input[:2]
+            for i, input_name in enumerate(compute_inputs):
+                if not input_name:
+                    continue
                 cast_output_name = cast_output_cache.get((input_name, precision_dtype))
                 if cast_output_name is None:
-                    cast_output_suffix = "bf16" if precision_dtype == "BFloat16" else "f16"
                     cast_output_name = f"{input_name}_{cast_output_suffix}"
                     cast_output_cache[(input_name, precision_dtype)] = cast_output_name
 
@@ -364,7 +368,7 @@ def _cast_input_dtypes(node: onnx.NodeProto, precision_dtype: str):
                         "Cast",
                         inputs=[input_name],  # Original input of the target node
                         outputs=[cast_output_name],
-                        to=onnx_dtype_map[precision_dtype],  # Cast to FP16/BF16
+                        to=precision_onnx_dtype,
                     )
 
                     # Insert the Cast node into the graph
@@ -373,6 +377,35 @@ def _cast_input_dtypes(node: onnx.NodeProto, precision_dtype: str):
                 # Update the target node input to use the cast node output
                 node.input[i] = cast_output_name
 
+            for i, output_name in enumerate(node.output):
+                output_value_info = value_info_map.get(output_name)
+                if output_value_info is None:
+                    continue
+
+                output_dtype = output_value_info.type.tensor_type.elem_type
+                if precision_dtype == "BFloat16" or output_dtype == precision_onnx_dtype:
+                    output_value_info.type.tensor_type.elem_type = precision_onnx_dtype
+                    continue
+
+                precision_output_name = f"{output_name}_{cast_output_suffix}_output"
+                precision_output_value_info = onnx.ValueInfoProto()
+                precision_output_value_info.CopyFrom(output_value_info)
+                precision_output_value_info.name = precision_output_name
+                precision_output_value_info.type.tensor_type.elem_type = precision_onnx_dtype
+                graph.value_info.append(precision_output_value_info)
+                value_info_map[precision_output_name] = precision_output_value_info
+                node.output[i] = precision_output_name
+                graph.node.extend(
+                    [
+                        onnx.helper.make_node(
+                            "Cast",
+                            inputs=[precision_output_name],
+                            outputs=[output_name],
+                            to=output_dtype,
+                        )
+                    ]
+                )
+
         precision_dtype = _get_precision_dtype()
         logger.debug(f"Using precision dtype: {precision_dtype}")
 

modelopt/torch/_deploy/utils/torch_onnx.py

@@ -35,6 +35,7 @@
 from torch.nn.parallel import DataParallel, DistributedDataParallel
 
 from modelopt.onnx.autocast.convert import convert_to_f16
+from modelopt.onnx.autocast.graphsanitizer import GraphSanitizer
 from modelopt.onnx.export import (
     FP8QuantExporter,
     INT4QuantExporter,
@@ -52,9 +53,11 @@
     fold_qdq_scale_fp16_to_fp32_casts,
     get_input_names,
     get_input_shapes,
+    get_min_opset_for_precisions,
     get_node_names,
     get_output_names,
     get_output_shapes,
+    get_qdq_precisions,
     infer_shapes,
     is_model_too_large_for_protobuf,
     remove_node_training_mode,
@@ -507,10 +510,10 @@ def get_onnx_bytes_and_metadata(
             `torch.onnx.export <https://pytorch.org/docs/stable/onnx.html#torch.onnx.export>`_.
         onnx_opset: The onnx opset version to use for exporting the model.
         dq_only: If True, the exported onnx model is converted to a dq_only model.
-        weights_dtype: Requested high-precision dtype for exported weights. For an FP8 model,
-            ``"bf16"`` is accepted only when every floating parameter is already BF16. This is
-            a weight-focused no-op, not a graph-wide conversion: floating buffers are not
-            considered for eligibility and may preserve higher-precision regions.
+        weights_dtype: Requested high-precision dtype for exported weights. For an FP8 or NVFP4
+            model, ``"bf16"`` is accepted only when every floating parameter is already BF16.
+            This is a weight-focused no-op, not a graph-wide conversion: floating buffers are
+            not considered for eligibility and may preserve higher-precision regions.
 
     Returns:
         bytes: Onnx model in bytes.
@@ -539,12 +542,19 @@ def get_onnx_bytes_and_metadata(
     uses_fp8 = is_fp8_quantized(model)
     uses_int8 = is_int8_quantized(model)
     uses_other_unsupported_quantizer = is_int4_quantized(model) or uses_mxfp8 or uses_int8
+    is_bf16_fp4_noop = (
+        weights_dtype == "bf16"
+        and source_parameter_dtypes == {torch.bfloat16}
+        and uses_fp4
+        and not (uses_fp8 or uses_other_unsupported_quantizer)
+    )
     is_bf16_fp8_noop = (
         weights_dtype == "bf16"
         and source_parameter_dtypes == {torch.bfloat16}
         and uses_fp8
         and not (uses_fp4 or uses_other_unsupported_quantizer)
     )
+    is_bf16_quantized_noop = is_bf16_fp4_noop or is_bf16_fp8_noop
 
     # 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
@@ -590,16 +600,21 @@ def get_onnx_bytes_and_metadata(
         )
         return onnx_model.to_bytes(), model_metadata
 
-    if weights_dtype == "fp16" and uses_fp8 and torch.bfloat16 in source_parameter_dtypes:
+    if (
+        weights_dtype == "fp16"
+        and (uses_fp4 or uses_fp8)
+        and torch.bfloat16 in source_parameter_dtypes
+    ):
+        quantization_format = "NVFP4" if uses_fp4 else "FP8"
         raise ValueError(
-            "Converting a BF16 FP8 ONNX graph to FP16 is not supported yet "
+            f"Converting a BF16 {quantization_format} ONNX graph to FP16 is not supported yet "
             f"(source parameter dtypes: {source_parameter_dtype_names})"
         )
 
     if (
         weights_dtype == "bf16"
-        and (uses_fp8 or uses_other_unsupported_quantizer)
-        and not is_bf16_fp8_noop
+        and (uses_fp4 or uses_fp8 or uses_other_unsupported_quantizer)
+        and not is_bf16_quantized_noop
     ):
         raise ValueError(
             "Converting a quantized ONNX graph to BF16 is not supported yet "
@@ -658,11 +673,17 @@ def get_onnx_bytes_and_metadata(
 
     onnx_opt_graph = quantize_weights(model, onnx_opt_graph)
 
+    if uses_fp4:
+        qdq_min_opset = get_min_opset_for_precisions(get_qdq_precisions(onnx_opt_graph))
+        opset_sanitizer = GraphSanitizer(onnx_opt_graph, min_opset=qdq_min_opset)
+        opset_sanitizer.convert_opset()
+        onnx_opt_graph = opset_sanitizer.model
+
     if dq_only:
         onnx_opt_graph = qdq_to_dq(onnx_opt_graph)
 
-    if weights_dtype in ["fp16", "bf16"] and not is_bf16_fp8_noop:
-        if uses_other_unsupported_quantizer or uses_fp8:
+    if weights_dtype in ["fp16", "bf16"] and not is_bf16_quantized_noop:
+        if weights_dtype == "fp16" and (uses_fp4 or uses_other_unsupported_quantizer or uses_fp8):
             onnx_opt_graph = convert_float_to_float16(
                 onnx_opt_graph,
                 keep_io_types=False,

tests/unit/torch/deploy/utils/test_torch_onnx_utils.py

@@ -26,6 +26,7 @@
 import torch.nn as nn
 from _test_utils.torch.deploy.lib_test_models import BaseDeployModel, get_deploy_models
 
+import modelopt.torch._deploy.utils.torch_onnx as torch_onnx
 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 (
@@ -299,6 +300,70 @@ def test_fp8_export_rejects_unsupported_dtype_conversion(
     assert not any(tmp_path.iterdir())
 
 
+@pytest.mark.parametrize(
+    ("source_dtype", "weights_dtype"),
+    [(torch.bfloat16, "fp16"), (torch.float32, "bf16")],
+    ids=["bf16-to-fp16", "fp32-to-bf16"],
+)
+def test_nvfp4_export_rejects_unsupported_dtype_conversion(
+    source_dtype, weights_dtype, monkeypatch, tmp_path
+):
+    monkeypatch.setattr(tempfile, "tempdir", str(tmp_path))
+    monkeypatch.setattr(torch_onnx, "is_fp4_quantized", lambda _: True)
+    model = nn.Linear(4, 4).eval().to(source_dtype)
+
+    with pytest.raises(
+        ValueError,
+        match=rf"Converting .* to {weights_dtype.upper()}.*source parameter dtypes: {source_dtype}",
+    ):
+        get_onnx_bytes_and_metadata(
+            model,
+            (torch.ones(1, 4, dtype=source_dtype),),
+            weights_dtype=weights_dtype,
+        )
+    assert not any(tmp_path.iterdir())
+
+
+@pytest.mark.parametrize(
+    ("source_dtype", "weights_dtype", "expected_calls"),
+    [
+        (torch.float32, "fp16", ["onnxconverter"]),
+        (torch.bfloat16, "bf16", []),
+    ],
+    ids=["fp32-to-fp16", "bf16-noop"],
+)
+def test_nvfp4_export_selects_precision_converter(
+    source_dtype, weights_dtype, expected_calls, monkeypatch
+):
+    calls = []
+
+    def record_onnxconverter(model, **kwargs):
+        calls.append("onnxconverter")
+        return model
+
+    def record_autocast(model, **kwargs):
+        calls.append("autocast")
+        return model
+
+    monkeypatch.setattr(torch_onnx, "is_fp4_quantized", lambda _: True)
+    monkeypatch.setattr(
+        torch_onnx, "configure_linear_module_onnx_quantizers", lambda _: nullcontext()
+    )
+    monkeypatch.setattr(torch_onnx, "quantize_weights", lambda _, graph: graph)
+    monkeypatch.setattr(torch_onnx, "convert_float_to_float16", record_onnxconverter)
+    monkeypatch.setattr(torch_onnx, "convert_to_f16", record_autocast)
+
+    model = nn.Linear(4, 4).eval().to(source_dtype)
+    get_onnx_bytes_and_metadata(
+        model,
+        (torch.ones(1, 4, dtype=source_dtype),),
+        weights_dtype=weights_dtype,
+        onnx_opset=23,
+    )
+
+    assert calls == expected_calls
+
+
 class SingleArgModel(nn.Module):
     def forward(self, x: torch.Tensor):
         return torch.add(x, x) - x