From 6ca1400e87a14c48b1d22699181cc42877321b6a Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Thu, 13 Aug 2026 03:13:41 -0700 Subject: [PATCH 1/9] feat(recipes): modernize DeepSeek V4 library configs Signed-off-by: Chen Cui --- examples/models/deepseek_v4/README.md | 48 ++- scripts/training/run_recipe.py | 12 +- .../references/recipe-index.md | 2 + .../deepseek/gb200/deepseek_v4.py | 25 +- .../deepseek/gb300/deepseek_v4.py | 25 +- .../bridge/recipes/deepseek/__init__.py | 4 + .../bridge/recipes/deepseek/gb200/__init__.py | 2 + .../recipes/deepseek/gb200/deepseek_v4.py | 67 +++- .../bridge/recipes/deepseek/gb300/__init__.py | 2 + .../recipes/deepseek/gb300/deepseek_v4.py | 58 +++- .../test_deepseek_v4_hardware_recipes.py | 300 ++++++++++++++++++ .../scripts/training/test_run_recipe.py | 11 +- 12 files changed, 482 insertions(+), 74 deletions(-) create mode 100644 tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py diff --git a/examples/models/deepseek_v4/README.md b/examples/models/deepseek_v4/README.md index bcc1e67525..f7cebe360d 100644 --- a/examples/models/deepseek_v4/README.md +++ b/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,44 @@ 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 [`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_128gpu_gb200_fp8mx_library_config`: 128-GPU GB200 + Adam MXFP8 with PP2/EP64/HybridEP +- `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 --mode pretrain` with +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. `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. diff --git a/scripts/training/run_recipe.py b/scripts/training/run_recipe.py index e3bd313e2a..654c7f0ae5 100755 --- a/scripts/training/run_recipe.py +++ b/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 diff --git a/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md b/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md index 5b58640d63..35a50a9b9e 100644 --- a/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md +++ b/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md @@ -74,6 +74,8 @@ 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_128gpu_gb200_fp8mx_library_config` | Pretrain | 1 | 2 | 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 diff --git a/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v4.py b/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v4.py index e7379dbb99..54af62b017 100644 --- a/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v4.py +++ b/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v4.py @@ -18,7 +18,7 @@ from megatron.bridge.perf_recipes._common import _benchmark_common from megatron.bridge.perf_recipes.environment import COMMON_PERF_ENV_VARS from megatron.bridge.recipes.deepseek.gb200.deepseek_v4 import ( - deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config, + deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config as deepseek_v4_flash_128gpu_library_config, ) from megatron.bridge.training.config import ConfigContainer from megatron.bridge.utils.cuda_graph import set_full_iteration_cuda_graph @@ -26,35 +26,20 @@ def deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_config() -> ConfigContainer: """DeepSeek V4 Flash pretrain: 128× GB200, MXFP8, full-iteration CUDA graph.""" - cfg = deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config() + cfg = deepseek_v4_flash_128gpu_library_config() - cfg.model.tensor_model_parallel_size = 1 cfg.model.pipeline_model_parallel_size = 1 - # The Sheet lists VPP=4, but the measured PP=1 W&B run resolves VPP to None. - cfg.model.virtual_pipeline_model_parallel_size = None - cfg.model.context_parallel_size = 1 - cfg.model.expert_model_parallel_size = 64 - cfg.model.expert_tensor_parallel_size = 1 - cfg.model.sequence_parallel = False cfg.model.pipeline_model_parallel_layout = None cfg.train.global_batch_size = 2048 - cfg.train.micro_batch_size = 1 + cfg.train.manual_gc_interval = 10 cfg.model.attention_backend = "auto" - cfg.model.moe_flex_dispatcher_backend = "hybridep" - cfg.model.moe_token_dispatcher_type = "flex" - cfg.model.moe_shared_expert_overlap = False - cfg.model.moe_hybridep_num_sms = 32 - cfg.model.moe_hybridep_num_sms_preprocessing = 108 cfg.model.moe_router_fusion = True cfg.model.moe_router_force_load_balancing = True cfg.model.moe_router_load_balancing_type = "seq_aux_loss" cfg.model.moe_aux_loss_coeff = 1.0e-4 - - cfg.model.recompute_granularity = "selective" cfg.model.recompute_modules = ["mla_up_proj"] - cfg.model.recompute_method = None - cfg.model.recompute_num_layers = None + cfg.model.fine_grained_activation_offloading = False cfg.model.offload_modules = [] cfg.model.fine_grained_offloading_max_inflight_offloads = None @@ -99,8 +84,6 @@ def deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_config() -> ConfigContainer: cfg.ddp.overlap_param_gather = True cfg.ddp.overlap_grad_reduce = True cfg.ddp.grad_reduce_in_fp32 = False - cfg.ddp.average_in_collective = False - cfg.comm_overlap.overlap_grad_reduce = True cfg.comm_overlap.overlap_moe_expert_parallel_comm = False cfg.comm_overlap.delay_wgrad_compute = False diff --git a/src/megatron/bridge/perf_recipes/deepseek/gb300/deepseek_v4.py b/src/megatron/bridge/perf_recipes/deepseek/gb300/deepseek_v4.py index 6c5d18c831..8998333555 100644 --- a/src/megatron/bridge/perf_recipes/deepseek/gb300/deepseek_v4.py +++ b/src/megatron/bridge/perf_recipes/deepseek/gb300/deepseek_v4.py @@ -21,7 +21,7 @@ ) from megatron.bridge.perf_recipes.environment import COMMON_PERF_ENV_VARS from megatron.bridge.recipes.deepseek.gb300.deepseek_v4 import ( - deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config, + deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config as deepseek_v4_pro_256gpu_library_config, ) from megatron.bridge.training.config import ConfigContainer @@ -52,24 +52,10 @@ def deepseek_v4_flash_pretrain_128gpu_gb300_fp8mx_config() -> ConfigContainer: def deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_config() -> ConfigContainer: """DeepSeek V4 Pro pretrain: 256× GB300, MXFP8, dev Megatron-Core required.""" - cfg = deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config() + cfg = deepseek_v4_pro_256gpu_library_config() - 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.expert_tensor_parallel_size = 1 - cfg.model.sequence_parallel = False cfg.train.global_batch_size = 4096 - cfg.train.micro_batch_size = 1 - - cfg.model.moe_flex_dispatcher_backend = "hybridep" - cfg.model.moe_token_dispatcher_type = "flex" - cfg.model.moe_shared_expert_overlap = False - cfg.model.pipeline_model_parallel_layout = "Et*4|(tttt|)*14tmL" - cfg.model.recompute_granularity = "selective" - cfg.model.recompute_modules = ["mla_up_proj", "mhc"] + cfg.train.manual_gc_interval = 100 _benchmark_common(cfg, cross_entropy_impl="native") @@ -104,11 +90,6 @@ def deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_config() -> ConfigContainer: cfg.dist.enable_megatron_core_experimental = True cfg.ddp.grad_reduce_in_fp32 = False - 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.comm_overlap.overlap_grad_reduce = True - cfg.env_vars = { **COMMON_PERF_ENV_VARS, "CUDA_DEVICE_MAX_CONNECTIONS": 32, diff --git a/src/megatron/bridge/recipes/deepseek/__init__.py b/src/megatron/bridge/recipes/deepseek/__init__.py index 1cf2ac77dd..b9dcb8c0bf 100644 --- a/src/megatron/bridge/recipes/deepseek/__init__.py +++ b/src/megatron/bridge/recipes/deepseek/__init__.py @@ -50,10 +50,12 @@ 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, ) 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, ) @@ -80,7 +82,9 @@ "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_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", ] diff --git a/src/megatron/bridge/recipes/deepseek/gb200/__init__.py b/src/megatron/bridge/recipes/deepseek/gb200/__init__.py index 696f6cae8b..7e93b92af4 100644 --- a/src/megatron/bridge/recipes/deepseek/gb200/__init__.py +++ b/src/megatron/bridge/recipes/deepseek/gb200/__init__.py @@ -17,10 +17,12 @@ 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, ) __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", diff --git a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py index aee81537eb..6341cb6a3e 100644 --- a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py @@ -71,7 +71,7 @@ def deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config() -> ConfigContainer: set_deepseek_v4_pipeline_model_parallel_layout(cfg.model) cfg.model.transformer_impl = "transformer_engine" - cfg.model.attention_backend = None + cfg.model.attention_backend = "auto" cfg.model.apply_dsa_kernel_fusion = True cfg.model.apply_rope_fusion = True cfg.model.use_fused_mhc = use_fused_mhc @@ -83,17 +83,20 @@ def deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config() -> ConfigContainer: cfg.model.moe_hybridep_num_sms = 16 cfg.model.moe_grouped_gemm = True cfg.model.moe_permute_fusion = True + cfg.model.moe_router_fusion = True cfg.model.moe_aux_loss_coeff = 0.0 cfg.model.moe_router_force_load_balancing = False + cfg.model.moe_pad_experts_for_cuda_graph_inference = True cfg.model.cross_entropy_loss_fusion = True - cfg.model.cross_entropy_fusion_impl = "te" + cfg.model.cross_entropy_fusion_impl = "native" cfg.model.recompute_granularity = "selective" cfg.model.recompute_modules = ["moe_act", "mhc", "mla_up_proj"] cfg.model.recompute_method = None cfg.model.recompute_num_layers = None - cfg.model.fine_grained_activation_offloading = False - cfg.model.offload_modules = None + 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.cuda_graph_impl = "none" cfg.model.cuda_graph_scope = "full" cfg.model.cuda_graph_warmup_steps = 3 @@ -127,6 +130,7 @@ def deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config() -> ConfigContainer: cfg.dist.enable_megatron_core_experimental = True cfg.comm_overlap = CommOverlapConfig(tp_comm_overlap=False) + cfg.comm_overlap.overlap_grad_reduce = True cfg.comm_overlap.delay_wgrad_compute = False cfg.comm_overlap.overlap_moe_expert_parallel_comm = False @@ -134,6 +138,7 @@ def deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config() -> ConfigContainer: cfg.ddp.use_megatron_fsdp = False cfg.env_vars = { **COMMON_RECIPE_ENV_VARS, + "NVTE_CPU_OFFLOAD_V1": 1, } return cfg @@ -145,10 +150,9 @@ def deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config() -> ConfigContainer: """ cfg = deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config() - cfg.model.apply_dsa_kernel_fusion = False cfg.model.dsa_indexer_loss_coeff = 0.0 cfg.model.dsa_indexer_use_sparse_loss = False - cfg.model.recompute_modules = ["mla_up_proj"] + cfg.model.recompute_modules = ["moe_act", "mhc", "mla_up_proj"] opt_cfg, scheduler_cfg = distributed_fused_adam_with_cosine_annealing( lr_warmup_iters=2000, @@ -188,6 +192,57 @@ def deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config() -> ConfigContainer: return cfg +def deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config() -> ConfigContainer: + """Return the real-training DeepSeek V4 Flash config for 128 GB200 GPUs. + + This variant uses one PP stage per 64-rank NVLink domain while preserving + the 64-GPU library recipe's optimizer, routing, loss, precision, validation, + and checkpoint contracts. + """ + cfg = deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config() + + cfg.model.tensor_model_parallel_size = 1 + cfg.model.pipeline_model_parallel_size = 2 + cfg.model.virtual_pipeline_model_parallel_size = None + cfg.model.context_parallel_size = 1 + cfg.model.expert_model_parallel_size = 64 + cfg.model.expert_tensor_parallel_size = 1 + cfg.model.sequence_parallel = False + set_deepseek_v4_pipeline_model_parallel_layout(cfg.model) + cfg.train.global_batch_size = 256 + cfg.train.micro_batch_size = 1 + + cfg.model.moe_token_dispatcher_type = "flex" + cfg.model.moe_flex_dispatcher_backend = "hybridep" + cfg.model.moe_shared_expert_overlap = False + cfg.model.moe_hybridep_num_sms = 32 + cfg.model.moe_hybridep_num_sms_preprocessing = 108 + cfg.model.moe_mlp_glu_interleave_size = 32 + cfg.model.use_transformer_engine_op_fuser = True + 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.ddp.average_in_collective = False + + cfg.env_vars = { + **cfg.env_vars, + "CUDA_DEVICE_MAX_CONNECTIONS": 32, + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 64, + "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, + "NVLINK_DOMAIN_SIZE": 72, + "USE_MNNVL": 1, + "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, + "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + "NVTE_NORM_BWD_USE_CUDNN": 1, + "NVTE_NORM_FWD_USE_CUDNN": 1, + "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, + "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1, + } + cfg.env_vars.pop("NVTE_CPU_OFFLOAD_V1", None) + return cfg + + def deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config() -> ConfigContainer: """DeepSeek-V4-Flash BF16 Muon pre-training on 64-GPU GB200 (PP=8, EP=8).""" cfg = deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config() diff --git a/src/megatron/bridge/recipes/deepseek/gb300/__init__.py b/src/megatron/bridge/recipes/deepseek/gb300/__init__.py index e4b9a7e485..db38cda838 100644 --- a/src/megatron/bridge/recipes/deepseek/gb300/__init__.py +++ b/src/megatron/bridge/recipes/deepseek/gb300/__init__.py @@ -17,10 +17,12 @@ from megatron.bridge.recipes.deepseek.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, ) __all__ = [ + "deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config", "deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config", "deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config", ] diff --git a/src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py index db0dcba42d..ea263b2370 100644 --- a/src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py @@ -63,7 +63,7 @@ def deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config() -> ConfigContainer: cfg.model.transformer_impl = "transformer_engine" cfg.model.attention_backend = None - cfg.model.apply_dsa_kernel_fusion = False + cfg.model.apply_dsa_kernel_fusion = True cfg.model.apply_rope_fusion = True cfg.model.use_fused_mhc = deepseek_v4_supports_blackwell_fused_kernels() cfg.model.dsa_indexer_loss_coeff = 0.0 @@ -72,15 +72,17 @@ def deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config() -> ConfigContainer: cfg.model.moe_token_dispatcher_type = "alltoall" cfg.model.moe_aux_loss_coeff = 0.0 cfg.model.moe_router_force_load_balancing = False + cfg.model.moe_pad_experts_for_cuda_graph_inference = True cfg.model.cross_entropy_loss_fusion = True - cfg.model.cross_entropy_fusion_impl = "te" + cfg.model.cross_entropy_fusion_impl = "native" cfg.model.recompute_granularity = "selective" cfg.model.recompute_modules = ["moe_act", "mhc"] cfg.model.recompute_method = None cfg.model.recompute_num_layers = None - cfg.model.fine_grained_activation_offloading = False - cfg.model.offload_modules = None + 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.cuda_graph_impl = "none" cfg.model.cuda_graph_scope = "full" cfg.model.cuda_graph_warmup_steps = 3 @@ -114,6 +116,7 @@ def deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config() -> ConfigContainer: cfg.dist.enable_megatron_core_experimental = True cfg.comm_overlap = CommOverlapConfig(tp_comm_overlap=False) + cfg.comm_overlap.overlap_grad_reduce = True cfg.comm_overlap.delay_wgrad_compute = False cfg.comm_overlap.overlap_moe_expert_parallel_comm = False @@ -121,6 +124,7 @@ def deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config() -> ConfigContainer: cfg.ddp.use_megatron_fsdp = False cfg.env_vars = { **COMMON_RECIPE_ENV_VARS, + "NVTE_CPU_OFFLOAD_V1": 1, } return cfg @@ -141,7 +145,6 @@ def deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config() -> ConfigContainer: cfg.train.train_iters = 1_000_000 cfg.train.global_batch_size = 128 cfg.train.micro_batch_size = 1 - cfg.model.apply_dsa_kernel_fusion = False cfg.model.apply_rope_fusion = True cfg.model.use_fused_mhc = deepseek_v4_supports_blackwell_fused_kernels() cfg.model.dsa_indexer_loss_coeff = 0.0 @@ -186,7 +189,50 @@ def deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config() -> ConfigContainer: cfg.model.moe_router_padding_for_fp8 = True cfg.model.mtp_eval_in_bf16 = True cfg.model.quant_recipe = _deepseek_v4_mxfp8_quant_recipe() + return cfg + + +def deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config() -> ConfigContainer: + """Return the real-training DeepSeek V4 Pro config for 256 GB300 GPUs. + + This variant adopts the measured PP4/VPP4/EP64 execution topology while + preserving the 32-GPU library recipe's optimizer, routing, loss, precision, + validation, and checkpoint contracts. + """ + cfg = deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config() + + 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.expert_tensor_parallel_size = 1 + cfg.model.sequence_parallel = False + cfg.model.pipeline_model_parallel_layout = "Et*4|(tttt|)*14tmL" + cfg.train.global_batch_size = 128 + cfg.train.micro_batch_size = 1 + + cfg.model.moe_token_dispatcher_type = "flex" + cfg.model.moe_flex_dispatcher_backend = "hybridep" + cfg.model.moe_shared_expert_overlap = False + cfg.model.moe_hybridep_num_sms = 32 + cfg.model.moe_mlp_glu_interleave_size = 32 + cfg.model.use_transformer_engine_op_fuser = True + cfg.model.recompute_granularity = "selective" + cfg.model.recompute_modules = ["mla_up_proj", "mhc"] cfg.env_vars = { - **COMMON_RECIPE_ENV_VARS, + **cfg.env_vars, + "CUDA_DEVICE_MAX_CONNECTIONS": 32, + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 64, + "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, + "NVLINK_DOMAIN_SIZE": 72, + "USE_MNNVL": 1, + "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, + "NVTE_CPU_OFFLOAD_V1": 1, + "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + "NVTE_NORM_BWD_USE_CUDNN": 1, + "NVTE_NORM_FWD_USE_CUDNN": 1, + "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, + "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1, } return cfg diff --git a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py new file mode 100644 index 0000000000..ba4f6693f0 --- /dev/null +++ b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py @@ -0,0 +1,300 @@ +# 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. + +"""Execution-parity and convergence guardrails for DeepSeek V4 hardware recipes.""" + +from collections.abc import Callable + +import pytest +import torch + +import megatron.bridge.recipes as recipes +from megatron.bridge.perf_recipes.deepseek.gb200.deepseek_v4 import ( + deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_config as flash_perf_config, +) +from megatron.bridge.perf_recipes.deepseek.gb300.deepseek_v4 import ( + deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_config as pro_perf_config, +) +from megatron.bridge.recipes.deepseek.gb200.deepseek_v4 import ( + deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config as flash_bf16_base_config, +) +from megatron.bridge.recipes.deepseek.gb200.deepseek_v4 import ( + deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config as flash_fp8_config, +) +from megatron.bridge.recipes.deepseek.gb200.deepseek_v4 import ( + deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config as flash_library_config, +) +from megatron.bridge.recipes.deepseek.gb300.deepseek_v4 import ( + deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config as pro_bf16_base_config, +) +from megatron.bridge.recipes.deepseek.gb300.deepseek_v4 import ( + deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config as pro_fp8_config, +) +from megatron.bridge.recipes.deepseek.gb300.deepseek_v4 import ( + deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config as pro_library_config, +) +from megatron.bridge.training.config import ConfigContainer +from tests.unit_tests.recipes.recipe_test_utils import patch_recipe_construction_dependencies + + +pytestmark = pytest.mark.unit + +_SHARED_MODEL_EXECUTION_FIELDS = ( + "tensor_model_parallel_size", + "virtual_pipeline_model_parallel_size", + "context_parallel_size", + "expert_model_parallel_size", + "expert_tensor_parallel_size", + "sequence_parallel", + "moe_token_dispatcher_type", + "moe_flex_dispatcher_backend", + "moe_shared_expert_overlap", + "moe_hybridep_num_sms", + "recompute_granularity", + "cross_entropy_fusion_impl", + "moe_pad_experts_for_cuda_graph_inference", + "moe_mlp_glu_interleave_size", + "use_transformer_engine_op_fuser", +) + +_SHARED_ENV_VARS = ( + "CUDA_DEVICE_MAX_CONNECTIONS", + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN", + "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API", + "NVLINK_DOMAIN_SIZE", + "USE_MNNVL", + "NVTE_BWD_LAYERNORM_SM_MARGIN", + "NVTE_FWD_LAYERNORM_SM_MARGIN", + "NVTE_NORM_BWD_USE_CUDNN", + "NVTE_NORM_FWD_USE_CUDNN", + "NVTE_ALLOW_NONDETERMINISTIC_ALGO", +) + + +@pytest.fixture(autouse=True) +def _keep_recipe_construction_offline(monkeypatch: pytest.MonkeyPatch) -> None: + patch_recipe_construction_dependencies(monkeypatch) + + +@pytest.mark.parametrize( + ("library_factory", "perf_factory"), + [ + (flash_library_config, flash_perf_config), + (pro_library_config, pro_perf_config), + ], + ids=["flash-gb200", "pro-gb300"], +) +def test_high_scale_library_recipes_share_approved_execution_fields( + library_factory: Callable[[], ConfigContainer], + perf_factory: Callable[[], ConfigContainer], +) -> None: + library_cfg = library_factory() + perf_cfg = perf_factory() + + for field_name in _SHARED_MODEL_EXECUTION_FIELDS: + library_value = getattr(library_cfg.model, field_name, None) + perf_value = getattr(perf_cfg.model, field_name, None) + assert library_value == perf_value, field_name + + assert library_cfg.train.micro_batch_size == perf_cfg.train.micro_batch_size == 1 + assert library_cfg.comm_overlap.overlap_grad_reduce is perf_cfg.comm_overlap.overlap_grad_reduce is True + assert library_cfg.ddp.average_in_collective == perf_cfg.ddp.average_in_collective + for variable_name in _SHARED_ENV_VARS: + assert library_cfg.env_vars[variable_name] == perf_cfg.env_vars[variable_name], variable_name + + +@pytest.mark.parametrize( + "base_factory", + [flash_bf16_base_config, pro_bf16_base_config], + ids=["flash-bf16", "pro-bf16"], +) +def test_base_recipes_enable_precision_independent_fusions( + base_factory: Callable[[], ConfigContainer], +) -> None: + cfg = base_factory() + + assert cfg.model.apply_dsa_kernel_fusion is True + assert cfg.model.moe_pad_experts_for_cuda_graph_inference is True + assert getattr(cfg.model, "moe_mlp_glu_interleave_size", None) is None + assert cfg.model.use_transformer_engine_op_fuser is False + assert cfg.model.cross_entropy_fusion_impl == "native" + assert cfg.comm_overlap.overlap_grad_reduce is True + assert "NVTE_CUTEDSL_FUSED_GROUPED_MLP" not in cfg.env_vars + + assert getattr(cfg.model, "moe_expert_rank_capacity_factor", None) is None + assert getattr(cfg.model, "moe_paged_stash", False) is False + assert cfg.model.cuda_graph_impl == "none" + assert cfg.rerun_state_machine.check_for_nan_in_loss is True + assert cfg.ddp.check_for_nan_in_grad is True + assert not getattr(cfg.model, "fp8", False) + assert getattr(cfg.model, "quant_recipe", None) is None + assert cfg.optimizer.optimizer_offload_fraction == 0.0 + + +def test_flash_base_recipe_ports_flash_fusions() -> None: + cfg = flash_bf16_base_config() + + assert cfg.model.attention_backend == "auto" + assert cfg.model.moe_router_fusion is True + assert cfg.train.manual_gc_interval == 5 + 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 cfg.env_vars["NVTE_CPU_OFFLOAD_V1"] == 1 + + +def test_pro_base_recipe_ports_activation_offload() -> None: + cfg = pro_bf16_base_config() + + assert cfg.model.attention_backend is None + assert cfg.train.manual_gc_interval == 5 + 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 cfg.env_vars["NVTE_CPU_OFFLOAD_V1"] == 1 + + +@pytest.mark.parametrize( + "fp8_factory", + [flash_fp8_config, pro_fp8_config], + ids=["flash-mxfp8", "pro-mxfp8"], +) +def test_mxfp8_recipes_keep_training_precision_contract( + fp8_factory: Callable[[], ConfigContainer], +) -> None: + cfg = fp8_factory() + + assert cfg.model.quant_recipe is not None + assert cfg.model.moe_router_padding_for_fp8 is True + assert cfg.mixed_precision.fp8_param_gather is False + assert cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag is False + assert cfg.mixed_precision.grad_reduce_in_fp32 is True + assert cfg.ddp.grad_reduce_in_fp32 is True + assert cfg.optimizer.main_grads_dtype == torch.float32 + + +def test_flash_mxfp8_recipe_uses_activation_offload_to_fit() -> None: + cfg = flash_fp8_config() + + assert cfg.model.recompute_modules == ["moe_act", "mhc", "mla_up_proj"] + 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 cfg.env_vars["NVTE_CPU_OFFLOAD_V1"] == 1 + + +def test_flash_high_scale_recipe_preserves_real_training_contract() -> None: + cfg = flash_library_config() + + assert cfg.train.train_iters == 1_000_000 + assert cfg.train.global_batch_size == 256 + assert cfg.validation.eval_interval == 2000 + assert cfg.validation.eval_iters == 32 + assert cfg.model.moe_router_force_load_balancing is False + assert cfg.model.dsa_indexer_loss_coeff == 0.0 + assert cfg.model.dsa_indexer_use_sparse_loss is False + assert cfg.model.apply_dsa_kernel_fusion is True + assert cfg.model.quant_recipe is not None + assert cfg.model.moe_router_padding_for_fp8 is True + assert cfg.mixed_precision.fp8_param_gather is False + assert cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag is False + assert cfg.mixed_precision.grad_reduce_in_fp32 is True + 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.model.cuda_graph_impl == "none" + assert cfg.model.use_transformer_engine_op_fuser is True + assert cfg.model.moe_mlp_glu_interleave_size == 32 + assert cfg.model.pipeline_model_parallel_size == 2 + assert len(cfg.model.pipeline_model_parallel_layout) == 2 + assert cfg.model.pipeline_model_parallel_layout[0][0] == "embedding" + assert cfg.model.pipeline_model_parallel_layout[-1][-2:] == ["mtp", "loss"] + assert sum(stage.count("decoder") for stage in cfg.model.pipeline_model_parallel_layout) == cfg.model.num_layers + assert cfg.model.recompute_modules == ["mhc", "mla_up_proj"] + 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 getattr(cfg.model, "moe_expert_rank_capacity_factor", None) is None + assert getattr(cfg.model, "moe_paged_stash", False) is False + assert cfg.env_vars["PYTORCH_CUDA_ALLOC_CONF"] == "expandable_segments:True" + assert cfg.env_vars["TORCH_NCCL_AVOID_RECORD_STREAMS"] == 1 + assert cfg.env_vars["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] == 1 + assert "NVTE_CPU_OFFLOAD_V1" not in cfg.env_vars + + perf_cfg = flash_perf_config() + assert perf_cfg.model.pipeline_model_parallel_size == 1 + assert perf_cfg.model.pipeline_model_parallel_layout is None + assert perf_cfg.train.global_batch_size == 2048 + assert perf_cfg.model.moe_router_force_load_balancing is True + assert perf_cfg.model.dsa_indexer_loss_coeff == 0.01 + assert perf_cfg.model.quant_recipe is None + assert perf_cfg.model.recompute_modules == ["mla_up_proj"] + assert perf_cfg.model.fine_grained_activation_offloading is False + assert perf_cfg.model.offload_modules == [] + assert perf_cfg.mixed_precision.fp8_param_gather is True + assert perf_cfg.ddp.grad_reduce_in_fp32 is False + assert perf_cfg.ddp.check_for_nan_in_grad is False + + +def test_pro_high_scale_recipe_preserves_real_training_contract() -> None: + cfg = pro_library_config() + + assert cfg.model.pipeline_model_parallel_size == 4 + assert cfg.model.virtual_pipeline_model_parallel_size == 4 + assert cfg.model.pipeline_model_parallel_layout == "Et*4|(tttt|)*14tmL" + assert cfg.train.train_iters == 1_000_000 + assert cfg.train.global_batch_size == 128 + assert cfg.validation.eval_interval == 2000 + assert cfg.validation.eval_iters == 32 + assert cfg.model.moe_router_force_load_balancing is False + assert cfg.model.dsa_indexer_loss_coeff == 0.0 + assert cfg.model.dsa_indexer_use_sparse_loss is False + assert cfg.model.apply_dsa_kernel_fusion is True + assert cfg.model.quant_recipe is not None + assert cfg.model.moe_router_padding_for_fp8 is True + assert cfg.mixed_precision.fp8_param_gather is False + assert cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag is False + assert cfg.mixed_precision.grad_reduce_in_fp32 is True + 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.model.cuda_graph_impl == "none" + assert cfg.model.use_transformer_engine_op_fuser is True + assert cfg.model.moe_mlp_glu_interleave_size == 32 + assert getattr(cfg.model, "moe_expert_rank_capacity_factor", None) is None + assert getattr(cfg.model, "moe_paged_stash", False) is False + assert cfg.model.fine_grained_activation_offloading is True + assert cfg.model.offload_modules == ["core_attn", "attn_proj"] + assert cfg.env_vars["NVTE_CPU_OFFLOAD_V1"] == 1 + assert cfg.env_vars["PYTORCH_CUDA_ALLOC_CONF"] == "expandable_segments:True" + assert cfg.env_vars["TORCH_NCCL_AVOID_RECORD_STREAMS"] == 1 + assert cfg.env_vars["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] == 1 + + perf_cfg = pro_perf_config() + assert perf_cfg.model.pipeline_model_parallel_size == 4 + assert perf_cfg.model.virtual_pipeline_model_parallel_size == 4 + assert perf_cfg.model.pipeline_model_parallel_layout == "Et*4|(tttt|)*14tmL" + assert perf_cfg.train.global_batch_size == 4096 + assert perf_cfg.model.moe_router_force_load_balancing is True + assert perf_cfg.model.dsa_indexer_loss_coeff == 0.01 + assert perf_cfg.model.quant_recipe is None + assert perf_cfg.mixed_precision.fp8_param_gather is True + assert perf_cfg.optimizer.main_grads_dtype == torch.bfloat16 + assert perf_cfg.ddp.grad_reduce_in_fp32 is False + assert perf_cfg.ddp.check_for_nan_in_grad is False + + +def test_high_scale_deepseek_v4_recipes_are_exported() -> None: + assert recipes.deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config is flash_library_config + assert recipes.deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config is pro_library_config diff --git a/tests/unit_tests/scripts/training/test_run_recipe.py b/tests/unit_tests/scripts/training/test_run_recipe.py index dc17ee22f2..1f8b540e2e 100644 --- a/tests/unit_tests/scripts/training/test_run_recipe.py +++ b/tests/unit_tests/scripts/training/test_run_recipe.py @@ -475,7 +475,7 @@ def test_benchmark_finetuning_recipes_use_unified_runner(monkeypatch, mode, task assert handles.recipe_runner.run_config.call_args.kwargs["mode"] == "pretrain" -def test_library_only_canonical_name_does_not_enable_benchmark_runtime(): +def test_library_only_canonical_name_bootstraps_recipe_environment(): module, handles = _load_module() handles.recipe_runner.load_recipe.return_value = SimpleNamespace() @@ -488,7 +488,10 @@ def test_library_only_canonical_name_does_not_enable_benchmark_runtime(): ] ) - handles.recipe_runner.bootstrap_recipe_environment.assert_not_called() + handles.recipe_runner.bootstrap_recipe_environment.assert_called_once() + bootstrap_call = handles.recipe_runner.bootstrap_recipe_environment.call_args + assert bootstrap_call.args == (handles.recipe_runner.load_recipe.return_value,) + assert bootstrap_call.kwargs["script_path"].endswith("scripts/training/run_recipe.py") handles.recipe_runner.load_forward_step.assert_called_once_with("llm_step", mode="pretrain") @@ -1444,4 +1447,6 @@ def test_config_container_overrides_are_forwarded_directly(): config, ["train.train_iters=3", "train.global_batch_size=8", "train.micro_batch_size=1"], ) - handles.recipe_runner.apply_runtime_environment.assert_called_once_with(config) + handles.recipe_runner.bootstrap_recipe_environment.assert_called_once() + bootstrap_call = handles.recipe_runner.bootstrap_recipe_environment.call_args + assert bootstrap_call.args == (config,) From 76a84bbf0775fad0f8975599465701d5975795cb Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Fri, 14 Aug 2026 14:00:56 -0700 Subject: [PATCH 2/9] feat(recipes): record DeepSeek V4 PP1 memory settings Signed-off-by: Chen Cui --- examples/models/deepseek_v4/README.md | 6 +++-- .../references/recipe-index.md | 2 +- .../recipes/deepseek/gb200/deepseek_v4.py | 21 ++++++++++-------- .../test_deepseek_v4_hardware_recipes.py | 22 +++++++++---------- 4 files changed, 27 insertions(+), 24 deletions(-) diff --git a/examples/models/deepseek_v4/README.md b/examples/models/deepseek_v4/README.md index f7cebe360d..236d5f73a4 100644 --- a/examples/models/deepseek_v4/README.md +++ b/examples/models/deepseek_v4/README.md @@ -80,7 +80,8 @@ Available Blackwell pretraining recipes: - `deepseek_v4_flash_pretrain_mxfp8_config`: Adam MXFP8 - `deepseek_v4_flash_pretrain_muon_config`: Muon BF16 - `deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config`: 128-GPU GB200 - Adam MXFP8 with PP2/EP64/HybridEP + Adam MXFP8 with PP1/EP64/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) @@ -103,7 +104,8 @@ 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. +reuse, and reduced-precision gradient reduction remain disabled except for the +validated FP8 parameter gather/buffer reuse in the high-scale GB200 recipe. `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. diff --git a/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md b/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md index 35a50a9b9e..a6f4d8dbca 100644 --- a/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md +++ b/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md @@ -74,7 +74,7 @@ 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_128gpu_gb200_fp8mx_library_config` | Pretrain | 1 | 2 | 64 | 128 (GB200; 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 diff --git a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py index 6341cb6a3e..f672c5253b 100644 --- a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py @@ -195,14 +195,16 @@ def deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config() -> ConfigContainer: def deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config() -> ConfigContainer: """Return the real-training DeepSeek V4 Flash config for 128 GB200 GPUs. - This variant uses one PP stage per 64-rank NVLink domain while preserving - the 64-GPU library recipe's optimizer, routing, loss, precision, validation, - and checkpoint contracts. + This variant uses PP1 with two 64-rank expert/data-parallel replicas. A + matched PP1 screen completed eight finite steps and reduced peak allocated + memory by 25.51 GiB with attention activation offload plus MXFP8 parameter + gather/storage. Expert capacity, paged stash, and CUDA graphs remain + disabled to preserve natural-routing training semantics. """ cfg = deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config() cfg.model.tensor_model_parallel_size = 1 - cfg.model.pipeline_model_parallel_size = 2 + cfg.model.pipeline_model_parallel_size = 1 cfg.model.virtual_pipeline_model_parallel_size = None cfg.model.context_parallel_size = 1 cfg.model.expert_model_parallel_size = 64 @@ -219,10 +221,12 @@ def deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config() -> ConfigCont cfg.model.moe_hybridep_num_sms_preprocessing = 108 cfg.model.moe_mlp_glu_interleave_size = 32 cfg.model.use_transformer_engine_op_fuser = True - 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.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.mixed_precision.fp8_param_gather = True + cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag = True cfg.ddp.average_in_collective = False cfg.env_vars = { @@ -239,7 +243,6 @@ def deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config() -> ConfigCont "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1, } - cfg.env_vars.pop("NVTE_CPU_OFFLOAD_V1", None) return cfg diff --git a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py index ba4f6693f0..863a6dc346 100644 --- a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py +++ b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py @@ -207,8 +207,8 @@ def test_flash_high_scale_recipe_preserves_real_training_contract() -> None: assert cfg.model.apply_dsa_kernel_fusion is True assert cfg.model.quant_recipe is not None assert cfg.model.moe_router_padding_for_fp8 is True - assert cfg.mixed_precision.fp8_param_gather is False - assert cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag is False + 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 True assert cfg.ddp.grad_reduce_in_fp32 is True assert cfg.optimizer.main_grads_dtype == torch.float32 @@ -216,21 +216,18 @@ def test_flash_high_scale_recipe_preserves_real_training_contract() -> None: 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 - assert cfg.model.pipeline_model_parallel_size == 2 - assert len(cfg.model.pipeline_model_parallel_layout) == 2 - assert cfg.model.pipeline_model_parallel_layout[0][0] == "embedding" - assert cfg.model.pipeline_model_parallel_layout[-1][-2:] == ["mtp", "loss"] - assert sum(stage.count("decoder") for stage in cfg.model.pipeline_model_parallel_layout) == cfg.model.num_layers - assert cfg.model.recompute_modules == ["mhc", "mla_up_proj"] - 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.model.pipeline_model_parallel_size == 1 + assert cfg.model.pipeline_model_parallel_layout is None + 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_expert_rank_capacity_factor", None) is None assert getattr(cfg.model, "moe_paged_stash", False) is False assert cfg.env_vars["PYTORCH_CUDA_ALLOC_CONF"] == "expandable_segments:True" assert cfg.env_vars["TORCH_NCCL_AVOID_RECORD_STREAMS"] == 1 assert cfg.env_vars["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] == 1 - assert "NVTE_CPU_OFFLOAD_V1" not in cfg.env_vars + assert cfg.env_vars["NVTE_CPU_OFFLOAD_V1"] == 1 perf_cfg = flash_perf_config() assert perf_cfg.model.pipeline_model_parallel_size == 1 @@ -243,6 +240,7 @@ def test_flash_high_scale_recipe_preserves_real_training_contract() -> None: assert perf_cfg.model.fine_grained_activation_offloading is False assert perf_cfg.model.offload_modules == [] assert perf_cfg.mixed_precision.fp8_param_gather is True + assert perf_cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag is True assert perf_cfg.ddp.grad_reduce_in_fp32 is False assert perf_cfg.ddp.check_for_nan_in_grad is False From 9f1fbe9b7a71cdd1ecf00c19726a4c4686e62f6e Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Fri, 14 Aug 2026 14:20:02 -0700 Subject: [PATCH 3/9] fix(recipes): keep DeepSeek V4 perf configs unchanged Signed-off-by: Chen Cui --- .../deepseek/gb200/deepseek_v4.py | 25 ++++- .../deepseek/gb300/deepseek_v4.py | 25 ++++- .../test_deepseek_v4_hardware_recipes.py | 94 +------------------ 3 files changed, 44 insertions(+), 100 deletions(-) diff --git a/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v4.py b/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v4.py index 54af62b017..e7379dbb99 100644 --- a/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v4.py +++ b/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v4.py @@ -18,7 +18,7 @@ from megatron.bridge.perf_recipes._common import _benchmark_common from megatron.bridge.perf_recipes.environment import COMMON_PERF_ENV_VARS from megatron.bridge.recipes.deepseek.gb200.deepseek_v4 import ( - deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config as deepseek_v4_flash_128gpu_library_config, + deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config, ) from megatron.bridge.training.config import ConfigContainer from megatron.bridge.utils.cuda_graph import set_full_iteration_cuda_graph @@ -26,20 +26,35 @@ def deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_config() -> ConfigContainer: """DeepSeek V4 Flash pretrain: 128× GB200, MXFP8, full-iteration CUDA graph.""" - cfg = deepseek_v4_flash_128gpu_library_config() + cfg = deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config() + cfg.model.tensor_model_parallel_size = 1 cfg.model.pipeline_model_parallel_size = 1 + # The Sheet lists VPP=4, but the measured PP=1 W&B run resolves VPP to None. + cfg.model.virtual_pipeline_model_parallel_size = None + cfg.model.context_parallel_size = 1 + cfg.model.expert_model_parallel_size = 64 + cfg.model.expert_tensor_parallel_size = 1 + cfg.model.sequence_parallel = False cfg.model.pipeline_model_parallel_layout = None cfg.train.global_batch_size = 2048 - cfg.train.manual_gc_interval = 10 + cfg.train.micro_batch_size = 1 cfg.model.attention_backend = "auto" + cfg.model.moe_flex_dispatcher_backend = "hybridep" + cfg.model.moe_token_dispatcher_type = "flex" + cfg.model.moe_shared_expert_overlap = False + cfg.model.moe_hybridep_num_sms = 32 + cfg.model.moe_hybridep_num_sms_preprocessing = 108 cfg.model.moe_router_fusion = True cfg.model.moe_router_force_load_balancing = True cfg.model.moe_router_load_balancing_type = "seq_aux_loss" cfg.model.moe_aux_loss_coeff = 1.0e-4 - cfg.model.recompute_modules = ["mla_up_proj"] + cfg.model.recompute_granularity = "selective" + cfg.model.recompute_modules = ["mla_up_proj"] + cfg.model.recompute_method = None + cfg.model.recompute_num_layers = None cfg.model.fine_grained_activation_offloading = False cfg.model.offload_modules = [] cfg.model.fine_grained_offloading_max_inflight_offloads = None @@ -84,6 +99,8 @@ def deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_config() -> ConfigContainer: cfg.ddp.overlap_param_gather = True cfg.ddp.overlap_grad_reduce = True cfg.ddp.grad_reduce_in_fp32 = False + cfg.ddp.average_in_collective = False + cfg.comm_overlap.overlap_grad_reduce = True cfg.comm_overlap.overlap_moe_expert_parallel_comm = False cfg.comm_overlap.delay_wgrad_compute = False diff --git a/src/megatron/bridge/perf_recipes/deepseek/gb300/deepseek_v4.py b/src/megatron/bridge/perf_recipes/deepseek/gb300/deepseek_v4.py index 8998333555..6c5d18c831 100644 --- a/src/megatron/bridge/perf_recipes/deepseek/gb300/deepseek_v4.py +++ b/src/megatron/bridge/perf_recipes/deepseek/gb300/deepseek_v4.py @@ -21,7 +21,7 @@ ) from megatron.bridge.perf_recipes.environment import COMMON_PERF_ENV_VARS from megatron.bridge.recipes.deepseek.gb300.deepseek_v4 import ( - deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config as deepseek_v4_pro_256gpu_library_config, + deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config, ) from megatron.bridge.training.config import ConfigContainer @@ -52,10 +52,24 @@ def deepseek_v4_flash_pretrain_128gpu_gb300_fp8mx_config() -> ConfigContainer: def deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_config() -> ConfigContainer: """DeepSeek V4 Pro pretrain: 256× GB300, MXFP8, dev Megatron-Core required.""" - cfg = deepseek_v4_pro_256gpu_library_config() + cfg = deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config() + 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.expert_tensor_parallel_size = 1 + cfg.model.sequence_parallel = False cfg.train.global_batch_size = 4096 - cfg.train.manual_gc_interval = 100 + cfg.train.micro_batch_size = 1 + + cfg.model.moe_flex_dispatcher_backend = "hybridep" + cfg.model.moe_token_dispatcher_type = "flex" + cfg.model.moe_shared_expert_overlap = False + cfg.model.pipeline_model_parallel_layout = "Et*4|(tttt|)*14tmL" + cfg.model.recompute_granularity = "selective" + cfg.model.recompute_modules = ["mla_up_proj", "mhc"] _benchmark_common(cfg, cross_entropy_impl="native") @@ -90,6 +104,11 @@ def deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_config() -> ConfigContainer: cfg.dist.enable_megatron_core_experimental = True cfg.ddp.grad_reduce_in_fp32 = False + 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.comm_overlap.overlap_grad_reduce = True + cfg.env_vars = { **COMMON_PERF_ENV_VARS, "CUDA_DEVICE_MAX_CONNECTIONS": 32, diff --git a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py index 863a6dc346..a3aa82022b 100644 --- a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py +++ b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Execution-parity and convergence guardrails for DeepSeek V4 hardware recipes.""" +"""Execution and convergence guardrails for DeepSeek V4 hardware recipes.""" from collections.abc import Callable @@ -20,12 +20,6 @@ import torch import megatron.bridge.recipes as recipes -from megatron.bridge.perf_recipes.deepseek.gb200.deepseek_v4 import ( - deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_config as flash_perf_config, -) -from megatron.bridge.perf_recipes.deepseek.gb300.deepseek_v4 import ( - deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_config as pro_perf_config, -) from megatron.bridge.recipes.deepseek.gb200.deepseek_v4 import ( deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config as flash_bf16_base_config, ) @@ -50,70 +44,12 @@ pytestmark = pytest.mark.unit -_SHARED_MODEL_EXECUTION_FIELDS = ( - "tensor_model_parallel_size", - "virtual_pipeline_model_parallel_size", - "context_parallel_size", - "expert_model_parallel_size", - "expert_tensor_parallel_size", - "sequence_parallel", - "moe_token_dispatcher_type", - "moe_flex_dispatcher_backend", - "moe_shared_expert_overlap", - "moe_hybridep_num_sms", - "recompute_granularity", - "cross_entropy_fusion_impl", - "moe_pad_experts_for_cuda_graph_inference", - "moe_mlp_glu_interleave_size", - "use_transformer_engine_op_fuser", -) - -_SHARED_ENV_VARS = ( - "CUDA_DEVICE_MAX_CONNECTIONS", - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN", - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API", - "NVLINK_DOMAIN_SIZE", - "USE_MNNVL", - "NVTE_BWD_LAYERNORM_SM_MARGIN", - "NVTE_FWD_LAYERNORM_SM_MARGIN", - "NVTE_NORM_BWD_USE_CUDNN", - "NVTE_NORM_FWD_USE_CUDNN", - "NVTE_ALLOW_NONDETERMINISTIC_ALGO", -) - @pytest.fixture(autouse=True) def _keep_recipe_construction_offline(monkeypatch: pytest.MonkeyPatch) -> None: patch_recipe_construction_dependencies(monkeypatch) -@pytest.mark.parametrize( - ("library_factory", "perf_factory"), - [ - (flash_library_config, flash_perf_config), - (pro_library_config, pro_perf_config), - ], - ids=["flash-gb200", "pro-gb300"], -) -def test_high_scale_library_recipes_share_approved_execution_fields( - library_factory: Callable[[], ConfigContainer], - perf_factory: Callable[[], ConfigContainer], -) -> None: - library_cfg = library_factory() - perf_cfg = perf_factory() - - for field_name in _SHARED_MODEL_EXECUTION_FIELDS: - library_value = getattr(library_cfg.model, field_name, None) - perf_value = getattr(perf_cfg.model, field_name, None) - assert library_value == perf_value, field_name - - assert library_cfg.train.micro_batch_size == perf_cfg.train.micro_batch_size == 1 - assert library_cfg.comm_overlap.overlap_grad_reduce is perf_cfg.comm_overlap.overlap_grad_reduce is True - assert library_cfg.ddp.average_in_collective == perf_cfg.ddp.average_in_collective - for variable_name in _SHARED_ENV_VARS: - assert library_cfg.env_vars[variable_name] == perf_cfg.env_vars[variable_name], variable_name - - @pytest.mark.parametrize( "base_factory", [flash_bf16_base_config, pro_bf16_base_config], @@ -229,21 +165,6 @@ def test_flash_high_scale_recipe_preserves_real_training_contract() -> None: assert cfg.env_vars["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] == 1 assert cfg.env_vars["NVTE_CPU_OFFLOAD_V1"] == 1 - perf_cfg = flash_perf_config() - assert perf_cfg.model.pipeline_model_parallel_size == 1 - assert perf_cfg.model.pipeline_model_parallel_layout is None - assert perf_cfg.train.global_batch_size == 2048 - assert perf_cfg.model.moe_router_force_load_balancing is True - assert perf_cfg.model.dsa_indexer_loss_coeff == 0.01 - assert perf_cfg.model.quant_recipe is None - assert perf_cfg.model.recompute_modules == ["mla_up_proj"] - assert perf_cfg.model.fine_grained_activation_offloading is False - assert perf_cfg.model.offload_modules == [] - assert perf_cfg.mixed_precision.fp8_param_gather is True - assert perf_cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag is True - assert perf_cfg.ddp.grad_reduce_in_fp32 is False - assert perf_cfg.ddp.check_for_nan_in_grad is False - def test_pro_high_scale_recipe_preserves_real_training_contract() -> None: cfg = pro_library_config() @@ -279,19 +200,6 @@ def test_pro_high_scale_recipe_preserves_real_training_contract() -> None: assert cfg.env_vars["TORCH_NCCL_AVOID_RECORD_STREAMS"] == 1 assert cfg.env_vars["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] == 1 - perf_cfg = pro_perf_config() - assert perf_cfg.model.pipeline_model_parallel_size == 4 - assert perf_cfg.model.virtual_pipeline_model_parallel_size == 4 - assert perf_cfg.model.pipeline_model_parallel_layout == "Et*4|(tttt|)*14tmL" - assert perf_cfg.train.global_batch_size == 4096 - assert perf_cfg.model.moe_router_force_load_balancing is True - assert perf_cfg.model.dsa_indexer_loss_coeff == 0.01 - assert perf_cfg.model.quant_recipe is None - assert perf_cfg.mixed_precision.fp8_param_gather is True - assert perf_cfg.optimizer.main_grads_dtype == torch.bfloat16 - assert perf_cfg.ddp.grad_reduce_in_fp32 is False - assert perf_cfg.ddp.check_for_nan_in_grad is False - def test_high_scale_deepseek_v4_recipes_are_exported() -> None: assert recipes.deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config is flash_library_config From 1709cac34e515b8c41cba734ee54fafc50fca13f Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Fri, 14 Aug 2026 14:30:37 -0700 Subject: [PATCH 4/9] fix(recipes): restore TE cross entropy Signed-off-by: Chen Cui --- src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py | 2 +- src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py | 2 +- tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py index f672c5253b..0f634237b8 100644 --- a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py @@ -88,7 +88,7 @@ def deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config() -> ConfigContainer: cfg.model.moe_router_force_load_balancing = False cfg.model.moe_pad_experts_for_cuda_graph_inference = True cfg.model.cross_entropy_loss_fusion = True - cfg.model.cross_entropy_fusion_impl = "native" + cfg.model.cross_entropy_fusion_impl = "te" cfg.model.recompute_granularity = "selective" cfg.model.recompute_modules = ["moe_act", "mhc", "mla_up_proj"] diff --git a/src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py index ea263b2370..c55b2fe13f 100644 --- a/src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py @@ -74,7 +74,7 @@ def deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config() -> ConfigContainer: cfg.model.moe_router_force_load_balancing = False cfg.model.moe_pad_experts_for_cuda_graph_inference = True cfg.model.cross_entropy_loss_fusion = True - cfg.model.cross_entropy_fusion_impl = "native" + cfg.model.cross_entropy_fusion_impl = "te" cfg.model.recompute_granularity = "selective" cfg.model.recompute_modules = ["moe_act", "mhc"] diff --git a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py index a3aa82022b..7a9946c450 100644 --- a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py +++ b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py @@ -64,7 +64,7 @@ def test_base_recipes_enable_precision_independent_fusions( assert cfg.model.moe_pad_experts_for_cuda_graph_inference is True assert getattr(cfg.model, "moe_mlp_glu_interleave_size", None) is None assert cfg.model.use_transformer_engine_op_fuser is False - assert cfg.model.cross_entropy_fusion_impl == "native" + assert cfg.model.cross_entropy_fusion_impl == "te" assert cfg.comm_overlap.overlap_grad_reduce is True assert "NVTE_CUTEDSL_FUSED_GROUPED_MLP" not in cfg.env_vars From 8f7017cad21a328f58932b8d0d3ce07c1a4d241b Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Fri, 14 Aug 2026 15:29:49 -0700 Subject: [PATCH 5/9] feat(recipes): add DeepSeek V4 GB200 packed SFT Signed-off-by: Chen Cui --- .../bridge/recipes/deepseek/__init__.py | 2 ++ .../bridge/recipes/deepseek/deepseek_v4.py | 3 ++ .../bridge/recipes/deepseek/gb200/__init__.py | 2 ++ .../recipes/deepseek/gb200/deepseek_v4.py | 34 ++++++++++++++++++- .../recipes/deepseek/gb300/deepseek_v4.py | 3 +- .../test_deepseek_v4_hardware_recipes.py | 25 +++++++++++++- 6 files changed, 66 insertions(+), 3 deletions(-) diff --git a/src/megatron/bridge/recipes/deepseek/__init__.py b/src/megatron/bridge/recipes/deepseek/__init__.py index b9dcb8c0bf..083772d64b 100644 --- a/src/megatron/bridge/recipes/deepseek/__init__.py +++ b/src/megatron/bridge/recipes/deepseek/__init__.py @@ -51,6 +51,7 @@ 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, @@ -83,6 +84,7 @@ "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", diff --git a/src/megatron/bridge/recipes/deepseek/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/deepseek_v4.py index 5ad648049b..43a255b8d0 100644 --- a/src/megatron/bridge/recipes/deepseek/deepseek_v4.py +++ b/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; diff --git a/src/megatron/bridge/recipes/deepseek/gb200/__init__.py b/src/megatron/bridge/recipes/deepseek/gb200/__init__.py index 7e93b92af4..3e1c7c09d7 100644 --- a/src/megatron/bridge/recipes/deepseek/gb200/__init__.py +++ b/src/megatron/bridge/recipes/deepseek/gb200/__init__.py @@ -18,6 +18,7 @@ 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, ) @@ -26,4 +27,5 @@ "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", ] diff --git a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py index 0f634237b8..ce2e8b336b 100644 --- a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py @@ -88,7 +88,8 @@ def deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config() -> ConfigContainer: cfg.model.moe_router_force_load_balancing = False cfg.model.moe_pad_experts_for_cuda_graph_inference = True cfg.model.cross_entropy_loss_fusion = True - cfg.model.cross_entropy_fusion_impl = "te" + # MCore warns of TE CE stability issues, and matched GB200 testing found no throughput or memory benefit. + cfg.model.cross_entropy_fusion_impl = "native" cfg.model.recompute_granularity = "selective" cfg.model.recompute_modules = ["moe_act", "mhc", "mla_up_proj"] @@ -286,3 +287,34 @@ def deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config() -> ConfigContainer cfg.mixed_precision = bf16_mixed() cfg.mixed_precision.grad_reduce_in_fp32 = True return cfg + + +def deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config() -> ConfigContainer: + """Return the GB200-optimized offline-packed OpenMath SFT config. + + This variant adds HybridEP dispatch, safe uneven-input padding, DSA indexer + training, and MoE fusions to the hardware-agnostic packed SFT recipe. It + preserves non-grouped expert GEMMs for checkpoint export compatibility. + """ + from megatron.bridge.recipes.deepseek.deepseek_v4 import ( + deepseek_v4_flash_sft_openmath_thinking_packed_config, + ) + + cfg = deepseek_v4_flash_sft_openmath_thinking_packed_config() + + cfg.model.apply_dsa_kernel_fusion = True + cfg.model.dsa_indexer_loss_coeff = 0.01 + cfg.model.dsa_indexer_use_sparse_loss = True + + cfg.model.moe_token_dispatcher_type = "flex" + cfg.model.moe_flex_dispatcher_backend = "hybridep" + cfg.model.moe_hybridep_num_sms = 16 + cfg.model.moe_shared_expert_overlap = False + cfg.model.moe_hybridep_pad_uneven_dispatch_inputs = True + cfg.model.moe_grouped_gemm = False + cfg.model.moe_permute_fusion = True + cfg.model.moe_router_fusion = True + + # MCore warns of TE CE stability issues, and matched GB200 testing found no throughput or memory benefit. + cfg.model.cross_entropy_fusion_impl = "native" + return cfg diff --git a/src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py index c55b2fe13f..3301a3ae15 100644 --- a/src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py @@ -74,7 +74,8 @@ def deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config() -> ConfigContainer: cfg.model.moe_router_force_load_balancing = False cfg.model.moe_pad_experts_for_cuda_graph_inference = True cfg.model.cross_entropy_loss_fusion = True - cfg.model.cross_entropy_fusion_impl = "te" + # MCore warns of TE CE stability issues, and matched GB200 testing found no throughput or memory benefit. + cfg.model.cross_entropy_fusion_impl = "native" cfg.model.recompute_granularity = "selective" cfg.model.recompute_modules = ["moe_act", "mhc"] diff --git a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py index 7a9946c450..c606777290 100644 --- a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py +++ b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py @@ -29,6 +29,9 @@ from megatron.bridge.recipes.deepseek.gb200.deepseek_v4 import ( deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config as flash_library_config, ) +from megatron.bridge.recipes.deepseek.gb200.deepseek_v4 import ( + deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config as flash_packed_sft_config, +) from megatron.bridge.recipes.deepseek.gb300.deepseek_v4 import ( deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config as pro_bf16_base_config, ) @@ -64,7 +67,7 @@ def test_base_recipes_enable_precision_independent_fusions( assert cfg.model.moe_pad_experts_for_cuda_graph_inference is True assert getattr(cfg.model, "moe_mlp_glu_interleave_size", None) is None assert cfg.model.use_transformer_engine_op_fuser is False - assert cfg.model.cross_entropy_fusion_impl == "te" + assert cfg.model.cross_entropy_fusion_impl == "native" assert cfg.comm_overlap.overlap_grad_reduce is True assert "NVTE_CUTEDSL_FUSED_GROUPED_MLP" not in cfg.env_vars @@ -130,6 +133,25 @@ def test_flash_mxfp8_recipe_uses_activation_offload_to_fit() -> None: assert cfg.env_vars["NVTE_CPU_OFFLOAD_V1"] == 1 +def test_flash_packed_sft_recipe_uses_gb200_training_contract() -> None: + cfg = flash_packed_sft_config() + + assert cfg.model.cp_partition_mode == "contiguous" + assert cfg.dataset.offline_packing_specs.pad_seq_to_mult == 2 + assert cfg.model.apply_dsa_kernel_fusion is True + assert cfg.model.dsa_indexer_loss_coeff == 0.01 + assert cfg.model.dsa_indexer_use_sparse_loss is True + assert cfg.model.moe_token_dispatcher_type == "flex" + assert cfg.model.moe_flex_dispatcher_backend == "hybridep" + assert cfg.model.moe_hybridep_num_sms == 16 + assert cfg.model.moe_hybridep_pad_uneven_dispatch_inputs is True + assert cfg.model.moe_shared_expert_overlap is False + assert cfg.model.moe_permute_fusion is True + assert cfg.model.moe_router_fusion is True + assert cfg.model.moe_grouped_gemm is False + assert cfg.model.cross_entropy_fusion_impl == "native" + + def test_flash_high_scale_recipe_preserves_real_training_contract() -> None: cfg = flash_library_config() @@ -203,4 +225,5 @@ def test_pro_high_scale_recipe_preserves_real_training_contract() -> None: def test_high_scale_deepseek_v4_recipes_are_exported() -> None: assert recipes.deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config is flash_library_config + assert recipes.deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config is flash_packed_sft_config assert recipes.deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config is pro_library_config From b3d67c91007a84d14ccbbfbb12fa6d16320c4348 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Fri, 14 Aug 2026 15:45:04 -0700 Subject: [PATCH 6/9] docs(cards): add DeepSeek V4 verification card Signed-off-by: Chen Cui --- .../deepseek-v4-flash/card.yaml | 373 ++++++++++++++++++ 1 file changed, 373 insertions(+) create mode 100644 examples/model_verification_cards/deepseek-v4-flash/card.yaml diff --git a/examples/model_verification_cards/deepseek-v4-flash/card.yaml b/examples/model_verification_cards/deepseek-v4-flash/card.yaml new file mode 100644 index 0000000000..1db3f8e8b0 --- /dev/null +++ b/examples/model_verification_cards/deepseek-v4-flash/card.yaml @@ -0,0 +1,373 @@ +# Agent-readable model verification card. +# status: unverified | verified | unsupported | not_applicable + +title: deepseek_v4_flash +summary: > + Performance disclaimer: timing and throughput metrics from functional training items + are sanity checks, not optimized performance results. pretrain_performance.GB200 uses + the tuned canonical GB200 performance recipe. + + DeepSeek-V4-Flash verification covers GPU conversion and GB200 SFT at the item-pinned + Bridge revision. The current 128-GPU GB200 MXFP8 library pretraining recipe uses + PP1/EP64 with natural routing, selective recompute, attention activation offload, and + MXFP8 parameter gather and gradient-buffer reuse. It has completed an eight-step + screen, but full pretraining and direct checkpoint resume remain unverified until the + bounded 100-step runs described below complete. CPU conversion is unsupported because + the 285B model exceeds node CPU RAM. Standard Megatron KV-cache inference is unsupported + for DSV4 hybrid attention; inference uses the HF-native path after export. + + The GB200 packed SFT recipe keeps moe_grouped_gemm=False for export compatibility and + uses the canonical MCore uneven-dispatch padding field. + +verification_index: + model_level: + verified: + - hf_to_megatron_gpu + - megatron_to_hf_gpu + unsupported: + - hf_to_megatron_cpu + - megatron_to_hf_cpu + - inference + unverified: + - manual_forward_pass + training: + GB200: + verified: [sft, sft_long_context, sft_export_inference] + unverified: [pretrain, checkpoint_resume] + unsupported: [peft] + performance: + GB200: verified + +model: + hf_id: deepseek-ai/DeepSeek-V4-Flash + hf_revision: "60d8d70770c6776ff598c94bb586a859a38244f1" # pragma: allowlist secret + architecture: DeepseekV4ForCausalLM + min_transformers_version: "5.12.0" + +verification_environment: + base_container: nvcr.io/nvidia/nemo:26.08.rc4 + bridge_commit: 8f7017cad21a328f58932b8d0d3ce07c1a4d241b # pragma: allowlist secret + +items: + hf_to_megatron_cpu: + status: unsupported + precision: null + command: null + last_verified: null + expected_result: > + CPU import requires approximately 570 GB of CPU RAM for the full BF16 weight + materialisation (285B parameters x 2 bytes). Standard compute nodes do not + provide this headroom; the import consistently exits with OOM before completing. + + hf_to_megatron_gpu: + status: verified + precision: bf16 + bridge_commit: 3ecb452af619bef88083e10a7738362df8f5d0bf # pragma: allowlist secret + command: > + ./scripts/conversion/convert.sh import --executor slurm --device gpu + --nodes 1 --gpus-per-node 4 --ep 4 + --hf-model deepseek-ai/DeepSeek-V4-Flash + --megatron-path work/model-verification/dsv4-flash/import-gpu + --torch-dtype bfloat16 --trust-remote-code + last_verified: 2026-08-04 + expected_result: > + The command exits successfully and creates iter_0000000. The checkpoint + is reloadable and paired GPU export produces correct HF weights. + + megatron_to_hf_cpu: + status: unsupported + precision: null + command: null + last_verified: null + expected_result: > + CPU export requires the full 570 GB BF16 weight set to reside in CPU RAM during + tensor merging; nodes consistently OOM before the export completes. + + megatron_to_hf_gpu: + status: verified + precision: bf16 + bridge_commit: 3ecb452af619bef88083e10a7738362df8f5d0bf # pragma: allowlist secret + command: > + ./scripts/conversion/convert.sh export --executor slurm --device gpu + --nodes 2 --gpus-per-node 4 --ep 8 --pp 4 + --hf-model deepseek-ai/DeepSeek-V4-Flash + --megatron-path work/model-verification/dsv4-flash/import-gpu/iter_0000000 + --hf-path work/model-verification/dsv4-flash/export-gpu + --torch-dtype bfloat16 --export-weight-dtype bfloat16 --trust-remote-code + last_verified: 2026-08-05 + expected_result: > + The command exits successfully and writes 46 BF16 safetensors shards covering + all transformer-layer parameters. The export completes 4176/4176 weight + conversions with zero errors. + + manual_forward_pass: + status: unverified + precision: bf16 + command: > + ./scripts/inference/infer.sh --executor slurm --task model-comparison + --nodes 2 --gpus-per-node 4 --ep 4 + --hf_model_path deepseek-ai/DeepSeek-V4-Flash + --megatron_model_path work/model-verification/dsv4-flash/import-gpu/iter_0000000 + --pp 1 --ep 4 + --prompt "The capital of France is the city of" + --hf-revision 60d8d70770c6776ff598c94bb586a859a38244f1 + last_verified: null + expected_result: > + Blocked: loading the HF model and Megatron model simultaneously on 4 GB200 + GPUs (4 x 192 GB = 768 GB) exceeds the available GPU memory for the combined + 285 B parameter model. Requires 8 or more nodes for side-by-side comparison. + + inference: + status: unsupported + precision: null + command: null + last_verified: null + expected_result: > + Standard Megatron KV-cache autoregressive inference is blocked by an assertion + in DSv4HybridAttention.forward (inference_context must be None). DSV4 hybrid + attention with CSA/DSA indexers requires a dedicated inference path with + MLA-specific KV caching. Inference is supported through HF native pipeline + (deepseek-ai/DeepSeek-V4-Flash) after Megatron-to-HF export. + + pretrain: + GB200: + status: unverified + precision: fp8_mx + enabled_features: + moe_dispatcher: hybridep + command: > + ./scripts/training/train.sh --nodes 32 --gpus-per-node 4 + --recipe deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config + --mode pretrain --dataset megatron-indexed --max_steps 100 + 'dataset.blend=[["work/data/rp2/head_01_text_document"],null]' + dataset.path_to_cache=work/cache/rp2 dataset.num_workers=0 + dataset.random_seed=1234 rng.seed=1234 + scheduler.lr_warmup_iters=10 scheduler.lr_decay_iters=100 + validation.eval_interval=0 validation.eval_iters=0 + ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true + rerun_state_machine.check_for_nan_in_loss=true + checkpoint.exit_on_missing_checkpoint=false checkpoint.load=null + --save_dir work/model-verification/dsv4-flash/pretrain-ref + --save_interval 50 + logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null + dist.distributed_timeout_minutes=120 + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + expected_result: > + The current library recipe resolves to TP1/PP1/EP64/CP1/DP2 with GBS/MBS 256/1, + HybridEP natural routing, native fused cross entropy, selective recompute over + moe, mhc, mla_up_proj, and layernorm, attention activation offload, and MXFP8 + parameter gather and gradient-buffer reuse. Expert capacity, paged stash, and + CUDA graphs remain disabled. Verification requires 100 finite steps with no + skipped or NaN iterations and complete checkpoints at steps 50 and 100; the + existing eight-step screen is insufficient for verified status. + + sft: + GB200: + status: verified + precision: bf16 + bridge_commit: 3ecb452af619bef88083e10a7738362df8f5d0bf # pragma: allowlist secret + enabled_features: + sequence_packing: offline + context_parallel_size: 2 + command: > + ./scripts/training/train.sh --nodes 16 --gpus-per-node 4 + --recipe deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config + --step-func dsv4_step + --pretrained_checkpoint work/models/deepseek-v4-flash-megatron + --save_dir work/model-verification/dsv4-flash/sft-ref + --save_interval 100 --max_steps 100 + model.pipeline_model_parallel_size=4 + 'model.pipeline_model_parallel_layout=Et*11|t*11|t*11|t*10mL' + model.context_parallel_size=2 + model.expert_model_parallel_size=8 + model.moe_grouped_gemm=false + model.recompute_granularity=full model.recompute_method=uniform + model.recompute_num_layers=1 + dataset.seq_length=1024 + scheduler.lr_warmup_iters=10 scheduler.lr_decay_iters=100 + validation.eval_interval=0 validation.eval_iters=0 + rng.seed=5678 + ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true + rerun_state_machine.check_for_nan_in_loss=true + checkpoint.exit_on_missing_checkpoint=false checkpoint.load=null + ddp.overlap_grad_reduce=false + logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null + dist.distributed_timeout_minutes=180 + last_verified: 2026-08-07 + metrics: + initial_loss: 0.94383 + final_loss: 0.28370 + last_10_steps_step_time_ms_avg: 13806.500 + last_10_steps_model_tflops_per_gpu_avg: 13.620 + expected_result: > + On 64 GB200 (16 nodes x 4 GPUs), TP=1/PP=4/EP=8/CP=2/DP=1, GBS/MBS 128/1. + OpenMathInstruct-2 thinking data with offline packing, seq_length=1024. + Full recompute (granularity=full) is required to fit in memory at seq=1024 + with CP=2. HybridEP dispatcher and DSA kernel fusion are enabled via the + GB200-specific recipe. The run completes 100 steps with finite LM and MTP + losses, no skipped or NaN iterations, and a complete step-100 checkpoint. + moe_grouped_gemm=False is preserved for export compatibility. + dist.distributed_timeout_minutes must be 180 or above to allow offline data + packing (approximately 63 minutes) before training begins. + + sft_export_inference: + GB200: + status: verified + precision: bf16 + bridge_commit: 3ecb452af619bef88083e10a7738362df8f5d0bf # pragma: allowlist secret + depends_on: sft + commands: + - > + ./scripts/conversion/convert.sh export --executor slurm --device gpu + --nodes 2 --gpus-per-node 4 --ep 8 + --hf-model deepseek-ai/DeepSeek-V4-Flash + --megatron-path work/model-verification/dsv4-flash/sft-ref/iter_0000100 + --hf-path work/model-verification/dsv4-flash/sft-export + --torch-dtype bfloat16 --export-weight-dtype bfloat16 --trust-remote-code + --not-strict + - > + uv run python + skills/create-model-verification-card/scripts/verify_hf_inference.py + --hf-model work/model-verification/dsv4-flash/sft-export + --prompt "In one short sentence, explain why Paris is important to France." + --max-new-tokens 32 --chat-template --disable-thinking + last_verified: 2026-08-06 + expected_result: > + The exported BF16 SFT checkpoint reloads successfully as DeepseekV4ForCausalLM. + Two independent greedy runs produce byte-identical token IDs and exactly 32 new + tokens with this literal completion: "Paris is important to France because it is + the country's capital and a major center for politics, culture, and economy. It + is also a symbol of French history" + + sft_long_context: + GB200: + status: verified + precision: bf16 + bridge_commit: 3ecb452af619bef88083e10a7738362df8f5d0bf # pragma: allowlist secret + enabled_features: + sequence_packing: offline + context_parallel_size: 2 + command: > + ./scripts/training/train.sh --nodes 16 --gpus-per-node 4 + --recipe deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config + --step-func dsv4_step + --pretrained_checkpoint work/models/deepseek-v4-flash-megatron + --save_dir work/model-verification/dsv4-flash/sft-ref + --save_interval 100 --max_steps 100 + model.pipeline_model_parallel_size=4 + 'model.pipeline_model_parallel_layout=Et*11|t*11|t*11|t*10mL' + model.context_parallel_size=2 model.cp_partition_mode=contiguous + model.expert_model_parallel_size=8 + model.moe_grouped_gemm=false + model.recompute_granularity=full model.recompute_method=uniform + model.recompute_num_layers=1 + dataset.seq_length=1024 + scheduler.lr_warmup_iters=10 scheduler.lr_decay_iters=100 + validation.eval_interval=0 validation.eval_iters=0 + rng.seed=5678 + ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true + rerun_state_machine.check_for_nan_in_loss=true + checkpoint.exit_on_missing_checkpoint=false checkpoint.load=null + ddp.overlap_grad_reduce=false + logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null + dist.distributed_timeout_minutes=180 + last_verified: 2026-08-07 + metrics: + initial_loss: 0.94383 + final_loss: 0.28370 + last_10_steps_step_time_ms_avg: 13806.500 + last_10_steps_model_tflops_per_gpu_avg: 13.620 + expected_result: > + Same run as sft.GB200. CP=2 with contiguous partitioning and offline-packed + OpenMathInstruct-2 at seq_length=1024 demonstrates sequence packing and + context parallelism working together over 100 training steps with finite loss + and no skipped or NaN iterations. + + peft: + all: + status: unsupported + precision: null + command: null + last_verified: null + enabled_features: {} + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + expected_result: > + No LoRA / PEFT recipe exists for DeepSeek-V4 in the current recipe set. + + checkpoint_resume: + GB200: + status: unverified + precision: fp8_mx + depends_on: pretrain + command: > + ./scripts/training/train.sh --nodes 32 --gpus-per-node 4 + --recipe deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config + --mode pretrain --dataset megatron-indexed --max_steps 100 + 'dataset.blend=[["work/data/rp2/head_01_text_document"],null]' + dataset.path_to_cache=work/cache/rp2 dataset.num_workers=0 + dataset.random_seed=1234 rng.seed=1234 + scheduler.lr_warmup_iters=10 scheduler.lr_decay_iters=100 + validation.eval_interval=0 validation.eval_iters=0 + ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true + rerun_state_machine.check_for_nan_in_loss=true + checkpoint.exit_on_missing_checkpoint=false + --load_dir work/model-verification/dsv4-flash/pretrain-ref + --save_dir work/model-verification/dsv4-flash/pretrain-resumed + --save_interval 50 checkpoint.ckpt_step=50 + logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null + dist.distributed_timeout_minutes=120 + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + resume_comparison: + reference_item: pretrain + sentinel_steps: [51, 100] + loss_relative_tolerance: 1.0e-2 + loss_absolute_tolerance: 1.0e-6 + sentinels_match: false + expected_result: > + Verification requires restoring optimizer, scheduler, data-order, and RNG state + from the current recipe's step-50 reference checkpoint, then running steps 51-100 + in a separate output root. Step-51 and step-100 losses must match the uninterrupted + reference within the declared 1% relative and 1e-6 absolute tolerances, and all + four metrics must be recorded. + + pretrain_performance: + GB200: + status: verified + precision: fp8_mx + bridge_commit: 3ecb452af619bef88083e10a7738362df8f5d0bf # pragma: allowlist secret + command: > + ./scripts/training/train.sh --nodes 32 --gpus-per-node 4 + --recipe deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_config + --max_steps 50 + scheduler.lr_warmup_iters=5 scheduler.lr_decay_iters=50 + validation.eval_interval=0 validation.eval_iters=0 + checkpoint.exit_on_missing_checkpoint=false + logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null + dist.distributed_timeout_minutes=120 + last_verified: 2026-08-06 + metrics: + initial_loss: 13.58750 + final_loss: 3.62683 + last_10_steps_step_time_ms_avg: 8134.600 + last_10_steps_model_tflops_per_gpu_avg: 731.400 + expected_result: > + On 128 GB200 (32 nodes x 4 GPUs) with contiguous NVLink-domain placement, + PP=1/EP=64/CP=1/TP=1/DP=2, GBS/MBS 2048/1, FP8-MX with HybridEP dispatcher + and full-iteration CUDA graphs. The 50-step run completes with finite losses, + no skipped or NaN iterations, and all four metrics recorded. First iteration + takes approximately 24 minutes for CUDA graph compilation and kernel warmup; + subsequent iterations run at approximately 8.1 seconds per step at steady state. From e459eb45473820d6b164a673d801f1022d17bdce Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Fri, 14 Aug 2026 16:00:53 -0700 Subject: [PATCH 7/9] docs(cards): correct DeepSeek V4 evidence labels Signed-off-by: Chen Cui --- .../deepseek-v4-flash/card.yaml | 81 +++++++------------ examples/models/deepseek_v4/README.md | 4 +- 2 files changed, 32 insertions(+), 53 deletions(-) diff --git a/examples/model_verification_cards/deepseek-v4-flash/card.yaml b/examples/model_verification_cards/deepseek-v4-flash/card.yaml index 1db3f8e8b0..70ff7d940a 100644 --- a/examples/model_verification_cards/deepseek-v4-flash/card.yaml +++ b/examples/model_verification_cards/deepseek-v4-flash/card.yaml @@ -32,8 +32,8 @@ verification_index: - manual_forward_pass training: GB200: - verified: [sft, sft_long_context, sft_export_inference] - unverified: [pretrain, checkpoint_resume] + verified: [sft, sft_export_inference] + unverified: [pretrain, sft_long_context, checkpoint_resume] unsupported: [peft] performance: GB200: verified @@ -158,13 +158,14 @@ items: last_10_steps_step_time_ms_avg: null last_10_steps_model_tflops_per_gpu_avg: null expected_result: > - The current library recipe resolves to TP1/PP1/EP64/CP1/DP2 with GBS/MBS 256/1, - HybridEP natural routing, native fused cross entropy, selective recompute over - moe, mhc, mla_up_proj, and layernorm, attention activation offload, and MXFP8 - parameter gather and gradient-buffer reuse. Expert capacity, paged stash, and - CUDA graphs remain disabled. Verification requires 100 finite steps with no - skipped or NaN iterations and complete checkpoints at steps 50 and 100; the - existing eight-step screen is insufficient for verified status. + The current library recipe resolves to TP1/PP1/EP64/CP1 with dense DP128, + expert DP2, and GBS/MBS 256/1. It uses HybridEP natural routing, native fused + cross entropy, selective recompute over moe, mhc, mla_up_proj, and layernorm, + attention activation offload, and MXFP8 parameter gather and gradient-buffer + reuse. Expert capacity, paged stash, and CUDA graphs remain disabled. + Verification requires 100 finite steps with no skipped or NaN iterations and + complete checkpoints at steps 50 and 100; the existing eight-step screen is + insufficient for verified status. sft: GB200: @@ -205,8 +206,9 @@ items: last_10_steps_step_time_ms_avg: 13806.500 last_10_steps_model_tflops_per_gpu_avg: 13.620 expected_result: > - On 64 GB200 (16 nodes x 4 GPUs), TP=1/PP=4/EP=8/CP=2/DP=1, GBS/MBS 128/1. - OpenMathInstruct-2 thinking data with offline packing, seq_length=1024. + On 64 GB200 (16 nodes x 4 GPUs), TP1/PP4/EP8/CP2 with dense DP8 and expert + DP1, GBS/MBS 128/1. OpenMathInstruct-2 thinking data uses offline packing + at seq_length=1024. Full recompute (granularity=full) is required to fit in memory at seq=1024 with CP=2. HybridEP dispatcher and DSA kernel fusion are enabled via the GB200-specific recipe. The run completes 100 steps with finite LM and MTP @@ -246,47 +248,23 @@ items: sft_long_context: GB200: - status: verified + status: unverified precision: bf16 - bridge_commit: 3ecb452af619bef88083e10a7738362df8f5d0bf # pragma: allowlist secret enabled_features: sequence_packing: offline context_parallel_size: 2 - command: > - ./scripts/training/train.sh --nodes 16 --gpus-per-node 4 - --recipe deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config - --step-func dsv4_step - --pretrained_checkpoint work/models/deepseek-v4-flash-megatron - --save_dir work/model-verification/dsv4-flash/sft-ref - --save_interval 100 --max_steps 100 - model.pipeline_model_parallel_size=4 - 'model.pipeline_model_parallel_layout=Et*11|t*11|t*11|t*10mL' - model.context_parallel_size=2 model.cp_partition_mode=contiguous - model.expert_model_parallel_size=8 - model.moe_grouped_gemm=false - model.recompute_granularity=full model.recompute_method=uniform - model.recompute_num_layers=1 - dataset.seq_length=1024 - scheduler.lr_warmup_iters=10 scheduler.lr_decay_iters=100 - validation.eval_interval=0 validation.eval_iters=0 - rng.seed=5678 - ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true - rerun_state_machine.check_for_nan_in_loss=true - checkpoint.exit_on_missing_checkpoint=false checkpoint.load=null - ddp.overlap_grad_reduce=false - logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null - dist.distributed_timeout_minutes=180 - last_verified: 2026-08-07 + command: null + last_verified: null metrics: - initial_loss: 0.94383 - final_loss: 0.28370 - last_10_steps_step_time_ms_avg: 13806.500 - last_10_steps_model_tflops_per_gpu_avg: 13.620 + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null expected_result: > - Same run as sft.GB200. CP=2 with contiguous partitioning and offline-packed - OpenMathInstruct-2 at seq_length=1024 demonstrates sequence packing and - context parallelism working together over 100 training steps with finite loss - and no skipped or NaN iterations. + The imported seq_length=1024 SFT run verifies offline packing and CP=2 plumbing, + but it is not a long-context result. Verification requires a separately sized + long-context run with finite loss, no skipped or NaN iterations, and all four + metrics recorded over at least 100 optimizer steps. peft: all: @@ -366,8 +344,9 @@ items: last_10_steps_model_tflops_per_gpu_avg: 731.400 expected_result: > On 128 GB200 (32 nodes x 4 GPUs) with contiguous NVLink-domain placement, - PP=1/EP=64/CP=1/TP=1/DP=2, GBS/MBS 2048/1, FP8-MX with HybridEP dispatcher - and full-iteration CUDA graphs. The 50-step run completes with finite losses, - no skipped or NaN iterations, and all four metrics recorded. First iteration - takes approximately 24 minutes for CUDA graph compilation and kernel warmup; - subsequent iterations run at approximately 8.1 seconds per step at steady state. + PP1/EP64/CP1/TP1 with dense DP128 and expert DP2, GBS/MBS 2048/1, FP8-MX + with HybridEP dispatcher and full-iteration CUDA graphs. The 50-step run + completes with finite losses, no skipped or NaN iterations, and all four + metrics recorded. First iteration takes approximately 24 minutes for CUDA + graph compilation and kernel warmup; subsequent iterations run at + approximately 8.1 seconds per step at steady state. diff --git a/examples/models/deepseek_v4/README.md b/examples/models/deepseek_v4/README.md index 236d5f73a4..9b5c87927f 100644 --- a/examples/models/deepseek_v4/README.md +++ b/examples/models/deepseek_v4/README.md @@ -80,8 +80,8 @@ Available Blackwell pretraining recipes: - `deepseek_v4_flash_pretrain_mxfp8_config`: Adam MXFP8 - `deepseek_v4_flash_pretrain_muon_config`: Muon BF16 - `deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config`: 128-GPU GB200 - Adam MXFP8 with PP1/EP64/DP2/HybridEP, selective recompute, attention - activation offload, and MXFP8 parameter gather/buffer reuse + 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) From fe394e7656a107fe3f0ac2972f239ad16a2fca6b Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Fri, 14 Aug 2026 16:06:54 -0700 Subject: [PATCH 8/9] docs(recipes): clarify DeepSeek V4 validation status Signed-off-by: Chen Cui --- examples/models/deepseek_v4/README.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/models/deepseek_v4/README.md b/examples/models/deepseek_v4/README.md index 9b5c87927f..14cc503e75 100644 --- a/examples/models/deepseek_v4/README.md +++ b/examples/models/deepseek_v4/README.md @@ -105,7 +105,8 @@ 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 -validated FP8 parameter gather/buffer reuse in the high-scale GB200 recipe. +provisional FP8 parameter gather/buffer reuse in the high-scale GB200 recipe; +that PP1 configuration still requires 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. From c5dc1a5a61224b2de902aed94c49c75c0259ba91 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Fri, 14 Aug 2026 17:31:48 -0700 Subject: [PATCH 9/9] feat(recipes): add DeepSeek V4 B200 config Signed-off-by: Chen Cui --- examples/models/deepseek_v4/README.md | 16 +++-- .../references/recipe-index.md | 1 + .../bridge/recipes/deepseek/__init__.py | 6 ++ .../bridge/recipes/deepseek/b200/__init__.py | 21 ++++++ .../recipes/deepseek/b200/deepseek_v4.py | 69 +++++++++++++++++++ .../test_deepseek_v4_hardware_recipes.py | 47 +++++++++++++ 6 files changed, 154 insertions(+), 6 deletions(-) create mode 100644 src/megatron/bridge/recipes/deepseek/b200/__init__.py create mode 100644 src/megatron/bridge/recipes/deepseek/b200/deepseek_v4.py diff --git a/examples/models/deepseek_v4/README.md b/examples/models/deepseek_v4/README.md index 14cc503e75..f2f16509ed 100644 --- a/examples/models/deepseek_v4/README.md +++ b/examples/models/deepseek_v4/README.md @@ -73,12 +73,15 @@ Run `bash conversion.sh` after setting `WORKSPACE` and `MODEL_VARIANT`. See each ## Pretraining Recipes -See [`slurm_pretrain.sh`](slurm_pretrain.sh) for the legacy 32-GPU Slurm launcher. Hardware-specific recipe definitions live in the [`gb200`](../../../src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py) and [`gb300`](../../../src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py) modules. +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_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 @@ -94,9 +97,9 @@ 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 --mode pretrain` with -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. +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 @@ -105,8 +108,9 @@ 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 GB200 recipe; -that PP1 configuration still requires 100-step checkpoint and resume validation. +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. diff --git a/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md b/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md index a6f4d8dbca..ee85c5d2bc 100644 --- a/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md +++ b/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md @@ -74,6 +74,7 @@ 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) | diff --git a/src/megatron/bridge/recipes/deepseek/__init__.py b/src/megatron/bridge/recipes/deepseek/__init__.py index 083772d64b..6afaca1080 100644 --- a/src/megatron/bridge/recipes/deepseek/__init__.py +++ b/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, @@ -80,6 +85,7 @@ "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", diff --git a/src/megatron/bridge/recipes/deepseek/b200/__init__.py b/src/megatron/bridge/recipes/deepseek/b200/__init__.py new file mode 100644 index 0000000000..d035278680 --- /dev/null +++ b/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"] diff --git a/src/megatron/bridge/recipes/deepseek/b200/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/b200/deepseek_v4.py new file mode 100644 index 0000000000..e4b7b21eb3 --- /dev/null +++ b/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 diff --git a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py index c606777290..9220c5884d 100644 --- a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py +++ b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py @@ -18,8 +18,12 @@ import pytest import torch +from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout import megatron.bridge.recipes as recipes +from megatron.bridge.recipes.deepseek.b200.deepseek_v4 import ( + deepseek_v4_flash_pretrain_64gpu_b200_fp8mx_library_config as flash_b200_library_config, +) from megatron.bridge.recipes.deepseek.gb200.deepseek_v4 import ( deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config as flash_bf16_base_config, ) @@ -188,6 +192,48 @@ def test_flash_high_scale_recipe_preserves_real_training_contract() -> None: assert cfg.env_vars["NVTE_CPU_OFFLOAD_V1"] == 1 +def test_flash_b200_nvl8_recipe_preserves_real_training_contract() -> None: + cfg = flash_b200_library_config() + + 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_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" + ) + parsed_layout = PipelineParallelLayerLayout( + cfg.model.pipeline_model_parallel_layout, + 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 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 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_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 + + def test_pro_high_scale_recipe_preserves_real_training_contract() -> None: cfg = pro_library_config() @@ -224,6 +270,7 @@ def test_pro_high_scale_recipe_preserves_real_training_contract() -> None: def test_high_scale_deepseek_v4_recipes_are_exported() -> None: + assert recipes.deepseek_v4_flash_pretrain_64gpu_b200_fp8mx_library_config is flash_b200_library_config assert recipes.deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config is flash_library_config assert recipes.deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config is flash_packed_sft_config assert recipes.deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config is pro_library_config