From 51493d4e8fdc3fc1eec443b1d5c59f7a0369f4bc Mon Sep 17 00:00:00 2001
From: ajrasane <131806219+ajrasane@users.noreply.github.com>
Date: Fri, 11 Sep 2026 20:59:28 +0000
Subject: [PATCH 1/4] [6721556] Fix NVFP4 FP16 ONNX conversion

Co-Authored-By: Codex <codex@openai.com>
Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
---
 CHANGELOG.rst                                 |  1 +
 modelopt/torch/_deploy/utils/torch_onnx.py    | 11 ++--
 .../deploy/utils/test_torch_onnx_utils.py     | 52 +++++++++++++++++++
 3 files changed, 61 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 3e10c786753..012c9748ee6 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -25,6 +25,7 @@ Changelog
 
 **Bug Fixes**
 
+- Fix FP16 ONNX exports of NVFP4-quantized models that failed TensorRT parsing due to mixed-precision elementwise inputs.
 - 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.
 - Add FP8 and INT8 recipes that quantize timm ResNet shortcut inputs immediately before residual adds. The torch ONNX example now accepts PTQ and AutoQuantize recipes through ``--recipe`` and uses ``--qformat`` when no recipe is provided. ResNet supports only FP8 and INT8 because TensorRT has limited convolution kernel support; AutoQuantize and other quantization formats are no longer supported for ResNet.
diff --git a/modelopt/torch/_deploy/utils/torch_onnx.py b/modelopt/torch/_deploy/utils/torch_onnx.py
index b217f1188b3..77b837ee18b 100644
--- a/modelopt/torch/_deploy/utils/torch_onnx.py
+++ b/modelopt/torch/_deploy/utils/torch_onnx.py
@@ -590,9 +590,14 @@ 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})"
         )
 
@@ -662,7 +667,7 @@ def get_onnx_bytes_and_metadata(
         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 == "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,
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 2fd1a1c8504..f9c49f00566 100644
--- a/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py
+++ b/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,57 @@ def test_fp8_export_rejects_unsupported_dtype_conversion(
     assert not any(tmp_path.iterdir())
 
 
+def test_nvfp4_export_rejects_bf16_to_fp16(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().bfloat16()
+
+    with pytest.raises(
+        ValueError,
+        match=r"Converting a BF16 NVFP4 ONNX graph to FP16.*torch.bfloat16",
+    ):
+        get_onnx_bytes_and_metadata(
+            model,
+            (torch.ones(1, 4, dtype=torch.bfloat16),),
+            weights_dtype="fp16",
+        )
+    assert not any(tmp_path.iterdir())
+
+
+@pytest.mark.parametrize(
+    ("weights_dtype", "expected_converter"),
+    [("fp16", "onnxconverter"), ("bf16", "autocast")],
+)
+def test_nvfp4_export_selects_precision_converter(weights_dtype, expected_converter, 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()
+    get_onnx_bytes_and_metadata(
+        model,
+        (torch.ones(1, 4),),
+        weights_dtype=weights_dtype,
+        onnx_opset=23,
+    )
+
+    assert calls == [expected_converter]
+
+
 class SingleArgModel(nn.Module):
     def forward(self, x: torch.Tensor):
         return torch.add(x, x) - x

From 1ffc34f89a28ba846dcea971d0af45cf6eccd868 Mon Sep 17 00:00:00 2001
From: ajrasane <131806219+ajrasane@users.noreply.github.com>
Date: Fri, 11 Sep 2026 22:48:18 +0000
Subject: [PATCH 2/4] fix: preserve NVFP4 export precision metadata

Co-Authored-By: Codex <codex@openai.com>
Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
---
 modelopt/onnx/export/nvfp4_exporter.py        |  8 +-
 modelopt/torch/_deploy/utils/torch_onnx.py    | 30 ++++--
 .../deploy/utils/test_torch_onnx_utils.py     | 35 ++++---
 .../quantization/test_onnx_export_cpu.py      | 93 ++++++++++++++++++-
 4 files changed, 142 insertions(+), 24 deletions(-)

diff --git a/modelopt/onnx/export/nvfp4_exporter.py b/modelopt/onnx/export/nvfp4_exporter.py
index 338e2725b14..512d051d108 100644
--- a/modelopt/onnx/export/nvfp4_exporter.py
+++ b/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] = {}
 
@@ -373,6 +373,12 @@ 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 output_name in node.output:
+                if output_name in value_info_map:
+                    value_info_map[output_name].type.tensor_type.elem_type = onnx_dtype_map[
+                        precision_dtype
+                    ]
+
         precision_dtype = _get_precision_dtype()
         logger.debug(f"Using precision dtype: {precision_dtype}")
 
diff --git a/modelopt/torch/_deploy/utils/torch_onnx.py b/modelopt/torch/_deploy/utils/torch_onnx.py
index 77b837ee18b..7b797d2299e 100644
--- a/modelopt/torch/_deploy/utils/torch_onnx.py
+++ b/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
@@ -603,8 +613,8 @@ def get_onnx_bytes_and_metadata(
 
     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 "
@@ -663,10 +673,16 @@ 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 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,
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 f9c49f00566..af61bcce553 100644
--- a/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py
+++ b/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py
@@ -300,28 +300,41 @@ def test_fp8_export_rejects_unsupported_dtype_conversion(
     assert not any(tmp_path.iterdir())
 
 
-def test_nvfp4_export_rejects_bf16_to_fp16(monkeypatch, tmp_path):
+@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().bfloat16()
+    model = nn.Linear(4, 4).eval().to(source_dtype)
 
     with pytest.raises(
         ValueError,
-        match=r"Converting a BF16 NVFP4 ONNX graph to FP16.*torch.bfloat16",
+        match=rf"Converting .* to {weights_dtype.upper()}.*source parameter dtypes: {source_dtype}",
     ):
         get_onnx_bytes_and_metadata(
             model,
-            (torch.ones(1, 4, dtype=torch.bfloat16),),
-            weights_dtype="fp16",
+            (torch.ones(1, 4, dtype=source_dtype),),
+            weights_dtype=weights_dtype,
         )
     assert not any(tmp_path.iterdir())
 
 
 @pytest.mark.parametrize(
-    ("weights_dtype", "expected_converter"),
-    [("fp16", "onnxconverter"), ("bf16", "autocast")],
+    ("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(weights_dtype, expected_converter, monkeypatch):
+def test_nvfp4_export_selects_precision_converter(
+    source_dtype, weights_dtype, expected_calls, monkeypatch
+):
     calls = []
 
     def record_onnxconverter(model, **kwargs):
@@ -340,15 +353,15 @@ def record_autocast(model, **kwargs):
     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()
+    model = nn.Linear(4, 4).eval().to(source_dtype)
     get_onnx_bytes_and_metadata(
         model,
-        (torch.ones(1, 4),),
+        (torch.ones(1, 4, dtype=source_dtype),),
         weights_dtype=weights_dtype,
         onnx_opset=23,
     )
 
-    assert calls == [expected_converter]
+    assert calls == expected_calls
 
 
 class SingleArgModel(nn.Module):
diff --git a/tests/unit/torch/quantization/test_onnx_export_cpu.py b/tests/unit/torch/quantization/test_onnx_export_cpu.py
index ce2ef626d63..3daa2ba6dab 100644
--- a/tests/unit/torch/quantization/test_onnx_export_cpu.py
+++ b/tests/unit/torch/quantization/test_onnx_export_cpu.py
@@ -35,6 +35,7 @@
 from modelopt.onnx.export import NVFP4QuantExporter
 from modelopt.onnx.export.nvfp4_exporter import _encode_nvfp4_block_scale
 from modelopt.onnx.quantization.qdq_utils import fp4qdq_to_2dq
+from modelopt.torch._deploy.utils import OnnxBytes, get_onnx_bytes_and_metadata
 from modelopt.torch.quantization.qtensor import NVFP4QTensor
 from modelopt.torch.quantization.utils import is_quantized_linear
 
@@ -59,7 +60,17 @@ def test_onnx_export_cpu(model_cls, num_bits, per_channel_quantization, constant
     )
 
 
-def test_nvfp4_exported_onnx_is_topologically_sorted(monkeypatch):
+class _NVFP4LinearWithExplicitBias(torch.nn.Module):
+    def __init__(self, dtype):
+        super().__init__()
+        self.linear = torch.nn.Linear(16, 16, bias=False, dtype=dtype)
+        self.bias = torch.nn.Parameter(torch.ones(16, dtype=dtype))
+
+    def forward(self, inputs):
+        return self.linear(inputs) + self.bias
+
+
+def _make_cpu_nvfp4_model(monkeypatch, model, sample_input, disable_input_quantizers=False):
     def forward_loop(model):
         model(sample_input)
 
@@ -67,17 +78,23 @@ def cpu_dynamic_block_quantize(inputs, *args):
         return inputs
 
     monkeypatch.setattr(tensor_quant, "dynamic_block_quantize_op", cpu_dynamic_block_quantize)
-
-    model = SimpleLinear().eval()
-    sample_input = model.get_input()
     model = mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, forward_loop=forward_loop)
 
     for module in model.modules():
         assert not isinstance(module, torch.nn.Linear) or is_quantized_linear(module)
         if isinstance(module, torch.nn.Linear):
-            module.input_quantizer.disable()
+            if disable_input_quantizers:
+                module.input_quantizer.disable()
             module.weight_quantizer._onnx_quantizer_type = "static"
 
+    return model
+
+
+def test_nvfp4_exported_onnx_is_topologically_sorted(monkeypatch):
+    model = SimpleLinear().eval()
+    sample_input = model.get_input()
+    model = _make_cpu_nvfp4_model(monkeypatch, model, sample_input, disable_input_quantizers=True)
+
     buffer = io.BytesIO()
     if "enable_onnx_checker" in inspect.signature(torch.onnx.export).parameters:
         kwargs = {"enable_onnx_checker": False}
@@ -105,6 +122,72 @@ def cpu_dynamic_block_quantize(inputs, *args):
     onnx.checker.check_model(converted_model)
 
 
+@pytest.mark.parametrize(
+    ("source_dtype", "weights_dtype", "expected_dtype"),
+    [
+        (torch.float32, "fp16", TensorProto.FLOAT16),
+        (torch.bfloat16, "bf16", TensorProto.BFLOAT16),
+    ],
+    ids=["fp32-to-fp16", "bf16-noop"],
+)
+def test_nvfp4_deploy_export_has_consistent_elementwise_types(
+    monkeypatch, source_dtype, weights_dtype, expected_dtype
+):
+    model = _NVFP4LinearWithExplicitBias(source_dtype).eval()
+    sample_input = torch.ones(1, 2, 16, dtype=source_dtype)
+    model = _make_cpu_nvfp4_model(monkeypatch, model, sample_input)
+
+    onnx_bytes, _ = get_onnx_bytes_and_metadata(
+        model,
+        (sample_input,),
+        weights_dtype=weights_dtype,
+    )
+    exported_model = onnx.load_model_from_string(
+        OnnxBytes.from_bytes(onnx_bytes).get_onnx_model_file_bytes()
+    )
+
+    assert utils.get_opset_version(exported_model) >= 23
+    onnx.checker.check_model(exported_model, full_check=True)
+    assert any(
+        node.op_type == "DequantizeLinear"
+        and any(attribute.name == "block_size" for attribute in node.attribute)
+        for node in exported_model.graph.node
+    )
+    assert any(node.op_type == "TRT_FP4DynamicQuantize" for node in exported_model.graph.node)
+
+    inferred_model = onnx.shape_inference.infer_shapes(exported_model, strict_mode=True)
+    tensor_types = {
+        initializer.name: initializer.data_type for initializer in inferred_model.graph.initializer
+    }
+    for value in [
+        *inferred_model.graph.input,
+        *inferred_model.graph.value_info,
+        *inferred_model.graph.output,
+    ]:
+        if value.type.HasField("tensor_type"):
+            tensor_types[value.name] = value.type.tensor_type.elem_type
+
+    floating_types = {TensorProto.FLOAT, TensorProto.FLOAT16, TensorProto.BFLOAT16}
+    elementwise_nodes = [
+        node
+        for node in inferred_model.graph.node
+        if node.op_type in {"Add", "Sub", "Mul", "Div", "Pow"}
+    ]
+    assert any(node.op_type == "Add" for node in elementwise_nodes)
+    for node in elementwise_nodes:
+        input_types = [tensor_types[input_name] for input_name in node.input]
+        assert len(set(input_types) & floating_types) <= 1, (
+            node.name,
+            [TensorProto.DataType.Name(input_type) for input_type in input_types],
+        )
+
+    add_node = next(node for node in elementwise_nodes if node.op_type == "Add")
+    assert [tensor_types[input_name] for input_name in add_node.input] == [
+        expected_dtype,
+        expected_dtype,
+    ]
+
+
 @pytest.mark.parametrize(
     ("convert", "deprecated"),
     [

From c97ae4a5d2b566f361ba73d497e5951ac358a5ab Mon Sep 17 00:00:00 2001
From: ajrasane <131806219+ajrasane@users.noreply.github.com>
Date: Sat, 12 Sep 2026 04:42:35 +0000
Subject: [PATCH 3/4] fix: preserve NVFP4 FP32 graph boundaries

Co-Authored-By: Codex <codex@openai.com>
Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
---
 CHANGELOG.rst                                 |  2 +-
 modelopt/onnx/export/nvfp4_exporter.py        | 43 +++++++++++++---
 .../quantization/test_onnx_export_cpu.py      | 49 ++++++++++++++++++-
 3 files changed, 84 insertions(+), 10 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 012c9748ee6..c56614fd579 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -25,7 +25,7 @@ Changelog
 
 **Bug Fixes**
 
-- Fix FP16 ONNX exports of NVFP4-quantized models that failed TensorRT parsing due to mixed-precision elementwise inputs.
+- 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 ``--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.
 - Add FP8 and INT8 recipes that quantize timm ResNet shortcut inputs immediately before residual adds. The torch ONNX example now accepts PTQ and AutoQuantize recipes through ``--recipe`` and uses ``--qformat`` when no recipe is provided. ResNet supports only FP8 and INT8 because TensorRT has limited convolution kernel support; AutoQuantize and other quantization formats are no longer supported for ResNet.
diff --git a/modelopt/onnx/export/nvfp4_exporter.py b/modelopt/onnx/export/nvfp4_exporter.py
index 512d051d108..cd6768c5063 100644
--- a/modelopt/onnx/export/nvfp4_exporter.py
+++ b/modelopt/onnx/export/nvfp4_exporter.py
@@ -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,11 +377,34 @@ 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 output_name in node.output:
-                if output_name in value_info_map:
-                    value_info_map[output_name].type.tensor_type.elem_type = onnx_dtype_map[
-                        precision_dtype
+            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}")
diff --git a/tests/unit/torch/quantization/test_onnx_export_cpu.py b/tests/unit/torch/quantization/test_onnx_export_cpu.py
index 3daa2ba6dab..e3d88cc0c74 100644
--- a/tests/unit/torch/quantization/test_onnx_export_cpu.py
+++ b/tests/unit/torch/quantization/test_onnx_export_cpu.py
@@ -125,10 +125,11 @@ def test_nvfp4_exported_onnx_is_topologically_sorted(monkeypatch):
 @pytest.mark.parametrize(
     ("source_dtype", "weights_dtype", "expected_dtype"),
     [
+        (torch.float32, "fp32", TensorProto.FLOAT),
         (torch.float32, "fp16", TensorProto.FLOAT16),
         (torch.bfloat16, "bf16", TensorProto.BFLOAT16),
     ],
-    ids=["fp32-to-fp16", "bf16-noop"],
+    ids=["fp32-preserved", "fp32-to-fp16", "bf16-noop"],
 )
 def test_nvfp4_deploy_export_has_consistent_elementwise_types(
     monkeypatch, source_dtype, weights_dtype, expected_dtype
@@ -188,6 +189,52 @@ def test_nvfp4_deploy_export_has_consistent_elementwise_types(
     ]
 
 
+@pytest.mark.parametrize(
+    ("source_dtype", "weights_dtype", "expected_gemm_dtype", "expected_output_dtype"),
+    [
+        (torch.float32, "fp32", TensorProto.FLOAT16, TensorProto.FLOAT),
+        (torch.float32, "fp16", TensorProto.FLOAT16, TensorProto.FLOAT16),
+        (torch.bfloat16, "bf16", TensorProto.BFLOAT16, TensorProto.BFLOAT16),
+    ],
+    ids=["fp32-preserved", "fp32-to-fp16", "bf16-noop"],
+)
+def test_nvfp4_deploy_export_has_consistent_gemm_types(
+    monkeypatch, source_dtype, weights_dtype, expected_gemm_dtype, expected_output_dtype
+):
+    model = torch.nn.Linear(16, 16, dtype=source_dtype).eval()
+    sample_input = torch.ones(2, 16, dtype=source_dtype)
+    model = _make_cpu_nvfp4_model(monkeypatch, model, sample_input)
+
+    onnx_bytes, _ = get_onnx_bytes_and_metadata(
+        model,
+        (sample_input,),
+        weights_dtype=weights_dtype,
+    )
+    exported_model = onnx.load_model_from_string(
+        OnnxBytes.from_bytes(onnx_bytes).get_onnx_model_file_bytes()
+    )
+
+    assert utils.get_opset_version(exported_model) >= 23
+    onnx.checker.check_model(exported_model, full_check=True)
+    assert any(node.op_type == "TRT_FP4DynamicQuantize" for node in exported_model.graph.node)
+    inferred_model = onnx.shape_inference.infer_shapes(exported_model, strict_mode=True)
+    tensor_types = {
+        initializer.name: initializer.data_type for initializer in inferred_model.graph.initializer
+    }
+    for value in [
+        *inferred_model.graph.input,
+        *inferred_model.graph.value_info,
+        *inferred_model.graph.output,
+    ]:
+        if value.type.HasField("tensor_type"):
+            tensor_types[value.name] = value.type.tensor_type.elem_type
+
+    gemm_node = next(node for node in inferred_model.graph.node if node.op_type == "Gemm")
+    assert [tensor_types[input_name] for input_name in gemm_node.input] == [expected_gemm_dtype] * 3
+    assert tensor_types[gemm_node.output[0]] == expected_gemm_dtype
+    assert tensor_types[inferred_model.graph.output[0].name] == expected_output_dtype
+
+
 @pytest.mark.parametrize(
     ("convert", "deprecated"),
     [

From 7038eb84d646cbeecdf14806d3074fc7e97b4786 Mon Sep 17 00:00:00 2001
From: ajrasane <131806219+ajrasane@users.noreply.github.com>
Date: Fri, 18 Sep 2026 16:49:43 +0000
Subject: [PATCH 4/4] fix: preserve mixed NVFP4 precision boundaries

Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
---
 modelopt/onnx/export/nvfp4_exporter.py        | 19 ++++----
 .../quantization/test_onnx_export_cpu.py      | 46 +++++++++++++++++++
 2 files changed, 54 insertions(+), 11 deletions(-)

diff --git a/modelopt/onnx/export/nvfp4_exporter.py b/modelopt/onnx/export/nvfp4_exporter.py
index cd6768c5063..932f27791dd 100644
--- a/modelopt/onnx/export/nvfp4_exporter.py
+++ b/modelopt/onnx/export/nvfp4_exporter.py
@@ -335,14 +335,10 @@ def post_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto:
         graph_inputs = {inp.name for inp in graph.input}
         cast_output_cache: dict[tuple[str, str], str] = {}
 
-        def _get_precision_dtype() -> str:
-            # Check initializers to determine the precision of the weights
-            precision_dtype = "Half"
-            for initializer in graph.initializer:
-                if initializer.data_type == 16:
-                    precision_dtype = "BFloat16"
-                    break  # Assuming all weights are of the same precision
-            return precision_dtype
+        def _get_precision_dtype(weight_initializer: onnx.TensorProto) -> str:
+            return (
+                "BFloat16" if weight_initializer.data_type == onnx.TensorProto.BFLOAT16 else "Half"
+            )
 
         def _cast_input_dtypes(node: onnx.NodeProto, precision_dtype: str):
             # Change the input types to match weight precision (precision_dtype)
@@ -383,6 +379,8 @@ def _cast_input_dtypes(node: onnx.NodeProto, precision_dtype: str):
                     continue
 
                 output_dtype = output_value_info.type.tensor_type.elem_type
+                # TRT_FP4QDQ leaves native-BF16 outputs annotated FLOAT; real FP32 boundaries are
+                # explicit Casts. FP16 can convert FP32 graphs, so it restores implicit boundaries.
                 if precision_dtype == "BFloat16" or output_dtype == precision_onnx_dtype:
                     output_value_info.type.tensor_type.elem_type = precision_onnx_dtype
                     continue
@@ -406,9 +404,6 @@ def _cast_input_dtypes(node: onnx.NodeProto, precision_dtype: str):
                     ]
                 )
 
-        precision_dtype = _get_precision_dtype()
-        logger.debug(f"Using precision dtype: {precision_dtype}")
-
         fp4_qdq_nodes = [node for node in graph.node if node.op_type == "TRT_FP4QDQ"]
         logger.debug(f"Found {len(fp4_qdq_nodes)} FP4QDQ nodes to convert")
 
@@ -416,6 +411,8 @@ def _cast_input_dtypes(node: onnx.NodeProto, precision_dtype: str):
             idx = initializer_indices.get(node.input[0])
             assert idx is not None, f"Initializer for weight '{node.input[0]}' not found."
             initializers_to_delete.append(graph.initializer[idx].name)
+            precision_dtype = _get_precision_dtype(graph.initializer[idx])
+            logger.debug(f"Using precision dtype {precision_dtype} for {node.input[0]}")
 
             # Retrieve compressed data from node attributes
             block_size = node.attribute[0].i
diff --git a/tests/unit/torch/quantization/test_onnx_export_cpu.py b/tests/unit/torch/quantization/test_onnx_export_cpu.py
index e3d88cc0c74..2cd82ee1b0c 100644
--- a/tests/unit/torch/quantization/test_onnx_export_cpu.py
+++ b/tests/unit/torch/quantization/test_onnx_export_cpu.py
@@ -70,6 +70,18 @@ def forward(self, inputs):
         return self.linear(inputs) + self.bias
 
 
+class _NVFP4MixedPrecisionLinear(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.fp32_linear = torch.nn.Linear(16, 16, dtype=torch.float32)
+        self.bf16_linear = torch.nn.Linear(16, 16, dtype=torch.bfloat16)
+
+    def forward(self, inputs):
+        fp32_output = self.fp32_linear(inputs)
+        bf16_output = self.bf16_linear(inputs.to(torch.bfloat16)).float()
+        return fp32_output + bf16_output
+
+
 def _make_cpu_nvfp4_model(monkeypatch, model, sample_input, disable_input_quantizers=False):
     def forward_loop(model):
         model(sample_input)
@@ -189,6 +201,40 @@ def test_nvfp4_deploy_export_has_consistent_elementwise_types(
     ]
 
 
+def test_nvfp4_deploy_export_preserves_mixed_precision_boundaries(monkeypatch):
+    model = _NVFP4MixedPrecisionLinear().eval()
+    sample_input = torch.ones(1, 16)
+    model = _make_cpu_nvfp4_model(monkeypatch, model, sample_input)
+
+    onnx_bytes, _ = get_onnx_bytes_and_metadata(
+        model,
+        (sample_input,),
+        weights_dtype="fp32",
+    )
+    exported_model = onnx.load_model_from_string(
+        OnnxBytes.from_bytes(onnx_bytes).get_onnx_model_file_bytes()
+    )
+
+    onnx.checker.check_model(exported_model, full_check=True)
+    inferred_model = onnx.shape_inference.infer_shapes(exported_model, strict_mode=True)
+    tensor_types = {
+        initializer.name: initializer.data_type for initializer in inferred_model.graph.initializer
+    }
+    for value in [
+        *inferred_model.graph.input,
+        *inferred_model.graph.value_info,
+        *inferred_model.graph.output,
+    ]:
+        if value.type.HasField("tensor_type"):
+            tensor_types[value.name] = value.type.tensor_type.elem_type
+
+    add_node = next(node for node in inferred_model.graph.node if node.op_type == "Add")
+    assert [tensor_types[input_name] for input_name in add_node.input] == [
+        TensorProto.FLOAT,
+        TensorProto.FLOAT,
+    ]
+
+
 @pytest.mark.parametrize(
     ("source_dtype", "weights_dtype", "expected_gemm_dtype", "expected_output_dtype"),
     [