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..70ff7d940a --- /dev/null +++ b/examples/model_verification_cards/deepseek-v4-flash/card.yaml @@ -0,0 +1,352 @@ +# 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_export_inference] + unverified: [pretrain, sft_long_context, 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 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: + 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), 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 + 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: unverified + precision: bf16 + enabled_features: + sequence_packing: offline + context_parallel_size: 2 + command: null + 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 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: + 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, + 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 bcc1e67525..f2f16509ed 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,51 @@ performance features, so it is not a supported runtime for that recipe. - `conversion.sh` imports HF weights into Megatron Bridge and exports Megatron checkpoints back to HF format. - `inference.sh` runs text generation against an HF or Megatron checkpoint. -- `slurm_pretrain.sh` runs the DeepSeek-V4-Flash pretraining recipes. +- `slurm_pretrain.sh` runs the legacy 32-GPU DeepSeek-V4-Flash pretraining recipes. - `slurm_sft.sh` runs DeepSeek-V4-Flash full SFT end to end (import, then fine-tune) on Hopper or Blackwell, with MTP on or off. Run `bash conversion.sh` after setting `WORKSPACE` and `MODEL_VARIANT`. See each script's header comments for the expected environment variables and `#SBATCH` directives to edit before submitting. ## Pretraining Recipes -See [`slurm_pretrain.sh`](slurm_pretrain.sh) for the Slurm launcher and [`deepseek_v4.py`](../../../src/megatron/bridge/recipes/deepseek/deepseek_v4.py) for recipe definitions. +See [`slurm_pretrain.sh`](slurm_pretrain.sh) for the legacy 32-GPU Slurm launcher. Hardware-specific recipe definitions live in the [`b200`](../../../src/megatron/bridge/recipes/deepseek/b200/deepseek_v4.py), [`gb200`](../../../src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py), and [`gb300`](../../../src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py) modules. Available Blackwell pretraining recipes: - `deepseek_v4_flash_pretrain_mxfp8_config`: Adam MXFP8 - `deepseek_v4_flash_pretrain_muon_config`: Muon BF16 -- `deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_config`: 256-GPU GB300 - performance configuration (requires the PR #4824 container and dev-MCore - stack described above) +- `deepseek_v4_flash_pretrain_64gpu_b200_fp8mx_library_config`: 64-GPU B200 + Adam MXFP8 with PP8/VPP2/EP8, all-to-all dispatch, selective recompute, and + attention activation offload plus MXFP8 parameter gather/buffer reuse +- `deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config`: 128-GPU GB200 + Adam MXFP8 with PP1/EP64/dense-DP128/expert-DP2/HybridEP, selective recompute, + attention activation offload, and MXFP8 parameter gather/buffer reuse +- `deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config`: 256-GPU GB300 + Adam MXFP8 with PP4/VPP4/EP64/HybridEP (requires the PR #4824 container and + dev-MCore stack described above) + +The `_library_config` suffix keeps these real-training definitions distinct +from the corresponding benchmark recipe families. They retain natural routing, +convergence batch sizes, correctness checks, and checkpoint behavior. + +The hardware-count-specific recipes are intentionally not accepted by +`slurm_pretrain.sh`: that legacy launcher overwrites the recipe's batch and +TP/PP/EP/CP settings. Launch them through a site multi-node runner that invokes +`scripts/training/run_recipe.py --recipe --mode pretrain` with +64 ranks for B200 Flash, 128 ranks for GB200 Flash, or 256 ranks for GB300 Pro. +Keep the topology and global batch size owned by the recipe; pass only dataset, +run-length, logging, and checkpoint overrides. + +The GB200 Flash and GB300 Pro library bases enable their supported DSA and +Transformer Engine fused grouped-MLP paths. The new hardware-count-specific +variants additionally enable GLU interleaving; the existing recipe identities +retain their checkpoint tensor layout. The library recipes retain unlimited +natural-routing capacity. All variants retain their precision-specific +training policies: full-iteration CUDA graphs, FP8 parameter gather/buffer +reuse, and reduced-precision gradient reduction remain disabled except for the +provisional FP8 parameter gather/buffer reuse in the high-scale B200 and GB200 +Flash recipes; both configurations still require 100-step checkpoint and resume +validation. `slurm_pretrain.sh` is a GB200 launcher with `TP=1,PP=4,EP=8,CP=1` by default. Indexer loss are disabled for now and is planned for a follow-up. 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..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,9 @@ All recipes live under `src/megatron/bridge/recipes/`. Each function returns a | `deepseek_v2_pretrain_config` | Pretrain | 1 | 4 | 32 | 128 | | `deepseek_v3_pretrain_config` | Pretrain | 2 | 16 | 64 | 2048 | | `deepseek_v3_pretrain_config_32nodes` | Pretrain | 2 | 8 | 32 | 256 | +| `deepseek_v4_flash_pretrain_64gpu_b200_fp8mx_library_config` | Pretrain | 1 | 8 | 8 | 64 (B200 NVL8; unlimited capacity) | +| `deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config` | Pretrain | 1 | 1 | 64 | 128 (GB200; unlimited capacity) | +| `deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config` | Pretrain | 1 | 4 | 64 | 256 (GB300; unlimited capacity) | ### GLM-4.5 diff --git a/src/megatron/bridge/recipes/deepseek/__init__.py b/src/megatron/bridge/recipes/deepseek/__init__.py index 1cf2ac77dd..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, @@ -50,10 +55,13 @@ deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config, deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config, deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config, + deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config, + deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config, ) from .gb300.deepseek_v4 import ( deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config, deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config, + deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config, ) @@ -77,10 +85,14 @@ "deepseek_v4_flash_no_mtp_sft_config", "deepseek_v4_pro_pretrain_config", "deepseek_v4_pro_pretrain_mxfp8_config", + "deepseek_v4_flash_pretrain_64gpu_b200_fp8mx_library_config", "deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config", "deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config", "deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config", + "deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config", + "deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config", "deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config", "deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config", + "deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config", "set_deepseek_v4_pipeline_model_parallel_layout", ] 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/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 696f6cae8b..3e1c7c09d7 100644 --- a/src/megatron/bridge/recipes/deepseek/gb200/__init__.py +++ b/src/megatron/bridge/recipes/deepseek/gb200/__init__.py @@ -17,11 +17,15 @@ deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config, deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config, deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config, + deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config, + deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config, ) __all__ = [ + "deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config", "deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config", "deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config", "deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config", + "deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config", ] diff --git a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py index aee81537eb..ce2e8b336b 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,21 @@ 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" + # 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"] 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 +131,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 +139,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 +151,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 +193,60 @@ 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 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 = 1 + 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 = ["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 = { + **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, + } + 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() @@ -228,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/__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..3301a3ae15 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,18 @@ 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" + # 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"] 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 +117,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 +125,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 +146,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 +190,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..9220c5884d --- /dev/null +++ b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py @@ -0,0 +1,276 @@ +# 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 and convergence guardrails for DeepSeek V4 hardware recipes.""" + +from collections.abc import Callable + +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, +) +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.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, +) +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 + + +@pytest.fixture(autouse=True) +def _keep_recipe_construction_offline(monkeypatch: pytest.MonkeyPatch) -> None: + patch_recipe_construction_dependencies(monkeypatch) + + +@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_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() + + 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 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 + 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 == 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 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() + + 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 + + +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 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,)