From d9d98b9e3563aff5ebdec7fd2ae1e96269aa5d12 Mon Sep 17 00:00:00 2001 From: tomlifu Date: Sun, 11 Oct 2026 07:47:50 -0700 Subject: [PATCH 1/2] perf(recipe): add an 8-GPU GB300 MXFP8 pretrain recipe for Qwen3.5-VL 35B-A3B Add qwen35_vl_35b_a3b_pretrain_8gpu_gb300_fp8mx_library_config, an MXFP8 Qwen3.5-VL 35B-A3B pretraining recipe for 8 GB300 GPUs. It starts from qwen35_vl_35b_a3b_pretrain_16gpu_gb300_bf16_config, keeps its parallel layout, CUDA graphs, recompute, batch shape and real routing, and adds: - MXFP8 with FP8 parameter all-gather (bf16_with_mxfp8_mixed); - Transformer Engine's CuTe DSL fused grouped MLP: the op fuser, the 32-wide GLU interleave and the NVTE_CUTEDSL_FUSED_GROUPED_MLP gate, which only work together; - 32 flex-dispatcher SMs, with HybridEP's fused permute off: it requires identical dispatch, combine and preprocessing chunk sizes, and this recipe uses 64/128/64; - the HybridEP preprocessing SMs, dispatch/expert-wgrad overlap, dispatch and preprocessing chunk sizes and cuDNN normalization of the configuration it was tuned on, which the 16-GPU recipe does not carry over; - cuDNN GDN (gdn_kernel_backend="transformer_engine"), set only when Megatron-Core declares the field (NVIDIA/Megatron-LM#7583), with the same guard as _enable_gdn_conv_fusion. The name carries a _library suffix because qwen35_vl_35b_a3b_pretrain_8gpu_gb300_fp8mx_config is already a perf recipe. Transformer Engine reads NVTE_CUTEDSL_FUSED_GROUPED_MLP once, when it registers the fused MLP at import time, and run_recipe.py applied a library recipe's env_vars after that import, so the gate was ignored. Add the variable to IMPORT_TIME_RECIPE_ENV_VARS, so that a recipe declaring it goes through the existing re-exec bootstrap, as for NVTE_CPU_OFFLOAD_V1. Through run_recipe.py this also changes two existing GB200 library recipes that declare the gate and no other import-time variable: qwen35_text_35b_a3b_sft_long_context_16gpu_gb200_fp8mx_config, which enables the op fuser and the GLU interleave, now runs the fused MLP instead of the unfused grouped GEMM (not measured here), and qwen3_30b_a3b_pretrain_8gpu_gb200_fp8mx_config, which keeps the op fuser off, now re-executes once. Tests check the recipe's settings, the exact set of model fields and environment variables in which it differs from the 16-GPU recipe, that building it leaves the 16-GPU recipe unchanged, that cuDNN GDN is selected for a model config that declares the field and that no attribute is created on one without it, and that the bootstrap covers both import-time variables. Signed-off-by: tomlifu Co-Authored-By: Claude Opus 5.5 --- scripts/training/run_recipe.py | 6 +- .../bridge/recipes/qwen_vl/__init__.py | 2 + .../bridge/recipes/qwen_vl/gb300/__init__.py | 2 + .../bridge/recipes/qwen_vl/gb300/qwen35_vl.py | 72 ++++++++++- .../recipes/qwen_vl/test_qwen35_vl_recipes.py | 115 ++++++++++++++++++ .../scripts/training/test_run_recipe.py | 14 ++- 6 files changed, 203 insertions(+), 8 deletions(-) diff --git a/scripts/training/run_recipe.py b/scripts/training/run_recipe.py index 02c307dcab..898b7a0b5d 100755 --- a/scripts/training/run_recipe.py +++ b/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 = ( diff --git a/src/megatron/bridge/recipes/qwen_vl/__init__.py b/src/megatron/bridge/recipes/qwen_vl/__init__.py index 2fcf7cf7e1..729f03eaf0 100644 --- a/src/megatron/bridge/recipes/qwen_vl/__init__.py +++ b/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", diff --git a/src/megatron/bridge/recipes/qwen_vl/gb300/__init__.py b/src/megatron/bridge/recipes/qwen_vl/gb300/__init__.py index 16867123f4..0431186495 100644 --- a/src/megatron/bridge/recipes/qwen_vl/gb300/__init__.py +++ b/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", ] diff --git a/src/megatron/bridge/recipes/qwen_vl/gb300/qwen35_vl.py b/src/megatron/bridge/recipes/qwen_vl/gb300/qwen35_vl.py index a2465f9501..653c752fff 100644 --- a/src/megatron/bridge/recipes/qwen_vl/gb300/qwen35_vl.py +++ b/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,77 @@ 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 env gate 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.env_vars["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] = 1 + + 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 (set below and in the 16-GPU env). 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 both 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 + cfg.env_vars["NUM_OF_TOKENS_PER_CHUNK_DISPATCH_API"] = 64 + cfg.env_vars["NUM_OF_TOKENS_PER_CHUNK_PREPROCESSING_API"] = 64 + cfg.env_vars["NVTE_NORM_BWD_USE_CUDNN"] = 1 + cfg.env_vars["NVTE_NORM_FWD_USE_CUDNN"] = 1 + + # 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" + + 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", ] diff --git a/tests/unit_tests/recipes/qwen_vl/test_qwen35_vl_recipes.py b/tests/unit_tests/recipes/qwen_vl/test_qwen35_vl_recipes.py index c768050643..c573583e24 100644 --- a/tests/unit_tests/recipes/qwen_vl/test_qwen35_vl_recipes.py +++ b/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) diff --git a/tests/unit_tests/scripts/training/test_run_recipe.py b/tests/unit_tests/scripts/training/test_run_recipe.py index 8fb1d4fedc..2140564133 100644 --- a/tests/unit_tests/scripts/training/test_run_recipe.py +++ b/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"]) From b4e53792a34927d58c38bbc1e7139ebb19fe6afd Mon Sep 17 00:00:00 2001 From: tomlifu Date: Mon, 12 Oct 2026 10:02:52 -0700 Subject: [PATCH 2/2] refactor(recipe): declare the 8-GPU Qwen3.5-VL recipe environment inline Address review: qwen35_vl_35b_a3b_pretrain_8gpu_gb300_fp8mx_library_config no longer edits the env_vars it inherits from qwen35_vl_35b_a3b_pretrain_16gpu_gb300_bf16_config. It declares its whole environment in one cfg.env_vars literal, COMMON_RECIPE_ENV_VARS first, right before it returns, as hardware recipes do, so the recipe shows every variable it sets. The environment is unchanged: the five variables the recipe used to overwrite or add (NVTE_CUTEDSL_FUSED_GROUPED_MLP, the dispatch and preprocessing chunk sizes and the two cuDNN normalization flags) keep their values, and every other variable keeps the 16-GPU value. test_representative_recipe_environment_is_visible now covers this recipe, so a recipe that edits inherited env_vars again fails there. Signed-off-by: tomlifu Co-Authored-By: Claude Opus 5.5 --- .../bridge/recipes/qwen_vl/gb300/qwen35_vl.py | 41 +++++++++++++------ .../test_recipe_environment_defaults.py | 10 +++++ 2 files changed, 39 insertions(+), 12 deletions(-) diff --git a/src/megatron/bridge/recipes/qwen_vl/gb300/qwen35_vl.py b/src/megatron/bridge/recipes/qwen_vl/gb300/qwen35_vl.py index 653c752fff..0f859fe5ba 100644 --- a/src/megatron/bridge/recipes/qwen_vl/gb300/qwen35_vl.py +++ b/src/megatron/bridge/recipes/qwen_vl/gb300/qwen35_vl.py @@ -366,28 +366,24 @@ def qwen35_vl_35b_a3b_pretrain_8gpu_gb300_fp8mx_library_config() -> ConfigContai 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 env gate 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. + # 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.env_vars["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] = 1 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 (set below and in the 16-GPU env). The tuned configuration ran this way too. + # 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 both chunk - # sizes to 128 and omits the cuDNN normalization flags), and was never measured with MXFP8. + # 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 - cfg.env_vars["NUM_OF_TOKENS_PER_CHUNK_DISPATCH_API"] = 64 - cfg.env_vars["NUM_OF_TOKENS_PER_CHUNK_PREPROCESSING_API"] = 64 - cfg.env_vars["NVTE_NORM_BWD_USE_CUDNN"] = 1 - cfg.env_vars["NVTE_NORM_FWD_USE_CUDNN"] = 1 # 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 @@ -396,6 +392,27 @@ def qwen35_vl_35b_a3b_pretrain_8gpu_gb300_fp8mx_library_config() -> ConfigContai 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 diff --git a/tests/unit_tests/recipes/test_recipe_environment_defaults.py b/tests/unit_tests/recipes/test_recipe_environment_defaults.py index 40235de6a2..1dc64e04e0 100644 --- a/tests/unit_tests/recipes/test_recipe_environment_defaults.py +++ b/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",