feat(recipes): modernize DeepSeek V4 library configs

examples/models/deepseek_v4/README.md

@@ -6,7 +6,13 @@ The bridge supports four published variants out of the same code path. The on-di
 
 ## MCore Checkout
 
-The pretraining recipes were tested with Megatron-LM `dev` commit `35f36c7c9dba` plus PR [#4839](https://github.com/NVIDIA/Megatron-LM/pull/4839) (`f04b762406f0` in the OCI test checkout). The Megatron-LM copy inside the current NeMo FW container is not expected to work for these recipes.
+DeepSeek V4 pretraining requires Megatron-LM `dev`; the Megatron-LM copy
+pinned by the current Megatron Bridge `main` branch is not supported. Stateless
+grouped-FP8 checkpoint saves additionally require Megatron-LM PR
+[#6509](https://github.com/NVIDIA/Megatron-LM/pull/6509) until that fix merges,
+including saves from the existing MXFP8 recipes. Earlier pretraining
+verification used Megatron-LM `dev` commit `35f36c7c9dba` plus PR
+[#4839](https://github.com/NVIDIA/Megatron-LM/pull/4839).
 
 The NeMo Framework container uses one shared `/opt/venv` for several source
 projects. A plain `uv sync` is exact by default and removes packages that are
@@ -38,9 +44,11 @@ In a standalone Megatron Bridge environment with its own virtual environment,
 exact sync is appropriate: use `uv sync` after switching to dev and restore the
 tracked lock file followed by `uv sync --locked` when switching back to main.
 
-The full-scale `deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_config` performance
-recipe preserves the stack validated by Megatron Bridge PR
-[#4824](https://github.com/NVIDIA-NeMo/Megatron-Bridge/pull/4824):
+The full-scale `deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config`
+library recipe and its corresponding performance recipe require the stack
+validated by Megatron Bridge PR
+[#4824](https://github.com/NVIDIA-NeMo/Megatron-Bridge/pull/4824). The performance
+definition preserves its exact benchmark stack:
 `nvcr.io/nvidia/nemo:26.06.01` with Megatron-LM dev commit
 `9d46c924dce3818f2b5f894f7380712c780d1801` and the capability-check patch
 documented in that PR. The Megatron-LM commit pinned by the current
@@ -58,22 +66,51 @@ performance features, so it is not a supported runtime for that recipe.
 
 - `conversion.sh` imports HF weights into Megatron Bridge and exports Megatron checkpoints back to HF format.
 - `inference.sh` runs text generation against an HF or Megatron checkpoint.
-- `slurm_pretrain.sh` runs the DeepSeek-V4-Flash pretraining recipes.
+- `slurm_pretrain.sh` runs the legacy 32-GPU DeepSeek-V4-Flash pretraining recipes.
 - `slurm_sft.sh` runs DeepSeek-V4-Flash full SFT end to end (import, then fine-tune) on Hopper or Blackwell, with MTP on or off.
 
 Run `bash conversion.sh` after setting `WORKSPACE` and `MODEL_VARIANT`. See each script's header comments for the expected environment variables and `#SBATCH` directives to edit before submitting.
 
 ## Pretraining Recipes
 
-See [`slurm_pretrain.sh`](slurm_pretrain.sh) for the Slurm launcher and [`deepseek_v4.py`](../../../src/megatron/bridge/recipes/deepseek/deepseek_v4.py) for recipe definitions.
+See [`slurm_pretrain.sh`](slurm_pretrain.sh) for the legacy 32-GPU Slurm launcher. Hardware-specific recipe definitions live in the [`b200`](../../../src/megatron/bridge/recipes/deepseek/b200/deepseek_v4.py), [`gb200`](../../../src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py), and [`gb300`](../../../src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py) modules.
 
 Available Blackwell pretraining recipes:
 
 - `deepseek_v4_flash_pretrain_mxfp8_config`: Adam MXFP8
 - `deepseek_v4_flash_pretrain_muon_config`: Muon BF16
-- `deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_config`: 256-GPU GB300
-  performance configuration (requires the PR #4824 container and dev-MCore
-  stack described above)
+- `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
+- `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
+- `deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config`: 256-GPU GB300
+  Adam MXFP8 with PP4/VPP4/EP64/HybridEP (requires the PR #4824 container and
+  dev-MCore stack described above)
+
+The `_library_config` suffix keeps these real-training definitions distinct
+from the corresponding benchmark recipe families. They retain natural routing,
+convergence batch sizes, correctness checks, and checkpoint behavior.
+
+The hardware-count-specific recipes are intentionally not accepted by
+`slurm_pretrain.sh`: that legacy launcher overwrites the recipe's batch and
+TP/PP/EP/CP settings. Launch them through a site multi-node runner that invokes
+`scripts/training/run_recipe.py --recipe <recipe-name> --mode pretrain` with
+64 ranks for B200 Flash, 128 ranks for GB200 Flash, or 256 ranks for GB300 Pro.
+Keep the topology and global batch size owned by the recipe; pass only dataset,
+run-length, logging, and checkpoint overrides.
+
+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
+retain their checkpoint tensor layout. The library recipes retain unlimited
+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.
 
 `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.
 

scripts/training/run_recipe.py

@@ -106,7 +106,6 @@
 from recipe_runner import (  # noqa: E402
     apply_cli_overrides,
     apply_determinism,
-    apply_runtime_environment,
     bootstrap_recipe_environment,
     load_forward_step,
     load_recipe,
@@ -496,17 +495,16 @@ def main(argv: list[str] | None = None) -> None:
             world_size=benchmark_world_size,
         )
     configuration_mode = _train_mode(args.mode)
+    recipe = bootstrap_recipe_environment(
+        recipe,
+        script_path=str(Path(__file__).resolve()),
+        argv=list(argv) if argv is not None else sys.argv[1:],
+    )
 
     if benchmark_metadata is not None:
-        recipe = bootstrap_recipe_environment(
-            recipe,
-            script_path=str(Path(__file__).resolve()),
-            argv=list(argv) if argv is not None else sys.argv[1:],
-        )
         execution_mode = "pretrain"
         step_mode = benchmark_metadata.task
     else:
-        recipe = apply_runtime_environment(recipe)
         execution_mode = configuration_mode
         step_mode = configuration_mode
 

skills/nemo-mbridge-recipe-recommender/references/recipe-index.md

@@ -74,6 +74,9 @@ All recipes live under `src/megatron/bridge/recipes/`. Each function returns a
 | `deepseek_v2_pretrain_config` | Pretrain | 1 | 4 | 32 | 128 |
 | `deepseek_v3_pretrain_config` | Pretrain | 2 | 16 | 64 | 2048 |
 | `deepseek_v3_pretrain_config_32nodes` | Pretrain | 2 | 8 | 32 | 256 |
+| `deepseek_v4_flash_pretrain_64gpu_b200_fp8mx_library_config` | Pretrain | 1 | 8 | 8 | 64 (B200 NVL8; unlimited capacity) |
+| `deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config` | Pretrain | 1 | 1 | 64 | 128 (GB200; unlimited capacity) |
+| `deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config` | Pretrain | 1 | 4 | 64 | 256 (GB300; unlimited capacity) |
 
 ### GLM-4.5
 

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

@@ -18,6 +18,11 @@
 models (V2, V2-Lite, V3, V4).
 """
 
+# DeepSeek V4 B200
+from .b200.deepseek_v4 import (
+    deepseek_v4_flash_pretrain_64gpu_b200_fp8mx_library_config,
+)
+
 # DeepSeek V2/V2-Lite
 from .deepseek_v2 import (
     deepseek_v2_lite_pretrain_config,
@@ -50,10 +55,13 @@
     deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config,
     deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config,
     deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config,
+    deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config,
+    deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config,
 )
 from .gb300.deepseek_v4 import (
     deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config,
     deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config,
+    deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config,
 )
 
 
@@ -77,10 +85,14 @@
     "deepseek_v4_flash_no_mtp_sft_config",
     "deepseek_v4_pro_pretrain_config",
     "deepseek_v4_pro_pretrain_mxfp8_config",
+    "deepseek_v4_flash_pretrain_64gpu_b200_fp8mx_library_config",
     "deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config",
     "deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config",
     "deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config",
+    "deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config",
+    "deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config",
     "deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config",
     "deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config",
+    "deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config",
     "set_deepseek_v4_pipeline_model_parallel_layout",
 ]

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

@@ -0,0 +1,21 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""B200 recipes for DeepSeek V4 Flash."""
+
+from megatron.bridge.recipes.deepseek.b200.deepseek_v4 import (
+    deepseek_v4_flash_pretrain_64gpu_b200_fp8mx_library_config,
+)
+
+
+__all__ = ["deepseek_v4_flash_pretrain_64gpu_b200_fp8mx_library_config"]

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

@@ -0,0 +1,69 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""B200 NVL8 recipes for DeepSeek V4 Flash."""
+
+from megatron.bridge.recipes.deepseek.gb200.deepseek_v4 import (
+    deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config,
+)
+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",
+}
+
+
+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.
+    """
+    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.context_parallel_size = 1
+    cfg.model.expert_model_parallel_size = 8
+    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.moe_token_dispatcher_type = "alltoall"
+    cfg.model.moe_flex_dispatcher_backend = None
+    cfg.model.moe_flex_dispatcher_num_sms = None
+    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_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.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}
+    return cfg

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

@@ -93,6 +93,9 @@ def deepseek_v4_flash_sft_openmath_thinking_packed_config() -> ConfigContainer:
     Pre-pack data with ``prepare_gpt_sft_packed_data.py`` before running SFT.
     When using CP>1, pass ``model.cp_partition_mode=contiguous`` (required for DSv4 CSA
     attention) and ``pad_seq_to_mult=4`` to ensure divisibility by cp_size.
+
+    For GB200-optimized training with HybridEP dispatcher and DSA kernel fusion,
+    use ``deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config`` instead.
     """
     cfg = deepseek_v4_flash_sft_config()
     # DSv4 hybrid attention requires contiguous CP partition when CP > 1;

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

@@ -17,11 +17,15 @@
     deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config,
     deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config,
     deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config,
+    deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config,
+    deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config,
 )
 
 
 __all__ = [
+    "deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config",
     "deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config",
     "deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config",
     "deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config",
+    "deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config",
 ]