fix(perf): restore non-qwen flat recipe parity

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>

Author: yaoyu-33

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

@@ -17,6 +17,7 @@
     ConfigContainer,
     _benchmark_common,
     _deepseek_v3_common,
+    _enable_deepseek_transformer_engine_graph,
     _perf_precision,
     deepseek_v3_pretrain_config,
     set_deepseek_v3_pipeline_model_parallel_layout,
@@ -82,34 +83,74 @@ def deepseek_v3_pretrain_256gpu_b200_fp8mx_config() -> ConfigContainer:
     _deepseek_v3_common(cfg)
 
     cfg.model.tensor_model_parallel_size = 1
-    cfg.model.pipeline_model_parallel_size = 16
-    cfg.model.virtual_pipeline_model_parallel_size = None
+    cfg.model.pipeline_model_parallel_size = 8
+    cfg.model.virtual_pipeline_model_parallel_size = 2
     cfg.model.context_parallel_size = 1
-    cfg.model.expert_model_parallel_size = 8
+    cfg.model.expert_model_parallel_size = 32
     cfg.model.sequence_parallel = False
     cfg.train.global_batch_size = 4096
     cfg.train.micro_batch_size = 1
 
-    cfg.model.recompute_modules = ["mla_up_proj"]
+    cfg.model.recompute_modules = ["mla_up_proj", "mlp"]
+
+    cfg.model.moe_flex_dispatcher_backend = "deepep"
+    cfg.model.moe_hybridep_num_sms = 16
 
     cfg.ddp.overlap_grad_reduce = True
     cfg.comm_overlap.overlap_grad_reduce = True
+    cfg.comm_overlap.delay_wgrad_compute = True
+    cfg.comm_overlap.overlap_moe_expert_parallel_comm = True
 
     set_deepseek_v3_pipeline_model_parallel_layout(cfg.model)
 
     _benchmark_common(cfg)
+    cfg.model.moe_flex_dispatcher_backend = "deepep"
+    cfg.model.moe_hybridep_num_sms = 16
+    cfg.comm_overlap.delay_wgrad_compute = True
+    cfg.comm_overlap.overlap_moe_expert_parallel_comm = True
+    _enable_deepseek_transformer_engine_graph(cfg)
     return cfg
 
 
 def deepseek_v3_pretrain_256gpu_b200_nvfp4_config() -> ConfigContainer:
     """DeepSeek V3 pretrain: 256× B200, NVFP4 (same layout as BF16)."""
     cfg = deepseek_v3_pretrain_256gpu_b200_bf16_config()
     cfg.mixed_precision = _perf_precision("nvfp4")
+    cfg.mixed_precision.fp4_param = False
+    cfg.mixed_precision.fp4_param_gather = False
+    cfg.model.pipeline_model_parallel_size = 8
+    cfg.model.virtual_pipeline_model_parallel_size = 2
+    cfg.model.expert_model_parallel_size = 32
+    cfg.model.recompute_modules = ["mla_up_proj", "layernorm", "moe_act"]
+    cfg.model.moe_flex_dispatcher_backend = "deepep"
+    cfg.model.moe_hybridep_num_sms = 16
+    cfg.comm_overlap.delay_wgrad_compute = True
+    cfg.comm_overlap.overlap_moe_expert_parallel_comm = True
+    set_deepseek_v3_pipeline_model_parallel_layout(cfg.model)
     return cfg
 
 
 def deepseek_v3_pretrain_256gpu_b200_fp8mx_large_scale_config() -> ConfigContainer:
     """DeepSeek V3 pretrain: 256× B200, MXFP8, large-scale proxy (GBS=256)."""
-    cfg = deepseek_v3_pretrain_256gpu_b200_fp8mx_config()
+    cfg = deepseek_v3_pretrain_config()
+    cfg.mixed_precision = _perf_precision("fp8_mx")
+    _deepseek_v3_common(cfg)
+
+    cfg.model.tensor_model_parallel_size = 1
+    cfg.model.pipeline_model_parallel_size = 16
+    cfg.model.virtual_pipeline_model_parallel_size = None
+    cfg.model.context_parallel_size = 1
+    cfg.model.expert_model_parallel_size = 8
+    cfg.model.sequence_parallel = False
     cfg.train.global_batch_size = 256
+    cfg.train.micro_batch_size = 1
+
+    cfg.model.recompute_modules = ["mla_up_proj"]
+
+    cfg.ddp.overlap_grad_reduce = True
+    cfg.comm_overlap.overlap_grad_reduce = True
+
+    set_deepseek_v3_pipeline_model_parallel_layout(cfg.model)
+
+    _benchmark_common(cfg)
     return cfg

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

@@ -17,6 +17,7 @@
     ConfigContainer,
     _benchmark_common,
     _deepseek_v3_common,
+    _enable_deepseek_transformer_engine_graph,
     _perf_precision,
     deepseek_v3_pretrain_config,
     set_deepseek_v3_pipeline_model_parallel_layout,
@@ -36,7 +37,7 @@ def deepseek_v3_pretrain_256gpu_b300_bf16_config() -> ConfigContainer:
     cfg.model.expert_model_parallel_size = 8
     cfg.model.sequence_parallel = False
     cfg.train.global_batch_size = 4096
-    cfg.train.micro_batch_size = 1
+    cfg.train.micro_batch_size = 2
 
     cfg.model.recompute_modules = ["mla_up_proj"]
 
@@ -46,6 +47,7 @@ def deepseek_v3_pretrain_256gpu_b300_bf16_config() -> ConfigContainer:
     set_deepseek_v3_pipeline_model_parallel_layout(cfg.model)
 
     _benchmark_common(cfg)
+    _enable_deepseek_transformer_engine_graph(cfg)
     return cfg
 
 
@@ -56,13 +58,13 @@ def deepseek_v3_pretrain_256gpu_b300_fp8cs_config() -> ConfigContainer:
     _deepseek_v3_common(cfg)
 
     cfg.model.tensor_model_parallel_size = 1
-    cfg.model.pipeline_model_parallel_size = 16
+    cfg.model.pipeline_model_parallel_size = 8
     cfg.model.virtual_pipeline_model_parallel_size = None
     cfg.model.context_parallel_size = 1
     cfg.model.expert_model_parallel_size = 8
     cfg.model.sequence_parallel = False
     cfg.train.global_batch_size = 4096
-    cfg.train.micro_batch_size = 1
+    cfg.train.micro_batch_size = 2
 
     cfg.model.recompute_modules = ["mla_up_proj"]
 
@@ -72,6 +74,7 @@ def deepseek_v3_pretrain_256gpu_b300_fp8cs_config() -> ConfigContainer:
     set_deepseek_v3_pipeline_model_parallel_layout(cfg.model)
 
     _benchmark_common(cfg)
+    _enable_deepseek_transformer_engine_graph(cfg)
     return cfg
 
 
@@ -82,13 +85,13 @@ def deepseek_v3_pretrain_256gpu_b300_fp8mx_config() -> ConfigContainer:
     _deepseek_v3_common(cfg)
 
     cfg.model.tensor_model_parallel_size = 1
-    cfg.model.pipeline_model_parallel_size = 16
+    cfg.model.pipeline_model_parallel_size = 8
     cfg.model.virtual_pipeline_model_parallel_size = None
     cfg.model.context_parallel_size = 1
     cfg.model.expert_model_parallel_size = 8
     cfg.model.sequence_parallel = False
     cfg.train.global_batch_size = 4096
-    cfg.train.micro_batch_size = 1
+    cfg.train.micro_batch_size = 2
 
     cfg.model.recompute_modules = ["mla_up_proj"]
 
@@ -98,21 +101,42 @@ def deepseek_v3_pretrain_256gpu_b300_fp8mx_config() -> ConfigContainer:
     set_deepseek_v3_pipeline_model_parallel_layout(cfg.model)
 
     _benchmark_common(cfg)
+    _enable_deepseek_transformer_engine_graph(cfg)
     return cfg
 
 
 def deepseek_v3_pretrain_256gpu_b300_nvfp4_config() -> ConfigContainer:
     """DeepSeek V3 pretrain: 256× B300, NVFP4 (PP=16 matching base layout)."""
     cfg = deepseek_v3_pretrain_256gpu_b300_bf16_config()
     cfg.mixed_precision = _perf_precision("nvfp4")
-    cfg.model.pipeline_model_parallel_size = 16
+    cfg.model.pipeline_model_parallel_size = 8
     cfg.model.virtual_pipeline_model_parallel_size = None
     set_deepseek_v3_pipeline_model_parallel_layout(cfg.model)
+    _enable_deepseek_transformer_engine_graph(cfg)
     return cfg
 
 
 def deepseek_v3_pretrain_256gpu_b300_fp8mx_large_scale_config() -> ConfigContainer:
     """DeepSeek V3 pretrain: 256× B300, MXFP8, large-scale proxy (GBS=256)."""
-    cfg = deepseek_v3_pretrain_256gpu_b300_fp8mx_config()
+    cfg = deepseek_v3_pretrain_config()
+    cfg.mixed_precision = _perf_precision("fp8_mx")
+    _deepseek_v3_common(cfg)
+
+    cfg.model.tensor_model_parallel_size = 1
+    cfg.model.pipeline_model_parallel_size = 16
+    cfg.model.virtual_pipeline_model_parallel_size = None
+    cfg.model.context_parallel_size = 1
+    cfg.model.expert_model_parallel_size = 8
+    cfg.model.sequence_parallel = False
     cfg.train.global_batch_size = 256
+    cfg.train.micro_batch_size = 1
+
+    cfg.model.recompute_modules = ["mla_up_proj"]
+
+    cfg.ddp.overlap_grad_reduce = True
+    cfg.comm_overlap.overlap_grad_reduce = True
+
+    set_deepseek_v3_pipeline_model_parallel_layout(cfg.model)
+
+    _benchmark_common(cfg)
     return cfg

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

@@ -35,6 +35,41 @@ def _deepseek_v3_common(cfg: ConfigContainer) -> None:
     cfg.model.moe_router_force_load_balancing = True
 
 
+def _enable_deepseek_full_iteration_mxfp8(
+    cfg: ConfigContainer,
+    *,
+    fp8_dot_product_attention: bool = False,
+    fp8_output_proj: bool = False,
+) -> None:
+    """Apply legacy DeepSeek V3 HybridEP full-iteration MXFP8 settings."""
+    cfg.model.cuda_graph_impl = "full_iteration"
+    cfg.model.cuda_graph_scope = []
+    cfg.model.high_priority_a2a_comm_stream = True
+    cfg.model.moe_expert_rank_capacity_factor = 1.5
+    cfg.model.moe_hybridep_num_sms_preprocessing = 32
+    cfg.model.moe_mlp_glu_interleave_size = 32
+    cfg.model.moe_pad_experts_for_cuda_graph_inference = True
+    cfg.model.moe_paged_stash = True
+    cfg.model.moe_paged_stash_buffer_size_factor_cpu = 1.0
+    cfg.model.moe_paged_stash_buffer_size_factor_cuda = 1.2
+    cfg.model.use_transformer_engine_op_fuser = True
+    cfg.model.fp8_output_proj = fp8_output_proj
+    cfg.model.use_te_rng_tracker = True
+    cfg.rng.te_rng_tracker = True
+
+    cfg.mixed_precision.fp8_dot_product_attention = fp8_dot_product_attention
+    cfg.comm_overlap.delay_wgrad_compute = True
+    cfg.comm_overlap.overlap_moe_expert_parallel_comm = True
+
+
+def _enable_deepseek_transformer_engine_graph(cfg: ConfigContainer) -> None:
+    """Apply legacy DeepSeek V3 Transformer Engine graph capture settings."""
+    cfg.model.cuda_graph_impl = "transformer_engine"
+    cfg.model.cuda_graph_scope = ["attn", "moe_router", "moe_preprocess"]
+    cfg.model.use_te_rng_tracker = True
+    cfg.rng.te_rng_tracker = True
+
+
 def _apply_deepseek_v3_64gpu_gb300_fsdp_configs(cfg: ConfigContainer) -> None:
     """Apply shared DeepSeek V3 64-GPU GB300 Megatron FSDP settings."""
     _deepseek_v3_common(cfg)

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

@@ -17,6 +17,7 @@
     ConfigContainer,
     _benchmark_common,
     _deepseek_v3_common,
+    _enable_deepseek_full_iteration_mxfp8,
     _enable_overlap_param_gather_with_optimizer_step,
     _perf_precision,
     deepseek_v3_pretrain_config,
@@ -99,7 +100,7 @@ def deepseek_v3_pretrain_256gpu_gb200_fp8mx_config() -> ConfigContainer:
     cfg.train.global_batch_size = 4096
     cfg.train.micro_batch_size = 1
 
-    cfg.model.recompute_modules = ["mlp"]
+    cfg.model.recompute_modules = ["mla_up_proj"]
 
     cfg.model.cuda_graph_impl = "transformer_engine"
     cfg.model.cuda_graph_scope = ["attn", "moe_router", "moe_preprocess"]
@@ -110,6 +111,7 @@ def deepseek_v3_pretrain_256gpu_gb200_fp8mx_config() -> ConfigContainer:
     set_deepseek_v3_pipeline_model_parallel_layout(cfg.model)
 
     _benchmark_common(cfg)
+    _enable_deepseek_full_iteration_mxfp8(cfg, fp8_output_proj=True)
     return cfg
 
 
@@ -125,6 +127,29 @@ def deepseek_v3_pretrain_256gpu_gb200_nvfp4_config() -> ConfigContainer:
 
 def deepseek_v3_pretrain_256gpu_gb200_fp8mx_large_scale_config() -> ConfigContainer:
     """DeepSeek V3 pretrain: 256× GB200, MXFP8, large-scale proxy (GBS=256)."""
-    cfg = deepseek_v3_pretrain_256gpu_gb200_fp8mx_config()
+    cfg = deepseek_v3_pretrain_config()
+    cfg.mixed_precision = _perf_precision("fp8_mx")
+    _deepseek_v3_common(cfg)
+
+    cfg.model.tensor_model_parallel_size = 1
+    cfg.model.pipeline_model_parallel_size = 4
+    cfg.model.virtual_pipeline_model_parallel_size = 4
+    cfg.model.context_parallel_size = 1
+    cfg.model.expert_model_parallel_size = 64
+    cfg.model.sequence_parallel = False
     cfg.train.global_batch_size = 256
+    cfg.train.micro_batch_size = 1
+
+    cfg.model.recompute_modules = ["mlp"]
+
+    cfg.model.cuda_graph_impl = "transformer_engine"
+    cfg.model.cuda_graph_scope = ["attn", "moe_router", "moe_preprocess"]
+
+    cfg.ddp.overlap_grad_reduce = True
+    cfg.comm_overlap.overlap_grad_reduce = True
+
+    set_deepseek_v3_pipeline_model_parallel_layout(cfg.model)
+
+    _benchmark_common(cfg)
+    cfg.model.fp8_output_proj = True
     return cfg

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

@@ -18,6 +18,7 @@
     _apply_deepseek_v3_64gpu_gb300_fsdp_configs,
     _benchmark_common,
     _deepseek_v3_common,
+    _enable_deepseek_full_iteration_mxfp8,
     _enable_overlap_param_gather_with_optimizer_step,
     _perf_precision,
     deepseek_v3_pretrain_config,
@@ -96,9 +97,9 @@ def deepseek_v3_pretrain_256gpu_gb300_fp8mx_config() -> ConfigContainer:
     cfg.model.expert_model_parallel_size = 32
     cfg.model.sequence_parallel = False
     cfg.train.global_batch_size = 4096
-    cfg.train.micro_batch_size = 2
+    cfg.train.micro_batch_size = 1
 
-    cfg.model.recompute_modules = ["mla_up_proj"]
+    cfg.model.recompute_modules = []
 
     cfg.model.cuda_graph_scope = []
     cfg.ddp.overlap_grad_reduce = True
@@ -107,6 +108,7 @@ def deepseek_v3_pretrain_256gpu_gb300_fp8mx_config() -> ConfigContainer:
     set_deepseek_v3_pipeline_model_parallel_layout(cfg.model, "Et*4|(t*4|)*14tmL")
 
     _benchmark_common(cfg)
+    _enable_deepseek_full_iteration_mxfp8(cfg, fp8_dot_product_attention=True, fp8_output_proj=True)
     return cfg
 
 
@@ -164,4 +166,7 @@ def deepseek_v3_pretrain_256gpu_gb300_fp8mx_large_scale_config() -> ConfigContai
     cfg = deepseek_v3_pretrain_256gpu_gb300_bf16_config()
     cfg.mixed_precision = _perf_precision("fp8_mx")
     cfg.train.global_batch_size = 256
+    cfg.model.fp8_output_proj = True
+    cfg.comm_overlap.overlap_param_gather_with_optimizer_step = None
+    cfg.optimizer.overlap_param_gather_with_optimizer_step = False
     return cfg

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

@@ -55,16 +55,11 @@ def deepseek_v3_pretrain_1024gpu_h100_fp8cs_config() -> ConfigContainer:
     """DeepSeek V3 pretrain: 1024× H100, FP8 current-scaling."""
     cfg = deepseek_v3_pretrain_config()
     cfg.mixed_precision = _perf_precision("fp8_cs")
-    cfg.mixed_precision.fp8_recipe = "blockwise"
-    cfg.mixed_precision.fp8_param = False
-    cfg.mixed_precision.fp8_param_gather = False
-    cfg.mixed_precision.num_layers_at_start_in_bf16 = 0
-    cfg.mixed_precision.num_layers_at_end_in_bf16 = 0
     _deepseek_v3_common(cfg)
 
     cfg.model.tensor_model_parallel_size = 2
     cfg.model.pipeline_model_parallel_size = 8
-    cfg.model.virtual_pipeline_model_parallel_size = 2
+    cfg.model.virtual_pipeline_model_parallel_size = 4
     cfg.model.context_parallel_size = 1
     cfg.model.expert_model_parallel_size = 64
     cfg.model.sequence_parallel = True
@@ -76,7 +71,7 @@ def deepseek_v3_pretrain_1024gpu_h100_fp8cs_config() -> ConfigContainer:
     cfg.ddp.overlap_grad_reduce = False
     cfg.comm_overlap.overlap_grad_reduce = False
 
-    cfg.model.pipeline_model_parallel_layout = None
+    set_deepseek_v3_pipeline_model_parallel_layout(cfg.model, "Et|(tt|)*30mL")
 
     _benchmark_common(cfg)
     _enable_overlap_param_gather_with_optimizer_step(cfg)
@@ -86,6 +81,11 @@ def deepseek_v3_pretrain_1024gpu_h100_fp8cs_config() -> ConfigContainer:
 def deepseek_v3_pretrain_1024gpu_h100_fp8sc_config() -> ConfigContainer:
     """DeepSeek V3 pretrain: 1024× H100, FP8-SC (VP=2, auto-applied default PP layout)."""
     cfg = deepseek_v3_pretrain_1024gpu_h100_fp8cs_config()
+    cfg.mixed_precision.fp8_recipe = "blockwise"
+    cfg.mixed_precision.fp8_param = False
+    cfg.mixed_precision.fp8_param_gather = False
+    cfg.mixed_precision.num_layers_at_start_in_bf16 = 0
+    cfg.mixed_precision.num_layers_at_end_in_bf16 = 0
     cfg.model.virtual_pipeline_model_parallel_size = 2
     # DeepSeek-V3 has 61 layers; (61 // PP) % VP != 0 for (PP=8, VP=2), so a custom layout
     # is required. The helper's default layout map provides one for this (pp, vp) pair.