From 69b3580db6ed186cae97c76f84e45fc6f0ffb576 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Tue, 27 Oct 2026 10:06:55 -0700
Subject: [PATCH 1/7] Test Megatron-Bridge quantized HF export against the
 Megatron exporter

Bump the Megatron-Bridge commit CI overlays on the NeMo container to
NVIDIA-NeMo/Megatron-Bridge#6657, which adds save_hf_pretrained_modelopt, and
check that it writes the same weights, scales and quantization config as
export_mcore_gpt_to_hf.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .../actions/megatron-from-source/action.yml   |   2 +-
 examples/megatron_bridge/README.md            |   2 +-
 .../export/test_quantized_mbridge_export.py   | 120 ++++++++++++++++++
 3 files changed, 122 insertions(+), 2 deletions(-)
 create mode 100644 tests/gpu_megatron/torch/export/test_quantized_mbridge_export.py

diff --git a/.github/actions/megatron-from-source/action.yml b/.github/actions/megatron-from-source/action.yml
index 50eef77d41d..d8d94592cb1 100644
--- a/.github/actions/megatron-from-source/action.yml
+++ b/.github/actions/megatron-from-source/action.yml
@@ -6,7 +6,7 @@ description: >-
 inputs:
   commit:
     description: Megatron-Bridge commit. Keep in sync with examples/megatron_bridge/README.md.
-    default: 1d9caf7324e608aca7181de40be26a6387976210
+    default: 0652ea797ab629e6451c64e1c92f976501b3b8b3
 
 runs:
   using: composite
diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md
index 81b6f11ff7e..975db49a61e 100644
--- a/examples/megatron_bridge/README.md
+++ b/examples/megatron_bridge/README.md
@@ -51,7 +51,7 @@ Some features (e.g. [`--vllm_fakequant_export`](#evaluating-a-qad-checkpoint-in-
 ```bash
 export MBRIDGE_DIR=${PWD}/Megatron-Bridge
 git clone https://github.com/NVIDIA-NeMo/Megatron-Bridge.git ${MBRIDGE_DIR}
-git -C ${MBRIDGE_DIR} checkout 1d9caf7324e608aca7181de40be26a6387976210
+git -C ${MBRIDGE_DIR} checkout 0652ea797ab629e6451c64e1c92f976501b3b8b3
 git -C ${MBRIDGE_DIR} submodule update --init 3rdparty/Megatron-LM
 ```
 
diff --git a/tests/gpu_megatron/torch/export/test_quantized_mbridge_export.py b/tests/gpu_megatron/torch/export/test_quantized_mbridge_export.py
new file mode 100644
index 00000000000..1bcb7fc9f5a
--- /dev/null
+++ b/tests/gpu_megatron/torch/export/test_quantized_mbridge_export.py
@@ -0,0 +1,120 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Quantized HF export through Megatron-Bridge matches ModelOpt's Megatron exporter."""
+
+import fnmatch
+import json
+import os
+from functools import partial
+from pathlib import Path
+
+import pytest
+import torch
+from _test_utils.torch.transformers_models import (
+    create_tiny_llama_dir,
+    create_tiny_qwen3_dir,
+    create_tiny_qwen3_moe_dir,
+)
+from safetensors.torch import load_file
+
+import modelopt.torch.quantization as mtq
+from modelopt.torch.export import export_mcore_gpt_to_hf
+from modelopt.torch.utils.plugins.mbridge import load_mbridge_model_from_hf, use_moe_grouped_gemm
+from modelopt.torch.utils.plugins.megatron_generate import megatron_prefill
+
+QUANT_CFGS = {
+    "nvfp4": mtq.NVFP4_DEFAULT_CFG,
+    "fp8": mtq.FP8_DEFAULT_CFG,
+    "w4a16_nvfp4": mtq.W4A16_NVFP4_CFG,
+    "nvfp4_mlp_only": mtq.NVFP4_MLP_ONLY_CFG,
+}
+
+
+def _load(export_dir: Path):
+    weights = {}
+    for shard in sorted(export_dir.glob("*.safetensors")):
+        weights.update(load_file(shard))
+    return weights, json.loads((export_dir / "config.json").read_text())
+
+
+def _bits(tensor: torch.Tensor) -> torch.Tensor:
+    return tensor.view(torch.uint8) if tensor.element_size() == 1 else tensor
+
+
+def _excluded(patterns, layers):
+    return {layer for layer in layers if any(fnmatch.fnmatch(layer, p) for p in patterns)}
+
+
+def _test_matches_megatron_exporter(source, quant_cfg_name, tmp_path, rank, size):
+    # Models in a reused worker pick their attention backend through process-wide variables.
+    for name in ("NVTE_FLASH_ATTN", "NVTE_FUSED_ATTN", "NVTE_UNFUSED_ATTN"):
+        os.environ.pop(name, None)
+    bridge, _, model, unwrapped, _ = load_mbridge_model_from_hf(
+        hf_model_name_or_path=str(source),
+        provider_overrides={"pipeline_dtype": torch.bfloat16},
+        moe_grouped_gemm=use_moe_grouped_gemm(str(source)),
+    )
+    torch.manual_seed(0)
+    input_ids = torch.randint(0, unwrapped.vocab_size, (2, 16), device="cuda")
+    mtq.quantize(unwrapped, QUANT_CFGS[quant_cfg_name], lambda m: megatron_prefill(m, input_ids))
+    reference_dir, export_dir = tmp_path / "megatron_exporter", tmp_path / "bridge"
+    export_mcore_gpt_to_hf(
+        unwrapped, str(source), dtype=torch.bfloat16, export_dir=str(reference_dir)
+    )
+    bridge.save_hf_pretrained_modelopt(model, str(export_dir))
+    if rank != 0:
+        return
+
+    ref_weights, ref_config = _load(reference_dir)
+    weights, config = _load(export_dir)
+    assert weights.keys() == ref_weights.keys()
+    for name, weight in weights.items():
+        # The Megatron exporter keeps per-tensor scales 1-d; both hold the same scalar.
+        reference = ref_weights[name].reshape(weight.shape)
+        assert weight.dtype == reference.dtype, name
+        assert torch.equal(_bits(weight), _bits(reference)), name
+
+    quant, ref_quant = config.pop("quantization_config"), ref_config.pop("quantization_config")
+    # The Megatron exporter also writes head_dim where the source config leaves it implied.
+    if "head_dim" not in config:
+        implied = config["hidden_size"] // config["num_attention_heads"]
+        assert ref_config.pop("head_dim", implied) == implied
+    assert config == ref_config
+    assert {k: v for k, v in quant.items() if k != "ignore"} == {
+        k: v for k, v in ref_quant.items() if k != "ignore"
+    }
+    # The exclusion patterns differ in spelling but must exclude the same linear layers.
+    layers = {
+        name.rpartition(".")[0]
+        for name, weight in weights.items()
+        if name.endswith(".weight") and weight.dim() == 2
+    }
+    assert _excluded(quant["ignore"], layers) == _excluded(ref_quant["ignore"], layers)
+
+
+@pytest.mark.parametrize(
+    ("create_source", "quant_cfg_name"),
+    [
+        (create_tiny_llama_dir, "nvfp4"),
+        (create_tiny_qwen3_dir, "fp8"),
+        (create_tiny_qwen3_dir, "w4a16_nvfp4"),
+        (create_tiny_qwen3_moe_dir, "nvfp4_mlp_only"),
+    ],
+)
+def test_matches_megatron_exporter(dist_workers_size_1, tmp_path, create_source, quant_cfg_name):
+    source = create_source(tmp_path, with_tokenizer=True)
+    dist_workers_size_1.run(
+        partial(_test_matches_megatron_exporter, source, quant_cfg_name, tmp_path)
+    )

From 85f305f70ada2db51623ac1e41fff9459c711856 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Tue, 27 Oct 2026 13:13:21 -0700
Subject: [PATCH 2/7] Check exported tensor shapes exactly, except per-tensor
 scales

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .../torch/export/test_quantized_mbridge_export.py            | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tests/gpu_megatron/torch/export/test_quantized_mbridge_export.py b/tests/gpu_megatron/torch/export/test_quantized_mbridge_export.py
index 1bcb7fc9f5a..ccc7fb62ff3 100644
--- a/tests/gpu_megatron/torch/export/test_quantized_mbridge_export.py
+++ b/tests/gpu_megatron/torch/export/test_quantized_mbridge_export.py
@@ -81,8 +81,11 @@ def _test_matches_megatron_exporter(source, quant_cfg_name, tmp_path, rank, size
     weights, config = _load(export_dir)
     assert weights.keys() == ref_weights.keys()
     for name, weight in weights.items():
+        reference = ref_weights[name]
         # The Megatron exporter keeps per-tensor scales 1-d; both hold the same scalar.
-        reference = ref_weights[name].reshape(weight.shape)
+        if reference.numel() == weight.numel() == 1:
+            reference = reference.reshape(weight.shape)
+        assert weight.shape == reference.shape, name
         assert weight.dtype == reference.dtype, name
         assert torch.equal(_bits(weight), _bits(reference)), name
 

From 311800ef64d8eca541d74f8a0efb142145f5e15c Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Tue, 27 Oct 2026 23:07:16 -0700
Subject: [PATCH 3/7] Export through the bridge at TP over all GPUs in the
 parity test

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .../export/test_quantized_mbridge_export.py   | 62 ++++++++++++-------
 1 file changed, 38 insertions(+), 24 deletions(-)

diff --git a/tests/gpu_megatron/torch/export/test_quantized_mbridge_export.py b/tests/gpu_megatron/torch/export/test_quantized_mbridge_export.py
index ccc7fb62ff3..6b73e0f9d86 100644
--- a/tests/gpu_megatron/torch/export/test_quantized_mbridge_export.py
+++ b/tests/gpu_megatron/torch/export/test_quantized_mbridge_export.py
@@ -57,25 +57,51 @@ def _excluded(patterns, layers):
     return {layer for layer in layers if any(fnmatch.fnmatch(layer, p) for p in patterns)}
 
 
-def _test_matches_megatron_exporter(source, quant_cfg_name, tmp_path, rank, size):
+def _quantize(source, quant_cfg_name, size):
     # Models in a reused worker pick their attention backend through process-wide variables.
     for name in ("NVTE_FLASH_ATTN", "NVTE_FUSED_ATTN", "NVTE_UNFUSED_ATTN"):
         os.environ.pop(name, None)
     bridge, _, model, unwrapped, _ = load_mbridge_model_from_hf(
         hf_model_name_or_path=str(source),
-        provider_overrides={"pipeline_dtype": torch.bfloat16},
+        provider_overrides={"tensor_model_parallel_size": size, "pipeline_dtype": torch.bfloat16},
         moe_grouped_gemm=use_moe_grouped_gemm(str(source)),
     )
     torch.manual_seed(0)
     input_ids = torch.randint(0, unwrapped.vocab_size, (2, 16), device="cuda")
     mtq.quantize(unwrapped, QUANT_CFGS[quant_cfg_name], lambda m: megatron_prefill(m, input_ids))
+    return bridge, model, unwrapped
+
+
+def _export_with_megatron_exporter(source, quant_cfg_name, export_dir, rank, size):
+    _, _, unwrapped = _quantize(source, quant_cfg_name, size)
+    export_mcore_gpt_to_hf(unwrapped, str(source), dtype=torch.bfloat16, export_dir=str(export_dir))
+
+
+def _export_with_bridge(source, quant_cfg_name, export_dir, rank, size):
+    bridge, model, _ = _quantize(source, quant_cfg_name, size)
+    bridge.save_hf_pretrained_modelopt(model, str(export_dir))
+
+
+@pytest.mark.parametrize(
+    ("create_source", "quant_cfg_name"),
+    [
+        (create_tiny_llama_dir, "nvfp4"),
+        (create_tiny_qwen3_dir, "fp8"),
+        (create_tiny_qwen3_dir, "w4a16_nvfp4"),
+        (create_tiny_qwen3_moe_dir, "nvfp4_mlp_only"),
+    ],
+)
+def test_matches_megatron_exporter(
+    dist_workers_size_1, dist_workers, tmp_path, create_source, quant_cfg_name
+):
+    source = create_source(tmp_path, with_tokenizer=True)
     reference_dir, export_dir = tmp_path / "megatron_exporter", tmp_path / "bridge"
-    export_mcore_gpt_to_hf(
-        unwrapped, str(source), dtype=torch.bfloat16, export_dir=str(reference_dir)
+    # The Megatron exporter needs TP=1. The bridge exports at TP over all GPUs, so multi-GPU runs
+    # also cover gathering its sharded weights and scales.
+    dist_workers_size_1.run(
+        partial(_export_with_megatron_exporter, source, quant_cfg_name, reference_dir)
     )
-    bridge.save_hf_pretrained_modelopt(model, str(export_dir))
-    if rank != 0:
-        return
+    dist_workers.run(partial(_export_with_bridge, source, quant_cfg_name, export_dir))
 
     ref_weights, ref_config = _load(reference_dir)
     weights, config = _load(export_dir)
@@ -87,7 +113,11 @@ def _test_matches_megatron_exporter(source, quant_cfg_name, tmp_path, rank, size
             reference = reference.reshape(weight.shape)
         assert weight.shape == reference.shape, name
         assert weight.dtype == reference.dtype, name
-        assert torch.equal(_bits(weight), _bits(reference)), name
+        if name.endswith("input_scale") and torch.cuda.device_count() > 1:
+            # Activations reduced across TP ranks calibrate to slightly different amaxes.
+            torch.testing.assert_close(weight, reference, rtol=2e-2, atol=0, msg=name)
+        else:
+            assert torch.equal(_bits(weight), _bits(reference)), name
 
     quant, ref_quant = config.pop("quantization_config"), ref_config.pop("quantization_config")
     # The Megatron exporter also writes head_dim where the source config leaves it implied.
@@ -105,19 +135,3 @@ def _test_matches_megatron_exporter(source, quant_cfg_name, tmp_path, rank, size
         if name.endswith(".weight") and weight.dim() == 2
     }
     assert _excluded(quant["ignore"], layers) == _excluded(ref_quant["ignore"], layers)
-
-
-@pytest.mark.parametrize(
-    ("create_source", "quant_cfg_name"),
-    [
-        (create_tiny_llama_dir, "nvfp4"),
-        (create_tiny_qwen3_dir, "fp8"),
-        (create_tiny_qwen3_dir, "w4a16_nvfp4"),
-        (create_tiny_qwen3_moe_dir, "nvfp4_mlp_only"),
-    ],
-)
-def test_matches_megatron_exporter(dist_workers_size_1, tmp_path, create_source, quant_cfg_name):
-    source = create_source(tmp_path, with_tokenizer=True)
-    dist_workers_size_1.run(
-        partial(_test_matches_megatron_exporter, source, quant_cfg_name, tmp_path)
-    )

From 422b4edd6bdb90b6354b6535db03b5de7f226436 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Wed, 28 Oct 2026 01:51:16 -0700
Subject: [PATCH 4/7] Cover Nemotron-H and Qwen3.5-VL-MoE in the bridge export
 parity test

Swap the Qwen3 W4A16-NVFP4 and Qwen3-MoE cases for Nemotron-H and
Qwen3.5-VL-MoE (NVFP4 MLP-only, vision tower unquantized), reuse
load_safetensors_dir, and check that the bridge ignores exactly the
unquantized linear layers.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .../export/test_quantized_mbridge_export.py   | 50 +++++++++++--------
 1 file changed, 29 insertions(+), 21 deletions(-)

diff --git a/tests/gpu_megatron/torch/export/test_quantized_mbridge_export.py b/tests/gpu_megatron/torch/export/test_quantized_mbridge_export.py
index 6b73e0f9d86..5d704dc437b 100644
--- a/tests/gpu_megatron/torch/export/test_quantized_mbridge_export.py
+++ b/tests/gpu_megatron/torch/export/test_quantized_mbridge_export.py
@@ -22,31 +22,32 @@
 
 import pytest
 import torch
+from _test_utils.torch.export.unified_checkpoint import load_safetensors_dir
 from _test_utils.torch.transformers_models import (
     create_tiny_llama_dir,
+    create_tiny_nemotron_h_dir,
+    create_tiny_qwen3_5_moe_vl_dir,
     create_tiny_qwen3_dir,
-    create_tiny_qwen3_moe_dir,
 )
-from safetensors.torch import load_file
 
 import modelopt.torch.quantization as mtq
 from modelopt.torch.export import export_mcore_gpt_to_hf
-from modelopt.torch.utils.plugins.mbridge import load_mbridge_model_from_hf, use_moe_grouped_gemm
+from modelopt.torch.utils.plugins.mbridge import (
+    get_language_model,
+    load_mbridge_model_from_hf,
+    use_moe_grouped_gemm,
+)
 from modelopt.torch.utils.plugins.megatron_generate import megatron_prefill
 
 QUANT_CFGS = {
     "nvfp4": mtq.NVFP4_DEFAULT_CFG,
     "fp8": mtq.FP8_DEFAULT_CFG,
-    "w4a16_nvfp4": mtq.W4A16_NVFP4_CFG,
     "nvfp4_mlp_only": mtq.NVFP4_MLP_ONLY_CFG,
 }
 
 
 def _load(export_dir: Path):
-    weights = {}
-    for shard in sorted(export_dir.glob("*.safetensors")):
-        weights.update(load_file(shard))
-    return weights, json.loads((export_dir / "config.json").read_text())
+    return load_safetensors_dir(export_dir), json.loads((export_dir / "config.json").read_text())
 
 
 def _bits(tensor: torch.Tensor) -> torch.Tensor:
@@ -66,9 +67,15 @@ def _quantize(source, quant_cfg_name, size):
         provider_overrides={"tensor_model_parallel_size": size, "pipeline_dtype": torch.bfloat16},
         moe_grouped_gemm=use_moe_grouped_gemm(str(source)),
     )
+    language_model, is_vlm = get_language_model(unwrapped)
+    quant_cfg = QUANT_CFGS[quant_cfg_name]
+    if is_vlm:
+        # As in quantize.py, the vision tower stays unquantized.
+        disable_vision = {"quantizer_name": "vision_model.*", "enable": False}
+        quant_cfg = {**quant_cfg, "quant_cfg": [*quant_cfg["quant_cfg"], disable_vision]}
     torch.manual_seed(0)
-    input_ids = torch.randint(0, unwrapped.vocab_size, (2, 16), device="cuda")
-    mtq.quantize(unwrapped, QUANT_CFGS[quant_cfg_name], lambda m: megatron_prefill(m, input_ids))
+    input_ids = torch.randint(0, language_model.vocab_size, (2, 16), device="cuda")
+    mtq.quantize(unwrapped, quant_cfg, lambda m: megatron_prefill(language_model, input_ids))
     return bridge, model, unwrapped
 
 
@@ -85,16 +92,18 @@ def _export_with_bridge(source, quant_cfg_name, export_dir, rank, size):
 @pytest.mark.parametrize(
     ("create_source", "quant_cfg_name"),
     [
-        (create_tiny_llama_dir, "nvfp4"),
-        (create_tiny_qwen3_dir, "fp8"),
-        (create_tiny_qwen3_dir, "w4a16_nvfp4"),
-        (create_tiny_qwen3_moe_dir, "nvfp4_mlp_only"),
+        (partial(create_tiny_llama_dir, with_tokenizer=True), "nvfp4"),
+        (partial(create_tiny_qwen3_dir, with_tokenizer=True), "fp8"),
+        # The bridge can't export quantized Mamba or GatedDeltaNet projections yet.
+        (partial(create_tiny_nemotron_h_dir, with_tokenizer=True), "nvfp4_mlp_only"),
+        (partial(create_tiny_qwen3_5_moe_vl_dir, with_processor=True), "nvfp4_mlp_only"),
     ],
+    ids=["llama-nvfp4", "qwen3-fp8", "nemotron_h-nvfp4_mlp_only", "qwen3_5_moe_vl-nvfp4_mlp_only"],
 )
 def test_matches_megatron_exporter(
     dist_workers_size_1, dist_workers, tmp_path, create_source, quant_cfg_name
 ):
-    source = create_source(tmp_path, with_tokenizer=True)
+    source = create_source(tmp_path)
     reference_dir, export_dir = tmp_path / "megatron_exporter", tmp_path / "bridge"
     # The Megatron exporter needs TP=1. The bridge exports at TP over all GPUs, so multi-GPU runs
     # also cover gathering its sharded weights and scales.
@@ -120,18 +129,17 @@ def test_matches_megatron_exporter(
             assert torch.equal(_bits(weight), _bits(reference)), name
 
     quant, ref_quant = config.pop("quantization_config"), ref_config.pop("quantization_config")
-    # The Megatron exporter also writes head_dim where the source config leaves it implied.
-    if "head_dim" not in config:
-        implied = config["hidden_size"] // config["num_attention_heads"]
-        assert ref_config.pop("head_dim", implied) == implied
     assert config == ref_config
     assert {k: v for k, v in quant.items() if k != "ignore"} == {
         k: v for k, v in ref_quant.items() if k != "ignore"
     }
-    # The exclusion patterns differ in spelling but must exclude the same linear layers.
+    # The bridge must exclude exactly the unquantized linear layers. The Megatron exporter's
+    # patterns spell them differently and miss some, e.g. Qwen3.5's mlp.shared_expert_gate.
     layers = {
         name.rpartition(".")[0]
         for name, weight in weights.items()
         if name.endswith(".weight") and weight.dim() == 2
     }
-    assert _excluded(quant["ignore"], layers) == _excluded(ref_quant["ignore"], layers)
+    unquantized = {layer for layer in layers if f"{layer}.weight_scale" not in weights}
+    assert _excluded(quant["ignore"], layers) == unquantized
+    assert _excluded(ref_quant["ignore"], layers) <= unquantized

From c5c8a8627a588bd5ec34308e55be1a9911bf1fe6 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Wed, 28 Oct 2026 02:26:48 -0700
Subject: [PATCH 5/7] Cap the bridge export TP at 2 in the parity test

The tiny models have as few as 2 KV heads, so TP over all GPUs fails on
hosts with more than 2 GPUs. Pick the worker pool as other gpu_megatron
tests do.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .../torch/export/test_quantized_mbridge_export.py     | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/tests/gpu_megatron/torch/export/test_quantized_mbridge_export.py b/tests/gpu_megatron/torch/export/test_quantized_mbridge_export.py
index 5d704dc437b..d558d932dab 100644
--- a/tests/gpu_megatron/torch/export/test_quantized_mbridge_export.py
+++ b/tests/gpu_megatron/torch/export/test_quantized_mbridge_export.py
@@ -101,16 +101,19 @@ def _export_with_bridge(source, quant_cfg_name, export_dir, rank, size):
     ids=["llama-nvfp4", "qwen3-fp8", "nemotron_h-nvfp4_mlp_only", "qwen3_5_moe_vl-nvfp4_mlp_only"],
 )
 def test_matches_megatron_exporter(
-    dist_workers_size_1, dist_workers, tmp_path, create_source, quant_cfg_name
+    request, dist_workers_size_1, tmp_path, create_source, quant_cfg_name
 ):
     source = create_source(tmp_path)
     reference_dir, export_dir = tmp_path / "megatron_exporter", tmp_path / "bridge"
-    # The Megatron exporter needs TP=1. The bridge exports at TP over all GPUs, so multi-GPU runs
-    # also cover gathering its sharded weights and scales.
+    # The Megatron exporter needs TP=1. The bridge exports at TP=2 when two GPUs are available, so
+    # it also gathers sharded weights and scales. TP stays <= 2, the fewest KV heads of these models.
+    bridge_workers = request.getfixturevalue(
+        f"dist_workers_size_{min(torch.cuda.device_count(), 2)}"
+    )
     dist_workers_size_1.run(
         partial(_export_with_megatron_exporter, source, quant_cfg_name, reference_dir)
     )
-    dist_workers.run(partial(_export_with_bridge, source, quant_cfg_name, export_dir))
+    bridge_workers.run(partial(_export_with_bridge, source, quant_cfg_name, export_dir))
 
     ref_weights, ref_config = _load(reference_dir)
     weights, config = _load(export_dir)

From d23a1285d2f333a38bf979e3464b81dea9797ab7 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Wed, 28 Oct 2026 03:10:31 -0700
Subject: [PATCH 6/7] Start the bridge worker pool in a fixture in the parity
 test

pytest-timeout only times the test call, so pool startup belongs in fixture
setup rather than in request.getfixturevalue inside the test.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .../torch/export/test_quantized_mbridge_export.py | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/tests/gpu_megatron/torch/export/test_quantized_mbridge_export.py b/tests/gpu_megatron/torch/export/test_quantized_mbridge_export.py
index d558d932dab..46343b9fe96 100644
--- a/tests/gpu_megatron/torch/export/test_quantized_mbridge_export.py
+++ b/tests/gpu_megatron/torch/export/test_quantized_mbridge_export.py
@@ -89,6 +89,13 @@ def _export_with_bridge(source, quant_cfg_name, export_dir, rank, size):
     bridge.save_hf_pretrained_modelopt(model, str(export_dir))
 
 
+@pytest.fixture
+def bridge_workers(request):
+    # TP=2 on multi-GPU hosts, so the bridge also gathers sharded weights and scales. TP stays <= 2,
+    # the fewest KV heads of these models. A fixture starts the pool before the test timeout.
+    return request.getfixturevalue(f"dist_workers_size_{min(torch.cuda.device_count(), 2)}")
+
+
 @pytest.mark.parametrize(
     ("create_source", "quant_cfg_name"),
     [
@@ -101,15 +108,11 @@ def _export_with_bridge(source, quant_cfg_name, export_dir, rank, size):
     ids=["llama-nvfp4", "qwen3-fp8", "nemotron_h-nvfp4_mlp_only", "qwen3_5_moe_vl-nvfp4_mlp_only"],
 )
 def test_matches_megatron_exporter(
-    request, dist_workers_size_1, tmp_path, create_source, quant_cfg_name
+    dist_workers_size_1, bridge_workers, tmp_path, create_source, quant_cfg_name
 ):
     source = create_source(tmp_path)
     reference_dir, export_dir = tmp_path / "megatron_exporter", tmp_path / "bridge"
-    # The Megatron exporter needs TP=1. The bridge exports at TP=2 when two GPUs are available, so
-    # it also gathers sharded weights and scales. TP stays <= 2, the fewest KV heads of these models.
-    bridge_workers = request.getfixturevalue(
-        f"dist_workers_size_{min(torch.cuda.device_count(), 2)}"
-    )
+    # The Megatron exporter needs TP=1.
     dist_workers_size_1.run(
         partial(_export_with_megatron_exporter, source, quant_cfg_name, reference_dir)
     )

From 94e6d0c310037c6d58d64b31f9d7cb4606ad25a4 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Wed, 28 Oct 2026 12:08:05 -0700
Subject: [PATCH 7/7] Gate the input_scale tolerance on the bridge's TP size

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .../gpu_megatron/torch/export/test_quantized_mbridge_export.py  | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/gpu_megatron/torch/export/test_quantized_mbridge_export.py b/tests/gpu_megatron/torch/export/test_quantized_mbridge_export.py
index 46343b9fe96..9cfa1549835 100644
--- a/tests/gpu_megatron/torch/export/test_quantized_mbridge_export.py
+++ b/tests/gpu_megatron/torch/export/test_quantized_mbridge_export.py
@@ -128,7 +128,7 @@ def test_matches_megatron_exporter(
             reference = reference.reshape(weight.shape)
         assert weight.shape == reference.shape, name
         assert weight.dtype == reference.dtype, name
-        if name.endswith("input_scale") and torch.cuda.device_count() > 1:
+        if name.endswith("input_scale") and bridge_workers.world_size > 1:
             # Activations reduced across TP ranks calibrate to slightly different amaxes.
             torch.testing.assert_close(weight, reference, rtol=2e-2, atol=0, msg=name)
         else: