perf(recipes): optimize DeepSeek V4 B200 topology

Signed-off-by: Chen Cui <chcui@nvidia.com>

Author: cuichenx

examples/models/deepseek_v4/README.md

@@ -80,8 +80,10 @@ Available Blackwell pretraining recipes:
 - `deepseek_v4_flash_pretrain_mxfp8_config`: Adam MXFP8
 - `deepseek_v4_flash_pretrain_muon_config`: Muon BF16
 - `deepseek_v4_flash_pretrain_64gpu_b200_fp8mx_library_config`: 64-GPU B200
-  Adam MXFP8 with PP8/VPP2/EP8, all-to-all dispatch, selective recompute, and
-  attention activation offload plus MXFP8 parameter gather/buffer reuse
+  Adam MXFP8 with PP4/VPP4/EP16, HybridEP across pairs of physical NVL8
+  domains, selective `[mhc,mla_up_proj]` recompute, no activation offload, and
+  MXFP8 parameter gather/buffer reuse. This recipe requires eight physical
+  B200 NVL8 systems with eight contiguous, node-major ranks per system.
 - `deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config`: 128-GPU GB200
   Adam MXFP8 with PP1/EP64/dense-DP128/expert-DP2/HybridEP, selective recompute,
   attention activation offload, and MXFP8 parameter gather/buffer reuse
@@ -101,6 +103,14 @@ TP/PP/EP/CP settings. Launch them through a site multi-node runner that invokes
 Keep the topology and global batch size owned by the recipe; pass only dataset,
 run-length, logging, and checkpoint overrides.
 
+For the B200 recipe, preserve node-major global-rank assignment so each EP16
+group spans two physical NVL8 systems with eight ranks in each NVLink domain.
+Its runtime must provide an MCore-compatible HybridEP package exposing
+`deep_ep.HybridEPBuffer` with cross-domain DMA-BUF registration and completion
+notification support. The site launcher must also provide the topology-correct
+NIC mapping required by that package. Megatron Bridge's `uv sync` does not
+install or configure this external runtime.
+
 The GB200 Flash and GB300 Pro library bases enable their supported DSA and
 Transformer Engine fused grouped-MLP paths. The new hardware-count-specific
 variants additionally enable GLU interleaving; the existing recipe identities
@@ -109,8 +119,9 @@ natural-routing capacity. All variants retain their precision-specific
 training policies: full-iteration CUDA graphs, FP8 parameter gather/buffer
 reuse, and reduced-precision gradient reduction remain disabled except for the
 provisional FP8 parameter gather/buffer reuse in the high-scale B200 and GB200
-Flash recipes; both configurations still require 100-step checkpoint and resume
-validation.
+Flash recipes. These configurations require 100-step natural-routing training,
+full model/optimizer/RNG checkpoint save, and fresh-process resume validation on
+the documented Megatron-LM `dev` stack.
 
 `slurm_pretrain.sh` is a GB200 launcher with `TP=1,PP=4,EP=8,CP=1` by default. Indexer loss are disabled for now and is planned for a follow-up.
 

src/megatron/bridge/recipes/deepseek/b200/deepseek_v4.py

@@ -19,51 +19,53 @@
 from megatron.bridge.training.config import ConfigContainer
 
 
-_DSV4_FLASH_PP8_VP2_LAYOUT = "Et*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*2|t*2|t*2|t*2|t*2mL"
-_FLEX_DISPATCHER_ENV_VARS = {
-    "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN",
-    "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API",
-    "NVLINK_DOMAIN_SIZE",
-    "USE_MNNVL",
-}
+_DSV4_FLASH_PP4_VP4_LAYOUT = "Et*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*2|t*2|t*2|t*2|t*2mL"
 
 
 def deepseek_v4_flash_pretrain_64gpu_b200_fp8mx_library_config() -> ConfigContainer:
     """Return real-training DeepSeek V4 Flash for 64 B200 GPUs.
 
-    PP8/VPP2 keeps each eight-rank expert group within one NVL8 system. The
-    recipe uses the portable all-to-all dispatcher and preserves natural,
-    unlimited-capacity routing without paged stash or CUDA graphs.
+    Requires eight physical NVL8 systems with eight contiguous, node-major
+    ranks per system. Each EP16 group spans two NVL8 domains, with eight ranks
+    per domain. The runtime must expose ``deep_ep.HybridEPBuffer``. Natural,
+    unlimited-capacity routing is preserved without paged stash, activation
+    offload, or CUDA graphs.
     """
     cfg = deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config()
 
     cfg.model.tensor_model_parallel_size = 1
-    cfg.model.pipeline_model_parallel_size = 8
-    cfg.model.virtual_pipeline_model_parallel_size = 2
+    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 = 8
+    cfg.model.expert_model_parallel_size = 16
     cfg.model.expert_tensor_parallel_size = 1
     cfg.model.sequence_parallel = False
-    cfg.model.pipeline_model_parallel_layout = _DSV4_FLASH_PP8_VP2_LAYOUT
+    cfg.model.pipeline_model_parallel_layout = _DSV4_FLASH_PP4_VP4_LAYOUT
 
-    cfg.model.moe_token_dispatcher_type = "alltoall"
-    cfg.model.moe_flex_dispatcher_backend = None
-    cfg.model.moe_flex_dispatcher_num_sms = None
+    cfg.model.moe_token_dispatcher_type = "flex"
+    cfg.model.moe_flex_dispatcher_backend = "hybridep"
+    cfg.model.moe_flex_dispatcher_num_sms = 32
     cfg.model.moe_deepep_num_sms = None
     cfg.model.moe_hybridep_num_sms = None
-    cfg.model.moe_hybridep_num_sms_preprocessing = None
+    cfg.model.moe_hybridep_num_sms_preprocessing = 108
     cfg.model.moe_shared_expert_overlap = False
 
-    cfg.model.recompute_modules = ["moe", "mhc", "mla_up_proj", "layernorm"]
-    cfg.model.fine_grained_activation_offloading = True
-    cfg.model.offload_modules = ["core_attn", "attn_proj"]
-    cfg.model.fine_grained_offloading_max_inflight_offloads = 2
+    cfg.model.recompute_modules = ["mhc", "mla_up_proj"]
+    cfg.model.fine_grained_activation_offloading = False
+    cfg.model.offload_modules = []
+    cfg.model.fine_grained_offloading_max_inflight_offloads = None
     cfg.model.moe_pad_experts_for_cuda_graph_inference = False
     cfg.model.cuda_graph_impl = "none"
     cfg.model.cuda_graph_modules = []
     cfg.model.cuda_graph_scope = None
     cfg.model.use_te_rng_tracker = False
     cfg.rng.te_rng_tracker = False
 
-    cfg.env_vars = {key: value for key, value in cfg.env_vars.items() if key not in _FLEX_DISPATCHER_ENV_VARS}
+    cfg.env_vars = {
+        **{key: value for key, value in cfg.env_vars.items() if key != "NVTE_CPU_OFFLOAD_V1"},
+        "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 8,
+        "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128,
+        "NVLINK_DOMAIN_SIZE": 8,
+        "USE_MNNVL": 0,
+    }
     return cfg

src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py

@@ -44,8 +44,6 @@ def deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config() -> ConfigContainer:
     All features validated in 1k-step convergence runs on OCI GB200 NVL72:
     fused DSA indexer, DSA indexer loss, HybridEP dispatcher, GroupedGEMM,
     permute fusion, selective recompute with mla_up_proj.
-
-    Note: set checkpoint.save_optim=False when using HybridEP.
     """
     use_fused_mhc = deepseek_v4_supports_blackwell_fused_kernels()
     cfg = _pretrain_common()
@@ -145,10 +143,7 @@ def deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config() -> ConfigContainer:
 
 
 def deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config() -> ConfigContainer:
-    """DeepSeek-V4-Flash MXFP8 pre-training on 64-GPU GB200 (PP=8, EP=8).
-
-    Note: set checkpoint.save_optim=False when using HybridEP.
-    """
+    """DeepSeek-V4-Flash MXFP8 pre-training on 64-GPU GB200 (PP=8, EP=8)."""
     cfg = deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config()
 
     cfg.model.dsa_indexer_loss_coeff = 0.0

tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py

@@ -175,6 +175,10 @@ def test_flash_high_scale_recipe_preserves_real_training_contract() -> None:
     assert cfg.ddp.grad_reduce_in_fp32 is True
     assert cfg.optimizer.main_grads_dtype == torch.float32
     assert cfg.ddp.check_for_nan_in_grad is True
+    assert cfg.checkpoint.save_optim is True
+    assert cfg.checkpoint.save_rng is True
+    assert cfg.checkpoint.load_optim is True
+    assert cfg.checkpoint.load_rng is True
     assert cfg.model.cuda_graph_impl == "none"
     assert cfg.model.use_transformer_engine_op_fuser is True
     assert cfg.model.moe_mlp_glu_interleave_size == 32
@@ -198,9 +202,9 @@ def test_flash_b200_nvl8_recipe_preserves_real_training_contract() -> None:
     assert cfg.train.train_iters == 1_000_000
     assert cfg.train.global_batch_size == 256
     assert cfg.model.tensor_model_parallel_size == 1
-    assert cfg.model.pipeline_model_parallel_size == 8
-    assert cfg.model.virtual_pipeline_model_parallel_size == 2
-    assert cfg.model.expert_model_parallel_size == 8
+    assert cfg.model.pipeline_model_parallel_size == 4
+    assert cfg.model.virtual_pipeline_model_parallel_size == 4
+    assert cfg.model.expert_model_parallel_size == 16
     assert cfg.model.pipeline_model_parallel_layout == (
         "Et*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*2|t*2|t*2|t*2|t*2mL"
     )
@@ -209,29 +213,43 @@ def test_flash_b200_nvl8_recipe_preserves_real_training_contract() -> None:
         pipeline_model_parallel_size=cfg.model.pipeline_model_parallel_size,
     )
     parsed_layout.validate_layer_layout(num_layers=43, mtp_num_layers=1)
-    assert parsed_layout.virtual_pipeline_model_parallel_size == 2
-    assert cfg.model.moe_token_dispatcher_type == "alltoall"
-    assert cfg.model.moe_flex_dispatcher_backend is None
-    assert cfg.model.moe_flex_dispatcher_num_sms is None
+    assert parsed_layout.virtual_pipeline_model_parallel_size == 4
+    assert cfg.model.moe_token_dispatcher_type == "flex"
+    assert cfg.model.moe_flex_dispatcher_backend == "hybridep"
+    assert cfg.model.moe_flex_dispatcher_num_sms == 32
+    assert cfg.model.moe_deepep_num_sms is None
+    assert cfg.model.moe_hybridep_num_sms is None
+    assert cfg.model.moe_hybridep_num_sms_preprocessing == 108
+    assert cfg.model.moe_shared_expert_overlap is False
     assert cfg.model.moe_router_force_load_balancing is False
-    assert cfg.model.recompute_modules == ["moe", "mhc", "mla_up_proj", "layernorm"]
-    assert cfg.model.fine_grained_activation_offloading is True
-    assert cfg.model.offload_modules == ["core_attn", "attn_proj"]
-    assert cfg.model.fine_grained_offloading_max_inflight_offloads == 2
+    assert getattr(cfg.model, "moe_router_force_biased", None) is None
+    assert getattr(cfg.model, "moe_expert_capacity_factor", None) is None
+    assert cfg.model.recompute_modules == ["mhc", "mla_up_proj"]
+    assert cfg.model.recompute_granularity == "selective"
+    assert cfg.model.recompute_method is None
+    assert cfg.model.recompute_num_layers is None
+    assert cfg.model.fine_grained_activation_offloading is False
+    assert cfg.model.offload_modules == []
+    assert cfg.model.fine_grained_offloading_max_inflight_offloads is None
     assert cfg.mixed_precision.fp8_param_gather is True
     assert cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag is True
     assert getattr(cfg.model, "moe_expert_rank_capacity_factor", None) is None
+    assert getattr(cfg.model, "moe_token_dropping", False) is False
     assert getattr(cfg.model, "moe_paged_stash", False) is False
     assert cfg.model.cuda_graph_impl == "none"
     assert cfg.model.cuda_graph_modules == []
     assert cfg.model.cuda_graph_scope is None
     assert cfg.rerun_state_machine.check_for_nan_in_loss is True
     assert cfg.ddp.check_for_nan_in_grad is True
-    assert "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN" not in cfg.env_vars
-    assert "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API" not in cfg.env_vars
-    assert "NVLINK_DOMAIN_SIZE" not in cfg.env_vars
-    assert "USE_MNNVL" not in cfg.env_vars
-    assert cfg.env_vars["NVTE_CPU_OFFLOAD_V1"] == 1
+    assert cfg.checkpoint.save_optim is True
+    assert cfg.checkpoint.save_rng is True
+    assert cfg.checkpoint.load_optim is True
+    assert cfg.checkpoint.load_rng is True
+    assert cfg.env_vars["NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN"] == 8
+    assert cfg.env_vars["NUM_OF_TOKENS_PER_CHUNK_COMBINE_API"] == 128
+    assert cfg.env_vars["NVLINK_DOMAIN_SIZE"] == 8
+    assert cfg.env_vars["USE_MNNVL"] == 0
+    assert "NVTE_CPU_OFFLOAD_V1" not in cfg.env_vars
 
 
 def test_pro_high_scale_recipe_preserves_real_training_contract() -> None: