perf(recipe): add an 8-GPU GB300 MXFP8 pretrain recipe for Qwen3.5-VL 35B-A3B (#6413)

Signed-off-by: tomlifu <lifuz@nvidia.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

Author: tomlifu

scripts/training/run_recipe.py

@@ -132,7 +132,11 @@
 TrainMode = Literal["pretrain", "finetune"]
 
 
-IMPORT_TIME_RECIPE_ENV_VARS = frozenset({"NVTE_CPU_OFFLOAD_V1"})
+# Recipe environment variables that are read when a library is imported, before this script applies the recipe
+# environment. A library recipe that declares one of them (and the user has not set it) is re-executed with the
+# recipe environment in place. Transformer Engine reads NVTE_CUTEDSL_FUSED_GROUPED_MLP once, when it registers the
+# CuTe DSL fused grouped MLP at import time; set any later, it is ignored and the unfused grouped GEMM runs.
+IMPORT_TIME_RECIPE_ENV_VARS = frozenset({"NVTE_CPU_OFFLOAD_V1", "NVTE_CUTEDSL_FUSED_GROUPED_MLP"})
 
 
 COMMON_OVERRIDE_FIELDS = (

src/megatron/bridge/recipes/qwen_vl/__init__.py

@@ -23,6 +23,7 @@
 
 # Qwen3.5/Qwen3.6-VL GB300 models
 from .gb300.qwen35_vl import (
+    qwen35_vl_35b_a3b_pretrain_8gpu_gb300_fp8mx_library_config,
     qwen35_vl_35b_a3b_pretrain_16gpu_gb300_bf16_config,
     qwen35_vl_397b_a17b_pretrain_config,
 )
@@ -85,6 +86,7 @@
     "qwen35_vl_9b_pretrain_mock_config",
     "qwen35_vl_27b_pretrain_16gpu_gb200_bf16_mock_config",
     "qwen35_vl_35b_a3b_pretrain_config",
+    "qwen35_vl_35b_a3b_pretrain_8gpu_gb300_fp8mx_library_config",
     "qwen35_vl_35b_a3b_pretrain_16gpu_gb300_bf16_config",
     "qwen35_vl_397b_a17b_pretrain_config",
     "qwen35_vl_35b_a3b_pretrain_mock_config",

src/megatron/bridge/recipes/qwen_vl/gb300/__init__.py

@@ -13,12 +13,14 @@
 # limitations under the License.
 
 from megatron.bridge.recipes.qwen_vl.gb300.qwen35_vl import (
+    qwen35_vl_35b_a3b_pretrain_8gpu_gb300_fp8mx_library_config,
     qwen35_vl_35b_a3b_pretrain_16gpu_gb300_bf16_config,
     qwen35_vl_397b_a17b_pretrain_config,
 )
 
 
 __all__ = [
+    "qwen35_vl_35b_a3b_pretrain_8gpu_gb300_fp8mx_library_config",
     "qwen35_vl_35b_a3b_pretrain_16gpu_gb300_bf16_config",
     "qwen35_vl_397b_a17b_pretrain_config",
 ]

src/megatron/bridge/recipes/qwen_vl/gb300/qwen35_vl.py

@@ -25,7 +25,7 @@
 from megatron.bridge.recipes.utils.environment_utils import COMMON_RECIPE_ENV_VARS
 from megatron.bridge.training.comm_overlap import CommOverlapConfig
 from megatron.bridge.training.config import ConfigContainer
-from megatron.bridge.training.mixed_precision import bf16_mixed
+from megatron.bridge.training.mixed_precision import bf16_mixed, bf16_with_mxfp8_mixed
 from megatron.bridge.utils.cuda_graph import set_cuda_graph_modules
 
 
@@ -330,7 +330,94 @@ def qwen35_vl_35b_a3b_pretrain_16gpu_gb300_bf16_config() -> ConfigContainer:
     return cfg
 
 
+def qwen35_vl_35b_a3b_pretrain_8gpu_gb300_fp8mx_library_config() -> ConfigContainer:
+    """Return the MXFP8 full-pretraining config for Qwen3.5/Qwen3.6-VL 35B-A3B on 8 GB300 GPUs.
+
+    Starts from :func:`qwen35_vl_35b_a3b_pretrain_16gpu_gb300_bf16_config` and keeps its parallel layout (EP 4,
+    TP/PP/CP 1), CUDA graphs, recompute, batch shape (micro batch 4, global batch 2048) and real routing. On top it
+    sets MXFP8, Transformer Engine's CuTe DSL fused grouped MLP, cuDNN GDN, and the HybridEP, dispatcher and
+    normalization settings of the configuration it was tuned on. With TP, PP and CP at 1, data parallelism is 8 on
+    8 GPUs (expert data parallelism 2), so each rank runs 64 micro batches per step.
+
+    Like the 16-GPU recipe, it defaults to the mock VLM dataset. The measured runs used an Energon dataset
+    (``--dataset energon dataset.path=...``) with ``dataset.task_encoder.max_pixels=200704``, and the vision CUDA graph
+    cap inherited from the 16-GPU recipe (3712 patches per micro batch) was measured with that image budget at micro
+    batch 4. Keep both, or re-measure the cap: a larger image budget can exceed it and stop training with an error.
+
+    cuDNN GDN is selected only on a Megatron-Core that declares ``gdn_kernel_backend`` (NVIDIA/Megatron-LM#7583);
+    otherwise the GatedDeltaNet layers keep running on FLA.
+
+    Transformer Engine reads the fused-MLP gate ``NVTE_CUTEDSL_FUSED_GROUPED_MLP`` once, when it is imported.
+    ``scripts/training/run_recipe.py``, which ``train.sh`` runs, sets it in time by re-executing itself; with any
+    other entry point, export it before Python starts, or Transformer Engine silently runs the unfused grouped GEMM.
+
+    The name carries a ``_library`` suffix because ``qwen35_vl_35b_a3b_pretrain_8gpu_gb300_fp8mx_config`` is
+    already a perf recipe (a mock-data benchmark with EP 8 and global batch 512), and bare recipe lookup keeps
+    resolving that name to it.
+
+    Measured on GB300 with real DataComp data and real routing (steps 101-150, one run per arm, submitted together),
+    a launch configuration close to this recipe reached 11.3% more TFLOP/s per GPU on 8 GPUs than on 16 at the same
+    global batch size. The two runs differ in data-parallel size and in micro batches per rank (64 vs 32), and one
+    run per arm does not separate these effects, so no cause is claimed.
+    """
+    cfg = qwen35_vl_35b_a3b_pretrain_16gpu_gb300_bf16_config()
+
+    cfg.mixed_precision = bf16_with_mxfp8_mixed()
+    cfg.mixed_precision.grad_reduce_in_fp32 = False
+
+    # CuTe DSL fused grouped MLP (FC1 + SwiGLU + FC2 for MXFP8 on GPUs with compute capability 10.x). The op fuser,
+    # the 32-wide GLU interleave and the NVTE_CUTEDSL_FUSED_GROUPED_MLP env gate (cfg.env_vars below) are all
+    # required. Transformer Engine reads the gate when it is imported and registers the fusion only if it is set
+    # (scripts/training/run_recipe.py re-executes itself so the gate is set before that import; other entry points
+    # must export it); without the gate or the interleave it silently runs the unfused grouped GEMM.
+    cfg.model.use_transformer_engine_op_fuser = True
+    cfg.model.moe_mlp_glu_interleave_size = 32
+
+    cfg.model.moe_flex_dispatcher_num_sms = 32
+    # Off: HybridEP's fused permute requires identical dispatch, combine and preprocessing chunk sizes, and this
+    # recipe uses 64/128/64 (cfg.env_vars below). The tuned configuration ran this way too.
+    cfg.model.moe_permute_fusion_into_hybridep = False
+
+    # The tuned configuration started from qwen35_vl_35b_a3b_pretrain_config and kept these values. The 16-GPU
+    # recipe does not carry them over (it leaves the two model fields at Megatron-Core defaults, sets the dispatch
+    # and preprocessing chunk sizes to 128 and omits the cuDNN normalization flags), and was never measured with
+    # MXFP8.
+    cfg.model.moe_hybridep_num_sms_preprocessing = 16
+    cfg.model.overlap_dispatch_backward_with_experts_wgrad = True
+
+    # cuDNN GDN (FROST engine: cuDNN frontend >= 1.29 and nvidia-cutlass-dsl >= 4.7): measured +12.0% to +19.4% end
+    # to end at 35B-A3B on 16 GB300 GPUs (see _select_gdn_kernel_backend in perf_recipes/qwen_vl/common.py). Guarded
+    # like _enable_gdn_conv_fusion: assigning a field Megatron-Core lacks would not raise, and the recipe would look
+    # enabled while FLA keeps running.
+    if hasattr(type(cfg.model), "gdn_kernel_backend") or hasattr(cfg.model, "gdn_kernel_backend"):
+        cfg.model.gdn_kernel_backend = "transformer_engine"
+
+    # The complete environment, declared here rather than inherited from the 16-GPU recipe. It differs from that
+    # recipe in the fused-MLP gate, the dispatch and preprocessing chunk sizes (64 instead of 128) and the cuDNN
+    # normalization flags.
+    cfg.env_vars = {
+        **COMMON_RECIPE_ENV_VARS,
+        "CUDA_DEVICE_MAX_CONNECTIONS": 32,
+        # Must equal expert_model_parallel_size or HybridEP splits the NVLink
+        # domain wrongly. Unset does NOT auto-detect.
+        "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 4,
+        "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128,
+        "NUM_OF_TOKENS_PER_CHUNK_DISPATCH_API": 64,
+        "NUM_OF_TOKENS_PER_CHUNK_PREPROCESSING_API": 64,
+        "NVLINK_DOMAIN_SIZE": 72,
+        "USE_MNNVL": 1,
+        # 0, as in the 16-GPU recipe: the LayerNorm SM margin follows the communication share (see there).
+        "NVTE_BWD_LAYERNORM_SM_MARGIN": 0,
+        "NVTE_FWD_LAYERNORM_SM_MARGIN": 0,
+        "NVTE_NORM_BWD_USE_CUDNN": 1,
+        "NVTE_NORM_FWD_USE_CUDNN": 1,
+        "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1,
+    }
+    return cfg
+
+
 __all__ = [
+    "qwen35_vl_35b_a3b_pretrain_8gpu_gb300_fp8mx_library_config",
     "qwen35_vl_35b_a3b_pretrain_16gpu_gb300_bf16_config",
     "qwen35_vl_397b_a17b_pretrain_config",
 ]

tests/unit_tests/recipes/qwen_vl/test_qwen35_vl_recipes.py

@@ -115,6 +115,7 @@
 ]
 
 _QWEN35_VL_GB300_FUNCS = [
+    _qwen35_vl_gb300_module.qwen35_vl_35b_a3b_pretrain_8gpu_gb300_fp8mx_library_config,
     _qwen35_vl_gb300_module.qwen35_vl_35b_a3b_pretrain_16gpu_gb300_bf16_config,
     _qwen35_vl_gb300_module.qwen35_vl_397b_a17b_pretrain_config,
 ]
@@ -514,6 +515,120 @@ def test_qwen35_vl_35b_a3b_pretrain_16gpu_gb300_defaults(monkeypatch: pytest.Mon
     assert cfg.checkpoint.pretrained_checkpoint is None
 
 
+def test_qwen35_vl_35b_a3b_pretrain_8gpu_gb300_fp8mx_library_defaults(monkeypatch: pytest.MonkeyPatch):
+    """The 8-GB300 MXFP8 recipe adds the measured precision, kernel and dispatcher settings to the 16-GPU recipe."""
+    patch_recipe_module_global(monkeypatch, _qwen35_vl_h100_module, "AutoBridge", _FakeAutoBridge)
+
+    cfg = _qwen35_vl_gb300_module.qwen35_vl_35b_a3b_pretrain_8gpu_gb300_fp8mx_library_config()
+    base = _qwen35_vl_gb300_module.qwen35_vl_35b_a3b_pretrain_16gpu_gb300_bf16_config()
+
+    _assert_basic_config(cfg)
+    # Same parallel layout, CUDA graphs, recompute, batch shape and routing as the 16-GPU recipe.
+    assert cfg.model.expert_model_parallel_size == 4
+    assert cfg.env_vars["NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN"] == 4
+    assert cfg.train.global_batch_size == base.train.global_batch_size == 2048
+    assert cfg.train.micro_batch_size == base.train.micro_batch_size == 4
+    assert cfg.model.max_vision_cuda_graph_seq_length == 3712
+    assert cuda_graph_module_names(cfg.model) == ["attn", "moe_router", "moe_preprocess"]
+    assert cfg.model.recompute_modules == ["core_attn", "gdn_norm_out"]
+    assert cfg.model.moe_router_force_load_balancing is False
+    # MXFP8 with FP8 parameter all-gather.
+    assert cfg.mixed_precision.fp8 == "e4m3"
+    assert cfg.mixed_precision.fp8_recipe == "mxfp8"
+    assert cfg.mixed_precision.fp8_param_gather is True
+    assert cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag is True
+    assert cfg.mixed_precision.grad_reduce_in_fp32 is False
+    # CuTe DSL fused grouped MLP needs all three settings together.
+    assert cfg.model.use_transformer_engine_op_fuser is True
+    assert cfg.model.moe_mlp_glu_interleave_size == 32
+    assert cfg.env_vars["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] == 1
+    assert cfg.model.moe_flex_dispatcher_num_sms == 32
+    assert cfg.model.moe_permute_fusion_into_hybridep is False
+    assert cfg.model.moe_hybridep_num_sms_preprocessing == 16
+    assert cfg.model.overlap_dispatch_backward_with_experts_wgrad is True
+    assert cfg.env_vars["NUM_OF_TOKENS_PER_CHUNK_COMBINE_API"] == 128
+    assert cfg.env_vars["NUM_OF_TOKENS_PER_CHUNK_DISPATCH_API"] == 64
+    assert cfg.env_vars["NUM_OF_TOKENS_PER_CHUNK_PREPROCESSING_API"] == 64
+    assert cfg.env_vars["NVTE_NORM_FWD_USE_CUDNN"] == 1
+    assert cfg.env_vars["NVTE_NORM_BWD_USE_CUDNN"] == 1
+    # Building the 8-GPU recipe must not leak the env gate into the 16-GPU recipe.
+    assert "NVTE_CUTEDSL_FUSED_GROUPED_MLP" not in base.env_vars
+
+    # These are the only model-field and environment-variable differences from the 16-GPU recipe; an extra model
+    # field or environment variable fails here. Mixed precision is checked above on five fields only, and other
+    # sections, such as train, optimizer, ddp and comm_overlap, are not compared.
+    changed_model_fields = {
+        name
+        for name in vars(cfg.model).keys() | vars(base.model).keys()
+        if getattr(cfg.model, name, None) != getattr(base.model, name, None)
+    }
+    assert changed_model_fields == {
+        "use_transformer_engine_op_fuser",
+        "moe_mlp_glu_interleave_size",
+        "moe_flex_dispatcher_num_sms",
+        "moe_permute_fusion_into_hybridep",
+        "moe_hybridep_num_sms_preprocessing",
+        "overlap_dispatch_backward_with_experts_wgrad",
+    }
+    changed_env_vars = {
+        name
+        for name in cfg.env_vars.keys() | base.env_vars.keys()
+        if cfg.env_vars.get(name) != base.env_vars.get(name)
+    }
+    assert changed_env_vars == {
+        "NVTE_CUTEDSL_FUSED_GROUPED_MLP",
+        "NUM_OF_TOKENS_PER_CHUNK_DISPATCH_API",
+        "NUM_OF_TOKENS_PER_CHUNK_PREPROCESSING_API",
+        "NVTE_NORM_FWD_USE_CUDNN",
+        "NVTE_NORM_BWD_USE_CUDNN",
+    }
+
+
+class _FakeModelCfgWithGdnKernelBackend(_FakeModelCfg):
+    """Model config from a Megatron-Core that declares gdn_kernel_backend (NVIDIA/Megatron-LM#7583)."""
+
+    # Megatron-Core's default; a class attribute, as on the real dataclass.
+    gdn_kernel_backend = "fla"
+
+
+class _FakeAutoBridgeWithGdnKernelBackend(_FakeAutoBridge):
+    """AutoBridge yielding a provider that declares gdn_kernel_backend."""
+
+    @staticmethod
+    def from_hf_pretrained(hf_path: str, *, revision: str | None = None):
+        return _FakeAutoBridgeWithGdnKernelBackend()
+
+    def to_megatron_provider(self, load_weights: bool = False):
+        return _FakeModelCfgWithGdnKernelBackend()
+
+
+def test_qwen35_vl_35b_a3b_pretrain_8gpu_gb300_fp8mx_library_selects_cudnn_gdn(monkeypatch: pytest.MonkeyPatch):
+    """If Megatron-Core has gdn_kernel_backend, the 8-GPU recipe selects cuDNN GDN; the 16-GPU one keeps FLA."""
+    patch_recipe_module_global(monkeypatch, _qwen35_vl_h100_module, "AutoBridge", _FakeAutoBridgeWithGdnKernelBackend)
+
+    cfg = _qwen35_vl_gb300_module.qwen35_vl_35b_a3b_pretrain_8gpu_gb300_fp8mx_library_config()
+    base = _qwen35_vl_gb300_module.qwen35_vl_35b_a3b_pretrain_16gpu_gb300_bf16_config()
+
+    assert cfg.model.gdn_kernel_backend == "transformer_engine"
+    # The 16-GPU BF16 recipe keeps Megatron-Core's default.
+    assert base.model.gdn_kernel_backend == "fla"
+
+
+def test_qwen35_vl_35b_a3b_pretrain_8gpu_gb300_fp8mx_library_no_phantom_gdn_kernel_backend(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """On a Megatron-Core without gdn_kernel_backend the recipe must not invent the attribute.
+
+    Assigning an unknown field would not raise: it would create an unused attribute, so the recipe would look as if it
+    selected cuDNN GDN while the layers keep running on FLA.
+    """
+    patch_recipe_module_global(monkeypatch, _qwen35_vl_h100_module, "AutoBridge", _FakeAutoBridge)
+
+    cfg = _qwen35_vl_gb300_module.qwen35_vl_35b_a3b_pretrain_8gpu_gb300_fp8mx_library_config()
+
+    assert not hasattr(cfg.model, "gdn_kernel_backend")
+
+
 def test_qwen35_vl_35b_a3b_pretrain_16gpu_h100_defaults(monkeypatch: pytest.MonkeyPatch):
     """The 16-H100 library pretrain recipe should own the measured execution policy."""
     patch_recipe_module_global(monkeypatch, _qwen35_vl_h100_module, "AutoBridge", _FakeAutoBridge)

tests/unit_tests/recipes/test_recipe_environment_defaults.py

@@ -210,6 +210,16 @@ def test_explicit_recipe_environment_invariants():
                 "NVTE_FWD_LAYERNORM_SM_MARGIN": 0,
             },
         ),
+        (
+            "qwen_vl/gb300/qwen35_vl.py",
+            "qwen35_vl_35b_a3b_pretrain_8gpu_gb300_fp8mx_library_config",
+            {
+                "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 4,
+                "NUM_OF_TOKENS_PER_CHUNK_DISPATCH_API": 64,
+                "NUM_OF_TOKENS_PER_CHUNK_PREPROCESSING_API": 64,
+                "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1,
+            },
+        ),
         (
             "nemotronh_multimodal/h100/nemotron_35_super_vl.py",
             "nemotron_35_super_vl_sft_64gpu_h100_bf16_config",

tests/unit_tests/scripts/training/test_run_recipe.py

@@ -496,10 +496,11 @@ def test_library_only_canonical_name_applies_runtime_environment_without_restart
     handles.recipe_runner.load_forward_step.assert_called_once_with("llm_step", mode="pretrain")
 
 
-def test_library_recipe_bootstraps_unset_import_time_environment(monkeypatch):
+@pytest.mark.parametrize("env_name", ["NVTE_CPU_OFFLOAD_V1", "NVTE_CUTEDSL_FUSED_GROUPED_MLP"])
+def test_library_recipe_bootstraps_unset_import_time_environment(monkeypatch, env_name):
     module, handles = _load_module()
-    monkeypatch.delenv("NVTE_CPU_OFFLOAD_V1", raising=False)
-    config = SimpleNamespace(env_vars={"NVTE_CPU_OFFLOAD_V1": 1})
+    monkeypatch.delenv(env_name, raising=False)
+    config = SimpleNamespace(env_vars={env_name: 1})
     handles.recipe_runner.load_recipe.return_value = config
 
     module.main(["--recipe", "deepseek_v4_flash_pretrain_config", "--mode", "pretrain"])
@@ -511,10 +512,11 @@ def test_library_recipe_bootstraps_unset_import_time_environment(monkeypatch):
     handles.recipe_runner.apply_runtime_environment.assert_not_called()
 
 
-def test_library_recipe_preserves_explicit_import_time_environment(monkeypatch):
+@pytest.mark.parametrize("env_name", ["NVTE_CPU_OFFLOAD_V1", "NVTE_CUTEDSL_FUSED_GROUPED_MLP"])
+def test_library_recipe_preserves_explicit_import_time_environment(monkeypatch, env_name):
     module, handles = _load_module()
-    monkeypatch.setenv("NVTE_CPU_OFFLOAD_V1", "0")
-    config = SimpleNamespace(env_vars={"NVTE_CPU_OFFLOAD_V1": 1})
+    monkeypatch.setenv(env_name, "0")
+    config = SimpleNamespace(env_vars={env_name: 1})
     handles.recipe_runner.load_recipe.return_value = config
 
     module.main(["--recipe", "deepseek_v4_flash_pretrain_config", "--mode", "pretrain"])