[perf] feat: add Qwen3 and DeepSeek-V3 MoE recipes

src/megatron/bridge/perf_recipes/deepseek/__init__.py

@@ -17,6 +17,7 @@
     deepseek_v3_pretrain_256gpu_gb200_fp8cs_config,
     deepseek_v3_pretrain_256gpu_gb200_fp8mx_config,
     deepseek_v3_pretrain_256gpu_gb200_fp8mx_large_scale_config,
+    deepseek_v3_256gpu_gb200_fp8mx_partial_cg_pretrain_config,
     deepseek_v3_pretrain_256gpu_gb200_nvfp4_config,
 )
 from megatron.bridge.perf_recipes.deepseek.gb300.deepseek_v3 import (
@@ -26,6 +27,7 @@
     deepseek_v3_pretrain_256gpu_gb300_fp8cs_config,
     deepseek_v3_pretrain_256gpu_gb300_fp8mx_config,
     deepseek_v3_pretrain_256gpu_gb300_fp8mx_large_scale_config,
+    deepseek_v3_256gpu_gb300_fp8mx_partial_cg_pretrain_config,
     deepseek_v3_pretrain_256gpu_gb300_nvfp4_config,
 )
 from megatron.bridge.perf_recipes.deepseek.h100.deepseek_v3 import (

src/megatron/bridge/perf_recipes/deepseek/common.py

@@ -14,6 +14,8 @@
 # ruff: noqa: F401
 """Common helpers for deepseek performance recipes."""
 
+import torch
+
 from megatron.bridge.perf_recipes._common import (
     _benchmark_common,
     _enable_overlap_param_gather_with_optimizer_step,
@@ -35,6 +37,15 @@ def _deepseek_v3_common(cfg: ConfigContainer) -> None:
     cfg.model.moe_router_force_load_balancing = True
 
 
+def _enable_deepseek_precision_aware_optimizer(cfg: ConfigContainer) -> None:
+    """Enable the optimizer precision settings used by tuned DeepSeek recipes."""
+    cfg.optimizer.use_precision_aware_optimizer = True
+    cfg.optimizer.main_grads_dtype = torch.float32
+    cfg.optimizer.main_params_dtype = torch.float32
+    cfg.optimizer.exp_avg_dtype = torch.bfloat16
+    cfg.optimizer.exp_avg_sq_dtype = torch.bfloat16
+
+
 def _enable_deepseek_full_iteration_mxfp8(
     cfg: ConfigContainer,
     *,

src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v3.py

@@ -18,6 +18,7 @@
     _benchmark_common,
     _deepseek_v3_common,
     _enable_deepseek_full_iteration_mxfp8,
+    _enable_deepseek_precision_aware_optimizer,
     _enable_overlap_param_gather_with_optimizer_step,
     _perf_precision,
     deepseek_v3_pretrain_config,
@@ -153,3 +154,49 @@ def deepseek_v3_pretrain_256gpu_gb200_fp8mx_large_scale_config() -> ConfigContai
     _benchmark_common(cfg)
     cfg.model.fp8_output_proj = True
     return cfg
+
+
+def deepseek_v3_256gpu_gb200_fp8mx_partial_cg_pretrain_config() -> ConfigContainer:
+    """DeepSeek V3 pretrain: 256× GB200, MXFP8, scoped CUDA graphs and offloading."""
+    cfg = deepseek_v3_pretrain_config()
+    cfg.mixed_precision = _perf_precision("fp8_mx")
+    _benchmark_common(cfg)
+
+    cfg.model.tensor_model_parallel_size = 1
+    cfg.model.pipeline_model_parallel_size = 4
+    cfg.model.virtual_pipeline_model_parallel_size = 2
+    cfg.model.expert_model_parallel_size = 64
+    cfg.model.expert_tensor_parallel_size = 1
+    cfg.model.context_parallel_size = 1
+    cfg.model.kv_channels = 128
+    cfg.model.make_vocab_size_divisible_by = 1280
+    cfg.model.moe_router_force_load_balancing = True
+    cfg.model.moe_router_fusion = True
+    cfg.model.moe_flex_dispatcher_backend = "hybridep"
+    cfg.model.moe_token_dispatcher_type = "flex"
+    cfg.model.moe_router_padding_for_quantization = True
+    cfg.model.moe_hybridep_num_sms = 32
+    cfg.model.recompute_granularity = "selective"
+    cfg.model.recompute_modules = ["mla_up_proj"]
+    cfg.model.fine_grained_activation_offloading = True
+    cfg.model.offload_modules = ["expert_fc1"]
+    cfg.model.cuda_graph_impl = "transformer_engine"
+    cfg.model.cuda_graph_scope = ["attn", "moe_router", "moe_preprocess"]
+    cfg.model.cuda_graph_warmup_steps = 1
+    cfg.model.use_te_rng_tracker = True
+
+    cfg.comm_overlap.delay_wgrad_compute = True
+    cfg.comm_overlap.overlap_moe_expert_parallel_comm = True
+    cfg.ddp.reuse_grad_buf_for_mxfp8_param_ag = True
+    cfg.ddp.overlap_grad_reduce = True
+    cfg.ddp.overlap_param_gather = True
+
+    cfg.train.micro_batch_size = 1
+    cfg.train.global_batch_size = 8192
+    cfg.train.exit_duration_in_mins = 220
+    cfg.train.manual_gc_interval = 10
+
+    cfg.rng.te_rng_tracker = True
+    set_deepseek_v3_pipeline_model_parallel_layout(cfg.model, "Et*4|(tttt|)*14tmL")
+    _enable_deepseek_precision_aware_optimizer(cfg)
+    return cfg

src/megatron/bridge/perf_recipes/deepseek/gb300/deepseek_v3.py

@@ -19,6 +19,7 @@
     _benchmark_common,
     _deepseek_v3_common,
     _enable_deepseek_full_iteration_mxfp8,
+    _enable_deepseek_precision_aware_optimizer,
     _enable_overlap_param_gather_with_optimizer_step,
     _perf_precision,
     deepseek_v3_pretrain_config,
@@ -170,3 +171,46 @@ def deepseek_v3_pretrain_256gpu_gb300_fp8mx_large_scale_config() -> ConfigContai
     cfg.comm_overlap.overlap_param_gather_with_optimizer_step = None
     cfg.optimizer.overlap_param_gather_with_optimizer_step = False
     return cfg
+
+
+def deepseek_v3_256gpu_gb300_fp8mx_partial_cg_pretrain_config() -> ConfigContainer:
+    """DeepSeek V3 pretrain: 256× GB300, MXFP8 and scoped CUDA graphs."""
+    cfg = deepseek_v3_pretrain_config()
+    cfg.mixed_precision = _perf_precision("fp8_mx")
+    _benchmark_common(cfg)
+
+    cfg.model.tensor_model_parallel_size = 1
+    cfg.model.pipeline_model_parallel_size = 4
+    cfg.model.virtual_pipeline_model_parallel_size = 2
+    cfg.model.expert_model_parallel_size = 64
+    cfg.model.expert_tensor_parallel_size = 1
+    cfg.model.context_parallel_size = 1
+    cfg.model.kv_channels = 128
+    cfg.model.make_vocab_size_divisible_by = 1280
+    cfg.model.moe_router_force_load_balancing = True
+    cfg.model.moe_router_fusion = True
+    cfg.model.moe_flex_dispatcher_backend = "hybridep"
+    cfg.model.moe_token_dispatcher_type = "flex"
+    cfg.model.moe_router_padding_for_quantization = True
+    cfg.model.moe_hybridep_num_sms = 32
+    cfg.model.cuda_graph_impl = "transformer_engine"
+    cfg.model.cuda_graph_scope = ["attn", "moe_router", "moe_preprocess"]
+    cfg.model.cuda_graph_warmup_steps = 1
+    cfg.model.use_te_rng_tracker = True
+
+    cfg.comm_overlap.delay_wgrad_compute = True
+    cfg.comm_overlap.overlap_moe_expert_parallel_comm = True
+    cfg.ddp.reuse_grad_buf_for_mxfp8_param_ag = True
+    cfg.ddp.overlap_grad_reduce = True
+    cfg.ddp.overlap_param_gather = True
+    cfg.ddp.check_for_nan_in_grad = False
+
+    cfg.train.micro_batch_size = 1
+    cfg.train.global_batch_size = 8192
+    cfg.train.exit_duration_in_mins = 220
+    cfg.train.manual_gc_interval = 10
+
+    cfg.rng.te_rng_tracker = True
+    set_deepseek_v3_pipeline_model_parallel_layout(cfg.model, "Et*4|(tttt|)*14tmL")
+    _enable_deepseek_precision_aware_optimizer(cfg)
+    return cfg

src/megatron/bridge/perf_recipes/qwen/__init__.py

@@ -8,6 +8,7 @@
     qwen3_235b_a22b_pretrain_64gpu_b200_fp8cs_config,
     qwen3_235b_a22b_pretrain_64gpu_b200_fp8mx_config,
     qwen3_235b_a22b_pretrain_64gpu_b200_nvfp4_config,
+    qwen3_235b_a22b_128gpu_b200_fp8mx_deepep_pretrain_config,
     qwen3_235b_a22b_pretrain_256gpu_b200_bf16_config,
     qwen3_235b_a22b_pretrain_256gpu_b200_fp8cs_config,
     qwen3_235b_a22b_pretrain_256gpu_b200_fp8mx_config,
@@ -24,6 +25,7 @@
     qwen3_235b_a22b_pretrain_64gpu_b300_fp8cs_config,
     qwen3_235b_a22b_pretrain_64gpu_b300_fp8mx_config,
     qwen3_235b_a22b_pretrain_64gpu_b300_nvfp4_config,
+    qwen3_235b_a22b_128gpu_b300_fp8mx_hybridep_pretrain_config,
     qwen3_235b_a22b_pretrain_256gpu_b300_bf16_config,
     qwen3_235b_a22b_pretrain_256gpu_b300_fp8cs_config,
     qwen3_235b_a22b_pretrain_256gpu_b300_fp8mx_config,
@@ -36,12 +38,17 @@
     qwen3_30b_a3b_pretrain_8gpu_gb200_bf16_config,
     qwen3_30b_a3b_pretrain_8gpu_gb200_fp8cs_config,
     qwen3_30b_a3b_pretrain_8gpu_gb200_fp8mx_config,
+    qwen3_30b_a3b_16gpu_gb200_bf16_pretrain_config,
+    qwen3_30b_a3b_16gpu_gb200_fp8mx_paged_stash_pretrain_config,
+    qwen3_30b_a3b_16gpu_gb200_fp8mx_partial_cg_pretrain_config,
     qwen3_30b_a3b_pretrain_32gpu_gb200_bf16_config,
     qwen3_30b_a3b_pretrain_32gpu_gb200_fp8cs_config,
     qwen3_235b_a22b_pretrain_64gpu_gb200_bf16_config,
     qwen3_235b_a22b_pretrain_64gpu_gb200_fp8cs_config,
     qwen3_235b_a22b_pretrain_64gpu_gb200_fp8mx_config,
+    qwen3_235b_a22b_128gpu_gb200_fp8mx_paged_stash_pretrain_config,
     qwen3_235b_a22b_pretrain_64gpu_gb200_nvfp4_config,
+    qwen3_235b_a22b_128gpu_gb200_fp8mx_partial_cg_pretrain_config,
     qwen3_235b_a22b_pretrain_256gpu_gb200_bf16_config,
     qwen3_235b_a22b_pretrain_256gpu_gb200_fp8cs_config,
     qwen3_235b_a22b_pretrain_256gpu_gb200_fp8mx_config,
@@ -60,6 +67,7 @@
     qwen3_235b_a22b_pretrain_64gpu_gb300_fp8cs_config,
     qwen3_235b_a22b_pretrain_64gpu_gb300_fp8mx_config,
     qwen3_235b_a22b_pretrain_64gpu_gb300_nvfp4_config,
+    qwen3_235b_a22b_128gpu_gb300_fp8mx_paged_stash_pretrain_config,
     qwen3_235b_a22b_pretrain_256gpu_gb300_bf16_config,
     qwen3_235b_a22b_pretrain_256gpu_gb300_fp8cs_config,
     qwen3_235b_a22b_pretrain_256gpu_gb300_fp8mx_config,
@@ -71,9 +79,12 @@
 from megatron.bridge.perf_recipes.qwen.h100.qwen3_moe import (
     qwen3_30b_a3b_pretrain_16gpu_h100_bf16_config,
     qwen3_30b_a3b_pretrain_16gpu_h100_fp8cs_config,
+    qwen3_30b_a3b_32gpu_h100_bf16_pretrain_config,
+    qwen3_30b_a3b_32gpu_h100_fp8sc_pretrain_config,
     qwen3_30b_a3b_pretrain_64gpu_h100_bf16_config,
     qwen3_30b_a3b_pretrain_64gpu_h100_fp8cs_config,
     qwen3_235b_a22b_pretrain_256gpu_h100_bf16_config,
+    qwen3_235b_a22b_256gpu_h100_bf16_hybridep_pretrain_config,
     qwen3_235b_a22b_pretrain_256gpu_h100_fp8cs_config,
     qwen3_235b_a22b_pretrain_256gpu_h100_fp8cs_large_scale_config,
     qwen3_next_80b_a3b_pretrain_128gpu_h100_bf16_config,

src/megatron/bridge/perf_recipes/qwen/b200/qwen3_moe.py

@@ -19,8 +19,10 @@
 from megatron.bridge.perf_recipes.qwen.common import (
     CommOverlapConfig,
     ConfigContainer,
+    _apply_qwen3_moe_tuned_defaults,
     _benchmark_common,
     _enable_overlap_param_gather_with_optimizer_step,
+    _enable_qwen_precision_aware_optimizer,
     _perf_precision,
     _with_global_batch_size,
     qwen3_30b_a3b_pretrain_config,
@@ -365,3 +367,35 @@ def qwen3_next_80b_a3b_pretrain_64gpu_b200_fp8mx_config() -> ConfigContainer:
     cfg = qwen3_next_80b_a3b_pretrain_64gpu_b300_bf16_config()
     cfg.mixed_precision = _perf_precision("fp8_mx")
     return cfg
+
+
+def qwen3_235b_a22b_128gpu_b200_fp8mx_deepep_pretrain_config() -> ConfigContainer:
+    """Qwen3 235B-A22B pretrain: 128× B200, MXFP8 and tuned DeepEP overlap."""
+    cfg = qwen3_235b_a22b_pretrain_config()
+    cfg.mixed_precision = _perf_precision("fp8_mx")
+    _apply_qwen3_moe_tuned_defaults(cfg, original_max_position_embeddings=4096)
+
+    cfg.model.tensor_model_parallel_size = 2
+    cfg.model.pipeline_model_parallel_size = 4
+    cfg.model.virtual_pipeline_model_parallel_size = 2
+    cfg.model.expert_model_parallel_size = 16
+    cfg.model.expert_tensor_parallel_size = 1
+    cfg.model.context_parallel_size = 1
+    cfg.model.sequence_parallel = True
+    cfg.model.moe_flex_dispatcher_backend = "deepep"
+    cfg.model.moe_router_padding_for_quantization = True
+    cfg.model.moe_deepep_num_sms = 20
+    cfg.model.overlap_p2p_comm = True
+    cfg.model.batch_p2p_comm = False
+
+    cfg.comm_overlap = CommOverlapConfig(
+        tp_comm_overlap=False,
+        overlap_moe_expert_parallel_comm=True,
+        delay_wgrad_compute=True,
+    )
+    cfg.train.micro_batch_size = 3
+    cfg.train.global_batch_size = 3072
+    cfg.ddp.reuse_grad_buf_for_mxfp8_param_ag = True
+    cfg.checkpoint.dist_ckpt_strictness = "log_all"
+    _enable_qwen_precision_aware_optimizer(cfg)
+    return cfg

src/megatron/bridge/perf_recipes/qwen/b300/qwen3_moe.py

@@ -16,8 +16,10 @@
 from megatron.bridge.perf_recipes.qwen.common import (
     CommOverlapConfig,
     ConfigContainer,
+    _apply_qwen3_moe_tuned_defaults,
     _benchmark_common,
     _enable_hybridep_full_iteration_mxfp8,
+    _enable_qwen_precision_aware_optimizer,
     _perf_precision,
     qwen3_30b_a3b_pretrain_config,
     qwen3_235b_a22b_pretrain_config,
@@ -370,3 +372,37 @@ def qwen3_next_80b_a3b_pretrain_64gpu_b300_fp8mx_config() -> ConfigContainer:
     cfg.mixed_precision = _perf_precision("fp8_mx")
     cfg.train.micro_batch_size = 2
     return cfg
+
+
+def qwen3_235b_a22b_128gpu_b300_fp8mx_hybridep_pretrain_config() -> ConfigContainer:
+    """Qwen3 235B-A22B pretrain: 128 × B300, MXFP8, HybridEP and scoped CUDA graphs."""
+    cfg = qwen3_235b_a22b_pretrain_config()
+    cfg.mixed_precision = _perf_precision("fp8_mx")
+    _apply_qwen3_moe_tuned_defaults(cfg, original_max_position_embeddings=4096)
+
+    cfg.model.tensor_model_parallel_size = 2
+    cfg.model.pipeline_model_parallel_size = 1
+    cfg.model.expert_model_parallel_size = 16
+    cfg.model.expert_tensor_parallel_size = 1
+    cfg.model.context_parallel_size = 1
+    cfg.model.sequence_parallel = True
+    cfg.model.moe_flex_dispatcher_backend = "hybridep"
+    cfg.model.moe_router_padding_for_quantization = True
+    cfg.model.moe_hybridep_num_sms = 32
+    cfg.model.cuda_graph_impl = "transformer_engine"
+    cfg.model.cuda_graph_scope = ["attn", "moe_router"]
+    cfg.model.cuda_graph_warmup_steps = 2
+    cfg.model.use_te_rng_tracker = True
+
+    cfg.comm_overlap = CommOverlapConfig(
+        tp_comm_overlap=False,
+        overlap_moe_expert_parallel_comm=True,
+        delay_wgrad_compute=True,
+    )
+    cfg.train.micro_batch_size = 3
+    cfg.train.global_batch_size = 9216
+    cfg.ddp.reuse_grad_buf_for_mxfp8_param_ag = True
+    cfg.checkpoint.dist_ckpt_strictness = "log_all"
+    cfg.rng.te_rng_tracker = True
+    _enable_qwen_precision_aware_optimizer(cfg)
+    return cfg

src/megatron/bridge/perf_recipes/qwen/common.py

@@ -14,6 +14,8 @@
 # ruff: noqa: F401
 """Common helpers for qwen performance recipes."""
 
+import torch
+
 from megatron.bridge.perf_recipes._common import (
     _benchmark_common,
     _enable_overlap_param_gather_with_optimizer_step,
@@ -33,6 +35,35 @@ def _with_global_batch_size(cfg: ConfigContainer, global_batch_size: int) -> Con
     return cfg
 
 
+def _apply_qwen3_moe_tuned_defaults(
+    cfg: ConfigContainer,
+    *,
+    original_max_position_embeddings: int,
+) -> None:
+    """Apply defaults shared by the hardware-tuned Qwen3 MoE recipes."""
+    _benchmark_common(cfg)
+
+    cfg.model.recompute_granularity = None
+    cfg.model.recompute_method = None
+    cfg.model.recompute_num_layers = None
+    cfg.model.yarn_original_max_position_embeddings = original_max_position_embeddings
+    cfg.model.make_vocab_size_divisible_by = 1187
+    cfg.model.moe_router_force_load_balancing = True
+    cfg.model.moe_router_fusion = True
+    cfg.model.moe_router_dtype = torch.float32
+    cfg.model.moe_token_dispatcher_type = "flex"
+
+    cfg.dataset.seq_length = cfg.model.seq_length
+    cfg.train.manual_gc_interval = 5
+
+
+def _enable_qwen_precision_aware_optimizer(cfg: ConfigContainer) -> None:
+    """Enable the BF16 optimizer-state settings used by tuned Qwen3 MoE recipes."""
+    cfg.optimizer.use_precision_aware_optimizer = True
+    cfg.optimizer.exp_avg_dtype = torch.bfloat16
+    cfg.optimizer.exp_avg_sq_dtype = torch.bfloat16
+
+
 def _enable_hybridep_full_iteration_mxfp8(cfg: ConfigContainer) -> None:
     cfg.model.cuda_graph_impl = "full_iteration"
     cfg.model.cuda_graph_scope = []