diff --git a/.dev.commit b/.dev.commit index 7bcf891a9f..a3512a4504 100644 --- a/.dev.commit +++ b/.dev.commit @@ -1 +1 @@ -2f2f8ebaef90f4d62a18a43808398223b1fc5809 +8a77f4a9b78cd58dee862f234cc5cc84e47311be diff --git a/examples/model_verification_cards/deepseek-v4-flash/card.yaml b/examples/model_verification_cards/deepseek-v4-flash/card.yaml index 1f4c5e866f..71d501d5e9 100644 --- a/examples/model_verification_cards/deepseek-v4-flash/card.yaml +++ b/examples/model_verification_cards/deepseek-v4-flash/card.yaml @@ -4,31 +4,44 @@ title: deepseek_v4_flash summary: > Performance scope: pretrain_performance.GB200 and pretrain_performance.GB300 - use their tuned canonical 128-GPU MXFP8 recipes. - Timing and throughput metrics from functional training items are sanity - checks rather than optimized performance results. + use their tuned canonical 128-GPU MXFP8 recipes. Timing and throughput metrics + from functional training items are sanity checks rather than optimized + performance results. The Megatron-backed DeepSeek V4 workflows in this card require the Megatron-Core dev branch; the default MCore submodule revision in Megatron Bridge r0.6.0 is insufficient. Follow the public setup instructions in examples/models/deepseek_v4/README.md. - DeepSeek-V4-Flash support verification covers GPU conversion, training, and - checkpoint resume on GB200. CPU conversion remains unverified because full - BF16 materialization requires a high-memory CPU node. Optimized Megatron - KV-cache inference is unavailable for DSv4HybridAttention, while the - maintained non-cached full-prefix compatibility path remains unverified for - this model. PEFT also remains unverified because no exact model recipe is - currently exported. SFT export requires nemo_26.08.rc4 or later (TE 2.18) - and the Bridge fix for stream_weights_megatron_to_hf None-task handling. + DeepSeek-V4-Flash verification covers GPU import at the item-pinned Bridge revision. + The verified 64-GPU GB200 MXFP8 library pretraining run uses PP4/VPP4/EP16 with + imported model weights, natural routing, grouped GEMM, TE fused cross entropy, and + selective recompute without activation offload. It completed 100 finite steps with + full-state checkpoints at steps 50 and 100. A direct step-50 continuation restored + optimizer, scheduler, data-order, and RNG state, completed steps 51-100, and passed + both loss-sentinel comparisons. A capacity-1.5 performance screen averaged 206.56 + TFLOP/s/GPU over steps 2-8, but dropped 66.16% of expert assignments and 64.23% of + routing probability mass; that result is benchmark-only and is excluded from the + natural-routing library recipe and verification metrics. CPU conversion remains + unverified because the available nodes lacked sufficient RAM. Standard Megatron + KV-cache inference is unsupported for DSV4 hybrid attention; inference uses the + HF-native path after export. - Known issues tracked for follow-up PRs: - - DSv4ExpertGatedMLPMapping: moe_grouped_gemm=True export produces malformed expert - weight shapes due to apply_swiglu_sharded_factory shard reassembly in the export path. - Workaround: use moe_grouped_gemm=False for export. Training with moe_grouped_gemm=True - is unaffected. - - MCore moe_hybridep_pad_uneven_dispatch_inputs missing from dev HEAD; apply 1-line patch - to megatron/core/transformer/transformer_config.py until upstream MCore PR lands. + The GB200 packed SFT recipe uses grouped GEMM, static packed-sequence shapes, + selective whole-MoE recompute, and attention activation offload while preserving + the base SFT objective. Its final objective-preserving configuration completed + 100 finite steps and a fresh-process checkpoint reload. GPU export of both the + imported model and the step-100 SFT checkpoint produced exact BF16 inventories + and passed independent strict GPU-only HF reloads. The public post-SFT workflow + also produced identical bounded greedy completions from two independent reloads. + A separate 100-step CP2 run verifies offline packing with contiguous context + parallelism. + + The GB200 packed PEFT recipe applies rank-32 LoRA to the DeepSeek MLA projections + and to separate routed/shared expert FC1/FC2 projections. Its execution path uses + scoped Transformer Engine CUDA graphs and tuned HybridEP dispatch. It completed + 100 finite steps, saved an adapter-only checkpoint, and reloaded that checkpoint + in a fresh process for finite step 101. verification_index: model_level: @@ -37,13 +50,12 @@ verification_index: - megatron_to_hf_gpu unverified: - hf_to_megatron_cpu - - megatron_to_hf_cpu - inference - manual_forward_pass + - megatron_to_hf_cpu training: GB200: - verified: [pretrain, sft, sft_long_context, sft_export_inference, checkpoint_resume] - unverified: [peft] + verified: [pretrain, sft, sft_long_context, sft_export_inference, peft, checkpoint_resume] performance: GB200: verified GB300: verified @@ -55,8 +67,8 @@ model: min_transformers_version: "5.12.0" verification_environment: - base_container: nvcr.io/nvidia/nemo:26.08.rc4 - bridge_commit: 3ecb452af619bef88083e10a7738362df8f5d0bf # pragma: allowlist secret + base_container: nvcr.io/nvidia/nemo:26.08 + bridge_commit: 07c5cc30d0510d437bfbbe7b12d9d98387295d89 # pragma: allowlist secret items: hf_to_megatron_cpu: @@ -64,72 +76,80 @@ items: precision: bf16 command: > ./scripts/conversion/convert.sh import --executor slurm --device cpu - --nodes 1 --hf-model deepseek-ai/DeepSeek-V4-Flash + --nodes 1 + --hf-model deepseek-ai/DeepSeek-V4-Flash --hf-revision 60d8d70770c6776ff598c94bb586a859a38244f1 --megatron-path work/model-verification/dsv4-flash/import-cpu --torch-dtype bfloat16 --trust-remote-code last_verified: null expected_result: > - CPU import requires approximately 570 GB merely for the full BF16 weight - materialization, plus conversion overhead. Prior attempts exhausted the - available node memory, but that resource limit does not establish a - product-level lack of support. A high-memory-node run must complete and - produce a reloadable checkpoint before this item is promoted. + CPU import requires approximately 570 GB of CPU RAM for the full BF16 weight + materialisation (285B parameters x 2 bytes). Verification requires a high-memory + node with enough additional headroom for conversion workspace and strict checkpoint reload. 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 + --hf-revision 60d8d70770c6776ff598c94bb586a859a38244f1 --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. + The command exits successfully and creates a reloadable iter_0000000 + Megatron checkpoint. megatron_to_hf_cpu: status: unverified precision: bf16 command: > ./scripts/conversion/convert.sh export --executor slurm --device cpu - --nodes 1 --hf-model deepseek-ai/DeepSeek-V4-Flash + --nodes 1 + --hf-model deepseek-ai/DeepSeek-V4-Flash --hf-revision 60d8d70770c6776ff598c94bb586a859a38244f1 - --megatron-path work/model-verification/dsv4-flash/import-gpu/iter_0000000 + --megatron-path work/model-verification/dsv4-flash/import-cpu/iter_0000000 --hf-path work/model-verification/dsv4-flash/export-cpu --torch-dtype bfloat16 --trust-remote-code last_verified: null expected_result: > - CPU export requires approximately 570 GB merely for the full BF16 weight - set, plus tensor-merging overhead. Prior attempts exhausted the available - node memory, but that resource limit does not establish a product-level - lack of support. A high-memory-node run must complete and produce a - reloadable Hugging Face checkpoint before this item is promoted. + CPU export requires a high-memory node with enough headroom for the full BF16 + weight set and tensor-merging workspace, followed by strict Hugging Face reload. megatron_to_hf_gpu: status: verified precision: bf16 + bridge_commit: 2859a3ee0f02743c54b390de91fee878c6de8909 # pragma: allowlist secret command: > ./scripts/conversion/convert.sh export --executor slurm --device gpu - --nodes 2 --gpus-per-node 4 --ep 8 --pp 4 + --nodes 2 --gpus-per-node 4 --tp 1 --pp 1 --ep 8 --etp 1 --hf-model deepseek-ai/DeepSeek-V4-Flash + --hf-revision 60d8d70770c6776ff598c94bb586a859a38244f1 --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 + --not-strict + last_verified: 2026-08-21 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. + Two GB200 nodes produced 46 complete safetensors shards containing all 35,020 + expected non-scale source keys. BF16 export intentionally omitted 34,167 + quantization-scale companions; 35,017 exported tensors are BF16 and the three + I32 tid2eid routing tables exactly match their source values after the expected + integer cast. Two independent Transformers processes each strictly reloaded all + 1,500 model modules across four GPUs with no CPU or disk placement and produced + the same bounded greedy token. --not-strict permits only the intentional scale + omission; exact key, shard, dtype, routing-table value, and reload checks remain + required correctness gates. 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 + ./scripts/inference/infer.sh --task model-comparison + --nodes 1 --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 @@ -137,9 +157,9 @@ items: --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. + The attempted four-GPU side-by-side load exceeded the available 768 GB of + aggregate GPU memory. Verification requires a larger allocation sized from + measured peak memory; no minimum working topology has been established. inference: status: unverified @@ -154,48 +174,58 @@ items: --tp 1 --pp 1 --ep 8 --etp 1 --trust-remote-code last_verified: null expected_result: > - Optimized KV-cache autoregressive inference currently conflicts with + Optimized KV-cache autoregressive inference conflicts with DSv4HybridAttention's requirement that inference_context be None. The - maintained legacy full-prefix task intentionally disables that context - and is therefore a plausible non-cached compatibility path, but no exact - DeepSeek V4 run has completed. The command must finish one synchronous - deterministic greedy generation before this item is promoted. + maintained legacy full-prefix task disables that context and remains a + plausible non-cached compatibility path, but the exact model command must + complete deterministic greedy generation before this item is promoted. pretrain: GB200: status: verified - precision: bf16 - enabled_features: {} + precision: fp8_mx + bridge_commit: dff6445009ebd054028a883ea6d84f9a16ad6b9c # pragma: allowlist secret + enabled_features: + moe_dispatcher: hybridep command: > ./scripts/training/train.sh --nodes 16 --gpus-per-node 4 - --recipe deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config - --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 + --recipe deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_library_config + --mode pretrain --dataset megatron-indexed --max_steps 100 + --pretrained_checkpoint work/model-verification/dsv4-flash/import-gpu/iter_0000000 + 'dataset.blend=[["work/data/the-pile/my-gpt3_08_text_document"],null]' + dataset.path_to_cache=work/cache/the-pile dataset.num_workers=0 dataset.random_seed=1234 rng.seed=1234 - scheduler.lr_warmup_iters=40 scheduler.lr_decay_iters=100 + 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 + checkpoint.exit_on_missing_checkpoint=false checkpoint.finetune=false + checkpoint.load=null checkpoint.load_optim=false checkpoint.load_rng=false + checkpoint.save_optim=true checkpoint.save_rng=true checkpoint.async_save=false --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=60 - last_verified: 2026-08-04 + dist.distributed_timeout_minutes=120 + last_verified: 2026-08-27 metrics: - initial_loss: 12.64786 - final_loss: 6.05023 - last_10_steps_step_time_ms_avg: 10415.760 - last_10_steps_model_tflops_per_gpu_avg: 142.910 - last_10_steps_tokens_per_second_per_gpu_avg: 1573.001 + initial_loss: 7.250204 + final_loss: 3.280722 + last_10_steps_step_time_ms_avg: 7856.370 + last_10_steps_model_tflops_per_gpu_avg: 189.690 + last_10_steps_tokens_per_second_per_gpu_avg: 2085.441495 expected_result: > - On 64 GB200 (16 nodes x 4 GPUs), the recipe runs with TP=1/PP=8/EP=8/CP=1/DP=1 - and GBS/MBS 256/1. Data is RP2 tokenized with Llama SentencePiece (token IDs - 0-31999 are within DSV4 vocab 0-129279; NullTokenizer is used). The run - completes 100 steps from random initialisation with finite loss, no skipped or - NaN iterations, and complete checkpoints at steps 50 and 100. Note: set - dataset.num_workers=0 to avoid an NCCL pipeline timeout during data loading. + On 64 GB200 (16 nodes x 4 GPUs), TP1/PP4/VPP4/EP16/CP1 uses dense DP16, + expert DP1, and GBS/MBS 256/1. The command loads the imported model weights + while starting optimizer and RNG state fresh. It uses HybridEP natural routing, + grouped GEMM, TE fused cross entropy, selective recompute over mhc and + mla_up_proj, and MXFP8 parameter gather and gradient-buffer reuse without + activation offload. Expert capacity, paged stash, CUDA graphs, and forced + load balancing remain disabled. The run completed 100 finite LM/MTP-loss + steps with no skipped or NaN iterations and wrote complete grouped-MXFP8 + checkpoints containing model, optimizer, scheduler, data-order, and RNG + state at steps 50 and 100. The item-specific Bridge revision pins the + compatible Megatron-LM dev revision used by this workload without changing + the card's default environment for other verification items. sft: GB200: @@ -203,46 +233,48 @@ items: precision: bf16 enabled_features: sequence_packing: offline - context_parallel_size: 2 + moe_dispatcher: hybridep command: > - ./scripts/training/train.sh --nodes 16 --gpus-per-node 4 - --recipe deepseek_v4_flash_sft_openmath_thinking_packed_config - --step-func dsv4_step + ./scripts/training/train.sh --nodes 8 --gpus-per-node 4 + --recipe deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config + --mode sft --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 + --seq_length 1024 + --pipeline_model_parallel_size 4 --context_parallel_size 1 + --expert_model_parallel_size 8 '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 + checkpoint.load_optim=false checkpoint.load_rng=false + checkpoint.save_optim=false 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 + last_verified: 2026-08-20 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 - last_10_steps_tokens_per_second_per_gpu_avg: 148.336 + initial_loss: 0.9435847 + final_loss: 0.2834835 + last_10_steps_step_time_ms_avg: 8575.85 + last_10_steps_model_tflops_per_gpu_avg: 42.18 + last_10_steps_tokens_per_second_per_gpu_avg: 477.620294 + peak_allocated_memory_gib: 184.96 + peak_reserved_memory_gib: 187.04 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. + On 32 GB200 (8 nodes x 4 GPUs), TP1/PP4/EP8/CP1 with dense DP8 and expert + DP1, GBS/MBS 128/1. OpenMathInstruct-2 thinking data uses offline packing + at seq_length=1024 with fixed token and cumulative-boundary shapes. The + GB200 recipe uses selective recompute over moe, mhc, mla_up_proj, and + layernorm, attention activation offload, HybridEP dispatch, and DSA kernel + fusion without enabling the optional DSA indexer loss. The run completed + 100 steps with finite LM and MTP losses, no skipped or NaN iterations, and a + complete model checkpoint at step 100. A fresh process reloaded that checkpoint + and completed finite step 101. moe_grouped_gemm=True is a recipe default. dist.distributed_timeout_minutes must be 180 or above to allow offline data packing (approximately 63 minutes) before training begins. @@ -250,34 +282,43 @@ items: GB200: status: verified precision: bf16 + bridge_commit: 2859a3ee0f02743c54b390de91fee878c6de8909 # 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 + --hf-revision 60d8d70770c6776ff598c94bb586a859a38244f1 --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 + ./scripts/inference/infer.sh --task hf-inference + --nodes 2 --gpus-per-node 4 --tasks-per-node 1 --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 + --max-new-tokens 32 --chat-template --disable-thinking --trust-remote-code + --dtype bfloat16 --device-map auto --autocast --require-gpu-only + last_verified: 2026-08-21 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" + The objective-preserving step-100 SFT checkpoint exported successfully on two + GB200 nodes. Its 46 complete safetensors shards contain all 35,020 expected + non-scale source keys: 35,017 BF16 tensors and three value-matched I32 tid2eid + routing tables, with 34,167 quantization-scale companions intentionally omitted + for BF16. Two maintained-verifier processes independently reloaded all 1,500 model + modules across four GPUs each, reported no missing, unexpected, mismatched, or + error loading entries, used no CPU or disk placement, and generated exactly 32 + new tokens at the 32-token maximum with identical output. The literal completion was: "Paris is + important to France because it is the country's capital, a major economic and + cultural center, and a symbol of French history, art, and architecture." sft_long_context: GB200: status: verified precision: bf16 + bridge_commit: 3ecb452af619bef88083e10a7738362df8f5d0bf # pragma: allowlist secret enabled_features: sequence_packing: offline context_parallel_size: 2 @@ -286,7 +327,7 @@ items: --recipe deepseek_v4_flash_sft_openmath_thinking_packed_config --step-func dsv4_step --pretrained_checkpoint work/models/deepseek-v4-flash-megatron - --save_dir work/model-verification/dsv4-flash/sft-ref + --save_dir work/model-verification/dsv4-flash/sft-long-context-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' @@ -313,59 +354,105 @@ items: last_10_steps_model_tflops_per_gpu_avg: 13.620 last_10_steps_tokens_per_second_per_gpu_avg: 148.336 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. + 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: GB200: - status: unverified + status: verified precision: bf16 - command: null - last_verified: null - enabled_features: {} + bridge_commit: 043d9cdaa1ac0cb52437d8d417a9f3150dfcae3b # pragma: allowlist secret + command: > + ./scripts/training/train.sh --nodes 8 --gpus-per-node 4 + --recipe deepseek_v4_flash_peft_openmath_thinking_packed_gb200_config + --mode lora --step-func dsv4_step + --pretrained_checkpoint work/models/deepseek-v4-flash-megatron + --save_dir work/model-verification/dsv4-flash/peft-ref + --save_interval 100 --max_steps 100 + --seq_length 1024 + --pipeline_model_parallel_size 4 --context_parallel_size 1 + --expert_model_parallel_size 8 + 'model.pipeline_model_parallel_layout=Et*11|t*11|t*11|t*10mL' + 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 + checkpoint.load_optim=false checkpoint.load_rng=false + checkpoint.save_optim=false + 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-25 + enabled_features: + sequence_packing: offline + cuda_graph: + implementation: transformer_engine + scopes: [moe_router, moe_preprocess] + moe_dispatcher: hybridep metrics: - initial_loss: null - final_loss: null - last_10_steps_step_time_ms_avg: null - last_10_steps_model_tflops_per_gpu_avg: null - last_10_steps_tokens_per_second_per_gpu_avg: null + initial_loss: 0.9435847 + final_loss: 0.2914867 + last_10_steps_step_time_ms_avg: 6944.33 + last_10_steps_model_tflops_per_gpu_avg: 52.00 + last_10_steps_tokens_per_second_per_gpu_avg: 589.833720 + peak_allocated_memory_gib: 53.712 + peak_reserved_memory_gib: 53.953 expected_result: > - No exact DeepSeek V4 PEFT recipe is currently exported. Recipe absence - does not establish an architectural lack of support; this item remains - unverified until a public recipe and bounded run pass the PEFT gates. + On 32 GB200 (8 nodes x 4 GPUs), TP1/PP4/EP8/CP1 with dense DP8 and expert + DP1, GBS/MBS 128/1. The recipe applies rank-32, alpha-32 LoRA with zero + dropout to linear_q_down_proj, linear_q_up_proj, linear_kv_proj, + linear_proj, linear_fc1, and linear_fc2. Routed experts use separate + per-expert adapters; shared experts are also adapted. The run preserves + the packed OpenMath thinking data, objective, natural routing, seed, and + warmup/decay horizon used by the verified SFT item while freezing the base + model and using PEFT's own 1e-4 peak learning rate. + It disables activation recompute and offload because the reduced PEFT + training-state footprint fits with 53.712 GiB peak allocated memory. + Transformer Engine CUDA graphs cover moe_router and moe_preprocess, while + HybridEP uses 32 flex-dispatcher SMs, 8 ranks per NVLink domain, 128-token + combine chunks, a 72-GPU domain, and MNNVL. All 100 steps completed with + finite LM/MTP losses and zero skipped or NaN iterations. LM/MTP losses + were 0.9435847/0.2250465 at step 1 and 0.2914867/0.02638751 at step 100. + The adapter-only checkpoint + reloaded after the same base checkpoint in a fresh process and completed + finite step 101 with LM loss 0.2827368 and MTP loss 0.02594211. checkpoint_resume: GB200: status: verified - precision: bf16 + precision: fp8_mx + bridge_commit: dff6445009ebd054028a883ea6d84f9a16ad6b9c # pragma: allowlist secret depends_on: pretrain command: > ./scripts/training/train.sh --nodes 16 --gpus-per-node 4 - --recipe deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config - --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 + --recipe deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_library_config + --mode pretrain --dataset megatron-indexed --max_steps 100 + 'dataset.blend=[["work/data/the-pile/my-gpt3_08_text_document"],null]' + dataset.path_to_cache=work/cache/the-pile dataset.num_workers=0 dataset.random_seed=1234 rng.seed=1234 - scheduler.lr_warmup_iters=40 scheduler.lr_decay_iters=100 + 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 - train.empty_unused_memory_level=2 + checkpoint.exit_on_missing_checkpoint=false checkpoint.finetune=false + checkpoint.load_optim=true checkpoint.load_rng=true + checkpoint.save_optim=true checkpoint.save_rng=true checkpoint.async_save=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=60 - last_verified: 2026-08-04 + dist.distributed_timeout_minutes=120 + last_verified: 2026-08-27 metrics: - initial_loss: 6.76733 - final_loss: 6.05425 - last_10_steps_step_time_ms_avg: 14648.110 - last_10_steps_model_tflops_per_gpu_avg: 101.560 - last_10_steps_tokens_per_second_per_gpu_avg: 1118.506 + initial_loss: 4.028752 + final_loss: 3.261145 + last_10_steps_step_time_ms_avg: 7995.850 + last_10_steps_model_tflops_per_gpu_avg: 186.640 + last_10_steps_tokens_per_second_per_gpu_avg: 2049.062951 resume_comparison: reference_item: pretrain sentinel_steps: [51, 100] @@ -373,17 +460,21 @@ items: loss_absolute_tolerance: 1.0e-6 sentinels_match: true expected_result: > - The command restores optimizer, scheduler, data-order, and RNG state from - pretrain-ref step 50 and runs steps 51-100 in a separate output root. - Step-51 loss 6.767332 matches the uninterrupted reference 6.767334 exactly - (absolute delta 0.000002, 0.00003%). Step-100 loss 6.054250 matches reference - 6.050228 with delta 0.004022 (0.066%), within the declared 1% tolerance. - Both sentinels match and all five metrics are recorded. + The direct 64-GPU TP1/PP4/VPP4/EP16 continuation restores model, optimizer, + scheduler, data-order, and RNG state from the complete step-50 reference + checkpoint, starts at step 51, and finishes at step 100 in a distinct output + root. All 50 resumed steps are finite with no skipped or NaN iterations, all + five metrics are recorded, the post-setup configuration persists, and a + complete full-state step-100 checkpoint is saved. Step-51 resumed/reference + LM loss is 4.028752/4.028773, an absolute difference of 0.000021. Step-100 + resumed/reference LM loss is 3.261145/3.280722, an absolute difference of + 0.019577. Both sentinels pass the declared 1% relative and 1e-6 absolute gate. 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 @@ -401,14 +492,13 @@ items: last_10_steps_model_tflops_per_gpu_avg: 731.400 last_10_steps_tokens_per_second_per_gpu_avg: 8056.450 expected_result: > - On 128 GB200 (32 nodes x 4 GPUs, --segment=16 for NVLink locality), - 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 five 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. - The --segment=16 flag is required to ensure nodes are allocated within NVLink - domains for MNNVL communication. + On 128 GB200, PP1/EP64/CP1/TP1/DP2 with GBS/MBS 2048/1, FP8-MX, + HybridEP dispatch, and full-iteration CUDA graphs completes 50 steps + with finite losses, no skipped or NaN iterations, and all five metrics + recorded. This benchmark uses forced routing, a static expert-rank + capacity, paged stash, and a different batch/objective contract, so its + metrics must not be compared with natural-routing library pretraining + as convergence evidence. GB300: status: verified diff --git a/examples/models/deepseek_v4/README.md b/examples/models/deepseek_v4/README.md index e57dda2e71..6ef4e3a2cc 100644 --- a/examples/models/deepseek_v4/README.md +++ b/examples/models/deepseek_v4/README.md @@ -1,181 +1,227 @@ # DeepSeek V4 -End-to-end conversion and inference scripts for the DeepSeek V4 family on Megatron Bridge. +Megatron Bridge supports checkpoint conversion and training for the DeepSeek V4 +family. See the +[DeepSeek-V4-Flash verification card](../../model_verification_cards/deepseek-v4-flash/card.yaml) +for the exact commands, revisions, metrics, and current verification status. -The bridge supports four published variants out of the same code path. The on-disk quantisation differs between post-trained (Flash, Pro) and pretrained-only (Flash-Base, Pro-Base) models — see [`docs/models/deepseek/deepseek-v4.md`](../../../docs/models/deepseek/deepseek-v4.md) for the per-variant scheme. +## Supported Variants -## MCore Checkout +All variants use the same bridge. Import dequantizes the published checkpoint to +the requested Megatron training dtype. -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. +| Variant | Hugging Face model | Published weight format | +|---------|--------------------|-------------------------| +| DeepSeek-V4-Flash | `deepseek-ai/DeepSeek-V4-Flash` | FP8 attention and MXFP4 experts | +| DeepSeek-V4-Flash-Base | `deepseek-ai/DeepSeek-V4-Flash-Base` | FP8 with float32 scales | +| DeepSeek-V4-Pro | `deepseek-ai/DeepSeek-V4-Pro` | FP8 attention and MXFP4 experts | +| DeepSeek-V4-Pro-Base | `deepseek-ai/DeepSeek-V4-Pro-Base` | FP8 with float32 scales | -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 -not in the Megatron Bridge dependency graph, including container-provided -packages such as NeMo Run. When switching MCore in this container, preserve -those packages and then restore the aggregate container pins: +## Runtime Requirements + +DeepSeek V4 training requires a compatible Megatron-LM `dev` revision; the +Megatron-LM revision pinned for the Megatron Bridge `main` branch is not a +supported training runtime. Switch the submodule before syncing the environment: ```bash ./scripts/switch_mcore.sh dev -uv sync --inexact -(cd /opt/NeMo-FW && uv sync --locked --all-groups --inexact) -export UV_NO_SYNC=1 +uv sync ``` -Keep `UV_NO_SYNC=1` set while running the examples. It makes `uv run` calls, -including those inside the example wrapper scripts, skip the implicit exact -sync. The unlocked dev sync updates `uv.lock`; do not commit that generated -change. From a clean checkout, return to the pinned main-branch submodule with: +`fast-hadamard-transform` is required by DSA and is installed from the pinned +source dependency by `uv sync`. Run the examples in a CUDA-enabled Megatron +Bridge container; see the +[repository installation instructions](../../../README.md#-installation). + +The Slurm launchers require an account, partition, and container image. Set +them before running any `--executor slurm` conversion or `train.sh` command: ```bash -./scripts/switch_mcore.sh main -git restore uv.lock -uv sync --locked --inexact -(cd /opt/NeMo-FW && uv sync --locked --all-groups --inexact) -export UV_NO_SYNC=1 +export SLURM_ACCOUNT=your-account +export SLURM_PARTITION=your-partition +export CONTAINER_IMAGE=your-container-image ``` -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): -`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 -Megatron Bridge `main` branch does not provide the required DeepSeek V4 -performance features, so it is not a supported runtime for that recipe. - -| Variant | HF path | Quant scheme | Validation | -|---------|---------|--------------|------------| -| DeepSeek-V4-Flash | `deepseek-ai/DeepSeek-V4-Flash` | FP8 attn + MXFP4 experts | Verified on GB200, last-token logit cosine 0.96-0.99 (short prompts ~0.98, long prompts >1024 tokens ~0.96-0.99) vs official inference | -| DeepSeek-V4-Flash-Base | `deepseek-ai/DeepSeek-V4-Flash-Base` | uniform FP8 (F32 scales) | Verified on GB200, last-real-token logit cosine 0.9866-0.9930, mean 0.9907 vs official inference | -| DeepSeek-V4-Pro | `deepseek-ai/DeepSeek-V4-Pro` | FP8 attn + MXFP4 experts | Import, export, inference verified on GB200 (PP=4 EP=8) and H100 (PP=16 EP=8) | -| DeepSeek-V4-Pro-Base | `deepseek-ai/DeepSeek-V4-Pro-Base` | uniform FP8 (F32 scales) | Same bridge code as Pro; end-to-end untested | - -## Examples - -- `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_sft.sh` runs DeepSeek-V4-Flash full SFT end to end (import, then fine-tune) on Hopper or Blackwell, with MTP on or off. +Pass `--mount HOST` or `--mount HOST:CONTAINER` for each dataset, checkpoint, +or output path that is not already visible inside the container. -Bridge-created DeepSeek-V4 providers default to the `flex`/HybridEP MoE -dispatcher with 16 dispatcher SMs and unfused HybridEP permutation. Native -Megatron checkpoints retain the dispatcher configuration with which they were -saved, so older checkpoints may still use `alltoall`. HybridEP requires a -runtime containing the HybridEP package. It auto-detects accessible ranks and -fabric support when its topology environment variables are unset. To override -detection, set -`NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN`, `NVLINK_DOMAIN_SIZE`, and -`USE_MNNVL` to values that describe the actual deployment. Hardware recipes -retain their independently qualified dispatcher overrides. +## Conversion -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. +The shared conversion CLI performs dtype-aware FP8/MXFP4 dequantization during +import; no preprocessing script is required. A DeepSeek-V4-Flash GPU import is: -## 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. - -Available Blackwell pretraining recipes: +```bash +./scripts/conversion/convert.sh import \ + --executor local --device gpu --gpus-per-node 4 \ + --hf-model deepseek-ai/DeepSeek-V4-Flash \ + --megatron-path work/models/deepseek-v4-flash \ + --tp 1 --pp 1 --ep 4 \ + --torch-dtype bfloat16 --trust-remote-code +``` -- `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) +Export accepts a different parallel layout because distributed checkpoints can +be resharded while loading: -`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. +```bash +./scripts/conversion/convert.sh export \ + --executor slurm --device gpu --nodes 2 --gpus-per-node 4 \ + --hf-model deepseek-ai/DeepSeek-V4-Flash \ + --megatron-path work/models/deepseek-v4-flash/iter_0000000 \ + --hf-path work/models/deepseek-v4-flash-hf \ + --tp 1 --pp 1 --ep 8 \ + --torch-dtype bfloat16 --export-weight-dtype bfloat16 \ + --trust-remote-code --not-strict +``` -Before submitting, set `CONTAINER_IMAGE`. For DCLM, also set `DCLM_DATA_DIR` and `DCLM_CACHE`. Use `CONTAINER_MOUNTS` and `EXTRA_PYTHONPATH` for cluster-specific data, checkouts, and Python dependencies. +BF16 export intentionally omits the source checkpoint's quantization-scale +companions, so it uses `--not-strict`. The verification card pairs that option +with exact key, shard, dtype, integer-routing-table, and strict HF reload checks. +GPU import and GPU export are verified. CPU conversion and manual HF/Megatron +forward correlation remain unverified at the revisions recorded in the card. -The bridge's `maybe_modify_loaded_hf_weight` hook dispatches dequantisation by tensor dtype: +[`conversion.sh`](conversion.sh) is a convenience wrapper for import, export, +and optional round-trip checks. Set `MODEL_VARIANT`, `WORKSPACE`, and the +parallelism and executor variables documented in its header before running it. -- `int8` -> MXFP4 packed nibbles -> `bfloat16` via the E2M1 lookup table and per-row 16-K-tile E8M0 scales -- `float8_e4m3fn` with companion `.scale` -> `bfloat16` via 128x128 block-scale expansion, handling both E8M0 and F32 scale dtypes +## Pretraining -No external dequantisation script is required. +Use the public [`scripts/training/train.sh`](../../../scripts/training/train.sh) +launcher. Hardware-qualified library recipes are defined under +[`recipes/deepseek`](../../../src/megatron/bridge/recipes/deepseek): -## SFT Recipes +| Hardware | Recipe | Precision and optimizer | +|----------|--------|-------------------------| +| H100, 32 GPUs | `deepseek_v4_flash_pretrain_32gpu_h100_bf16_config` | BF16 Adam | +| H100, 32 GPUs | `deepseek_v4_flash_pretrain_32gpu_h100_fp8mx_config` | MXFP8 Adam | +| H100, 32 GPUs | `deepseek_v4_flash_pretrain_32gpu_h100_bf16_muon_config` | BF16 Muon | +| GB200, 64 GPUs | `deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config` | BF16 Adam | +| GB200, 64 GPUs | `deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config` | MXFP8 Adam | +| GB200, 64 GPUs | `deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config` | BF16 Muon | +| GB200, 64 GPUs | `deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_library_config` | MXFP8 Adam | +| GB300, 32 GPUs | `deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config` | BF16 Adam | +| GB300, 32 GPUs | `deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config` | MXFP8 Adam | -Full-parameter SFT of DeepSeek-V4-Flash. See [`slurm_sft.sh`](slurm_sft.sh) for the end-to-end launcher and [`deepseek_v4.py`](../../../src/megatron/bridge/recipes/deepseek/deepseek_v4.py) for the recipe definitions. +For example, a short generated-data run of the GB200 library recipe is: -| Recipe | MTP | mHC kernel | Hardware | -|--------|-----|------------|----------| -| `deepseek_v4_flash_sft_config` | on | fused (cuTile, sm_100) | Blackwell (Hopper: set `use_fused_mhc=False`) | -| `deepseek_v4_flash_no_mtp_sft_config` | off | fused (cuTile, sm_100) | Blackwell (Hopper: set `use_fused_mhc=False`) | +```bash +./scripts/training/train.sh --nodes 16 --gpus-per-node 4 \ + --recipe deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_library_config \ + --mode pretrain --dataset mock --max_steps 10 +``` -Both recipes enable fused mHC and fused rope (`use_fused_mhc=True`, `apply_rope_fusion=True`), matching the pretrain recipes. The historical "fused-kernel SFT NaN" reports are both resolved: the fused-mHC report was a confound, and the fused-rope NaN was a bridge config-mapping bug (`partial_rotary_factor` double-applied) fixed by `rotary_percent=1.0` in #4271 — with that fix, full-model (43-layer, real weights) SFT with rope fusion on 8×GB300 matches the unfused control's loss trajectory. The fused mHC cuTile kernel is **sm_100 (Blackwell)**; on Hopper set `use_fused_mhc=False`. The `no_mtp` variant drops the MTP layer and trims `csa_compress_ratios` back to `num_layers` (the bridge appends an MTP-layer ratio that `transformer_config` would otherwise reject). +The historical `128gpu` name is retained for compatibility. The recipe owns its +TP1/PP4/VPP4/EP16/CP1 topology, global batch size, natural-routing semantics, +recompute, offload, and checkpoint settings. Its current qualification is at 64 +GPUs; larger data-parallel scales require separate performance validation. Keep +the recipe defaults and override only the dataset, run length, logging, pretrained +checkpoint, and output paths. The latest 100-step real-data candidate and its +checkpoint evidence are recorded in the verification card. -> There are intentionally **no MXFP8 or Muon SFT recipes**: both were prototyped (mirroring the pretrain recipes) but fail in full-model DSv4-Flash SFT — see Blockers. SFT ships **Adam/bf16**; the pretrain MXFP8/Muon recipes are unaffected. +Compatibility aliases such as `deepseek_v4_flash_pretrain_mxfp8_config` remain +exported, but new launches should use the hardware-qualified names above. +Canonical benchmark recipes under `src/megatron/bridge/perf_recipes/` are +performance references, not substitutes for the library recipes. -`slurm_sft.sh` selects the recipe from `MTP` (`on`|`off`); `HARDWARE` (`blackwell`|`hopper`) sets the node topology (on Hopper, also set `use_fused_mhc=False` — the fused mHC kernel is Blackwell-only). It runs the model at `TP=1, PP=4, EP=8` (32 GPUs), and: +## Supervised Fine-Tuning -1. Imports `deepseek-ai/DeepSeek-V4-Flash` into a Megatron checkpoint (skipped if already present). -2. Runs full SFT from that checkpoint via `scripts/training/run_recipe.py`. +DeepSeek-V4-Flash provides BF16 Adam full-parameter SFT recipes: -GPUs per node differ by hardware, so 32 GPUs means a different node count: +| Recipe | Sequence format | MTP | Target | +|--------|-----------------|-----|--------| +| `deepseek_v4_flash_sft_config` | Unpacked SBHD | On | Hopper or Blackwell | +| `deepseek_v4_flash_no_mtp_sft_config` | Unpacked SBHD | Off | Hopper or Blackwell | +| `deepseek_v4_flash_sft_openmath_thinking_packed_config` | Offline-packed THD | On | Portable base | +| `deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config` | Offline-packed THD | On | 32-GPU GB200 | -| Hardware | GPUs/node | Nodes for 32 GPUs | `#SBATCH` | -|----------|-----------|-------------------|-----------| -| GB200 NVL | 4 | 8 | `--nodes=8 --gpus-per-node=4` (default) | -| H100/H200 | 8 | 4 | `--nodes=4 --gpus-per-node=8` | +The recipes select fused mHC only when the runtime supports the Blackwell +kernel; Hopper uses the unfused fallback. The GB200 packed recipe additionally +enables HybridEP, uneven-dispatch padding, DSA fusion, grouped GEMM, selective +recompute, and attention activation offload. -**Sequences are unpacked (SBHD).** The CSA/DSA indexer asserts `packed_seq_params is None` (`csa.py`), so packed/THD sequences are not yet supported on the sparse layers. The recipes ship an unpacked SQuAD config; do not set `dataset.enable_offline_packing=true`. Select another built-in source with `--dataset gsm8k` or use `--dataset local-jsonl dataset.dataset_root=` for JSONL data — both stay unpacked by default. +Launch the verified GB200 packed recipe from an imported BF16 checkpoint: -**Eval sizing.** Each evaluation draws `validation.eval_iters × global_batch_size` samples; if that exceeds your validation/test split the eval hangs trying to form a batch. `slurm_sft.sh` defaults to a small `EVAL_ITERS=2` and `DO_TEST=false` (the end-of-run test eval is the usual culprit on small test sets) — raise them only when your splits are large enough. +```bash +./scripts/training/train.sh --nodes 8 --gpus-per-node 4 \ + --recipe deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config \ + --mode sft --step-func dsv4_step \ + --pretrained_checkpoint work/models/deepseek-v4-flash \ + --save_dir work/results/deepseek-v4-flash-sft \ + --save_interval 100 --max_steps 100 --seq_length 1024 \ + dist.distributed_timeout_minutes=180 +``` -## SFT Status, TODO & Blockers +The card records 100 finite steps, a fresh-process checkpoint reload, and +post-SFT GPU export with deterministic HF inference for this configuration at +CP=1. It also records a separate 100-step CP=2 packed-SFT validation at +sequence length 1024; longer sequence lengths remain unverified. MXFP8 and +Muon SFT recipes are intentionally not shipped because full-model tests did +not establish a stable supported configuration. -Validated end to end on **8× GB300 (32 GPU, TP1/PP4/EP8)** with real DeepSeek-V4-Flash weights: HF→Megatron import (FP8/MXFP4→bf16) + full SFT, **MTP on and off**, `lm loss` decreasing with no NaN — at SBHD / bf16 / Adam / 4K. (Those end-to-end runs used `use_fused_mhc=False`; the recipes now default `use_fused_mhc=True`, verified clean by an isolated GB300 re-run and confirmed by a reviewer on GB200.) The items below are **not** implemented and are gated on upstream Megatron-Core (verified against the code and PRs as of 2026-06-02; tracking: [NVIDIA/Megatron-LM#4468](https://github.com/NVIDIA/Megatron-LM/issues/4468)): +## Parameter-Efficient Fine-Tuning -| Capability | Status | Gating | Notes | -|------------|--------|--------|-------| -| Packed sequence (**THD**) for DSv4 attention | **Experimental** (unmerged) | mcore PR #5011 (open) | THD *is implemented* in #5011 — it removes the `packed_seq_params is None` asserts in `csa.py`/`dsa.py` and adds the THD index/topk + `cu_seqlens` path. It **grafts cleanly onto our pinned mcore** (0-conflict 3-way merge of #5011) and needs `nvidia-cudnn-frontend[cutedsl]>=1.24.0` installed **`--no-deps`** for the DSA THD kernels (the stock 26.04 container ships an older FE; a full install shadows the container CUDA and breaks TE). The shipped/supported path stays **SBHD**, but an experimental THD (packed-seq) SFT is now **self-validated on 8×GB300** (10-iter, `lm loss` 4.46→1.48, `mtp_1` 13.5→2.70, no NaN) — ahead of #5011 merging. Merged dispatcher-THD #4816 is MoE-side only and does not lift the attention block. | -| CUDA Graphs for DSv4 THD | TODO | (no PR) | Follows THD. | -| Context parallel / long-context (≥64K) | TODO | draft PR #5087 (depends on #5011) | SFT runs CP=1; this is the Phase-3 long-context target. | -| MXFP8 / Muon **SFT** | **Fails (upstream); no SFT recipe shipped** | fp8 numerics; Muon + expert-parallel | Both prototyped (mirroring the pretrain recipes) and tested at full scale on 8×GB300 with `use_fused_mhc=False`: **MXFP8 NaNs at iter-2** (fp8 × hash-MoE/ClampedSwiGLU numerics) and **Muon hits an iter-2 `AssertionError`** (Muon + EP-MoE grad bookkeeping not yet supported upstream). Removed from the shipped recipe set; the **pretrain** MXFP8/Muon recipes remain. SFT ships **Adam/bf16**. | -| Full SFT on **H100-80GB** | needs **≥64 GPU** | hardware | At 32 GPU (TP1/PP4/EP8 ⇒ DP=1) the fp32 master-param buffer can't shard and OOMs. Use **≥64 H100** (PP8 ⇒ DP≥2), or H200-140GB / Blackwell, or PEFT/LoRA. | +DeepSeek-V4-Flash provides packed OpenMath LoRA recipes that preserve the SFT +data and objective contract: -Already incorporated in the pinned mcore (no action): dense-loss + per-layer rope-type fix (#5018), CSA/HCA (#4458), Hash MoE/ClampedSwiGLU (#4481), MTP+mHC (#4518), fusion kernels (#4894). +| Recipe | Target | +|--------|--------| +| `deepseek_v4_flash_peft_openmath_thinking_packed_config` | Portable base | +| `deepseek_v4_flash_peft_openmath_thinking_packed_gb200_config` | 32-GPU GB200 | -## Container Image +The LoRA rank and alpha are both 32. Adapters cover the MLA query down/up, KV, +and output projections plus shared and routed expert FC1/FC2 projections. +Routed experts use separate adapters (`share_expert_adapters=False`) rather +than sharing one adapter across the experts local to an EP rank. -Run inside a container that has the DSv4 prerequisites: Megatron-Bridge on a **`main2dev`** Megatron-LM commit (validated on `ed6b1f65502aec7f2fe27e14a1245c29e435c2a6`; has both `safe_get_world_size` and the DSv4 `csa.py`/`dsa.py`), the DSA dependency `fast_hadamard_transform`, and a pre-built `helpers_cpp`. The [NeMo Framework container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/nemo) — or an image built from `docker/Dockerfile.ci` with the submodule checked out to a `main2dev` commit — provides this. Set `slurm_sft.sh`'s `CONTAINER_IMAGE` to the resulting `.sqsh`. +Launch the GB200 recipe from the same imported BF16 checkpoint and packed data +used by SFT: -The first import downloads ~285B of HF weights, so point `HF_HOME` at shared scratch (e.g. `/home/scratch./HF_HOME`) before submitting so the download persists across jobs, and set `HF_TOKEN` if the repo is gated. +```bash +./scripts/training/train.sh --nodes 8 --gpus-per-node 4 \ + --recipe deepseek_v4_flash_peft_openmath_thinking_packed_gb200_config \ + --mode lora --step-func dsv4_step \ + --pretrained_checkpoint work/models/deepseek-v4-flash \ + --save_dir work/results/deepseek-v4-flash-peft \ + --save_interval 100 --max_steps 100 --seq_length 1024 \ + dist.distributed_timeout_minutes=180 +``` -### Disk footprint +The verification card records the current loss, checkpoint, and performance +status for this exact recipe. -The full bf16 model is ~570 GB, so plan storage before submitting: +## Legacy Slurm Templates -| Artifact | Size | Notes | -|----------|------|-------| -| HF download cache (`HF_HOME`) | ~150–200 GB | quantized (FP8/MXFP4) | -| Imported Megatron checkpoint (loaded by SFT) | ~570 GB | bf16, `${MEGATRON_DIR}/iter_0000000` | -| Each saved SFT checkpoint | ~570 GB | bf16 model only | +[`slurm_pretrain.sh`](slurm_pretrain.sh) and +[`slurm_sft.sh`](slurm_sft.sh) are customizable 32-GPU site templates retained +for compatibility. They call the lower-level `run_recipe.py` entry point and +do not reproduce the current hardware-qualified validation commands. +`slurm_sft.sh` uses the unpacked SFT recipes; it does not launch the GB200 packed +recipe. Prefer `scripts/training/train.sh` for new runs. -`slurm_sft.sh` does **not** save checkpoints by default (`SAVE_CKPT=0`) — to get training running you only need the base weights *loaded*, not written back — so the whole run fits in ~750 GB (cache + base). Never save optimizer state: distributed Adam for a 285B model is multi-TB (`save_optim=false` is fixed in the script). For a real fine-tune, set `SAVE_CKPT=1` (saving is capped to `KEEP_CKPTS=1` latest checkpoint) and put `HF_HOME` + `MEGATRON_DIR` on shared/project storage so a small (e.g. 2 TB) personal scratch only holds the output. +## Storage -## Parallelism Configurations +DeepSeek-V4-Flash materializes approximately 570 GB of BF16 model weights. +Plan persistent storage before importing or saving checkpoints: -DSv4 currently requires **TP=1** because MLA tensor parallelism is not supported alongside the DSv4 hybrid attention path. Scale via expert and pipeline parallelism instead. +| Artifact | Approximate size | +|----------|------------------| +| Quantized Hugging Face cache | 150-200 GB | +| Imported BF16 Megatron checkpoint | 570 GB | +| Each BF16 model-only SFT checkpoint | 570 GB | +| Each exported BF16 HF checkpoint | 570 GB | -| Model | TP | PP | EP | GPUs | GPU | Verified | -|-------|---:|---:|---:|-----:|-----|----------| -| DeepSeek-V4-Flash | 1 | 1 | 4 | 4 | GB200 192GB | Import, export, inference | -| DeepSeek-V4-Flash | 1 | 1 | 16 | 16 | H100 80GB | Import, export, inference | -| DeepSeek-V4-Flash-Base | 1 | 1 | 4 | 4 | GB200 192GB | Import, export, inference | -| DeepSeek-V4-Pro | 1 | 4 | 8 | 32 | GB200 192GB | Import, export, inference | -| DeepSeek-V4-Pro | 1 | 16 | 8 | 128 | H100 80GB | Import, export, inference | +Optimizer state can add multiple terabytes. Disable optimizer-state saves only +when the workflow does not require resumable training; otherwise provision the +required storage and validate checkpoint resume explicitly. ## Known Limitations -- **MTP is disabled for inference** via `disable_mtp_for_inference()`. MTP weights are mapped end-to-end and loaded into the Megatron model. - -- **Fused mHC: use the unfused path for SFT.** The fused cuTile mHC kernel needs sm_100 (not H100) and is on by default in the bridge config; it works for import/inference on Blackwell, but **NaNs in SFT training** (see SFT Blockers). The SFT recipes therefore force `use_fused_mhc=False` — the validated unfused/reference path, which also runs on Hopper. - -- **`fast_hadamard_transform` is required by the DSA attention variant.** `csa.py` and `dsa.py` import `hadamard_transform` from this package and hard-assert availability — there is no in-tree PyTorch fallback. It is a required Megatron Bridge dependency, pinned to the complete Dao-AILab git source because the PyPI source distribution is incomplete, and is installed by the `uv sync` commands above. - -- **Logit parity is verified for Flash and Flash-Base** against the official inference stack at last-real-token logits. The remaining gap is structural, from different attention/HC kernel decompositions and accumulation precisions between MCore and official inference. +- DeepSeek V4 currently uses TP=1 with hybrid attention; scale with PP, EP, and + DP. Use only recipe-qualified topology changes. +- Standard Megatron KV-cache autoregressive inference is unsupported because + the hybrid-attention path does not accept an inference context. The + [`inference.sh`](inference.sh) example uses legacy full-prefix generation and + is not a verified KV-cache path. HF-native inference is the target after a + verified Megatron-to-HF export. +- Fused mHC requires Blackwell (`sm_100`). Hopper uses the unfused path. +- CPU import and export require enough RAM for the full BF16 model plus + conversion workspace and remain unverified in the current card. diff --git a/scripts/inference/setup_inference.py b/scripts/inference/setup_inference.py index 1a61ac6b9c..ba1f1ae41d 100755 --- a/scripts/inference/setup_inference.py +++ b/scripts/inference/setup_inference.py @@ -75,7 +75,12 @@ def _build_parser() -> argparse.ArgumentParser: "--gpus_per_node", type=int, dest="gpus_per_node", - help="GPUs and inference tasks per node.", + help="GPUs allocated per node.", + ) + execution.add_argument( + "--tasks-per-node", + type=int, + help="Inference tasks per node; defaults to one task per GPU.", ) execution.add_argument("--cpus-per-task", type=int, help="CPUs allocated to each inference task.") execution.add_argument("--mem", help="Optional Slurm memory request, such as 64G.") @@ -170,6 +175,8 @@ def _validate_args(args: argparse.Namespace) -> None: raise ValueError("--nodes must be at least 1.") if args.gpus_per_node is None or args.gpus_per_node < 1: raise ValueError("--gpus-per-node is required and must be at least 1.") + if args.tasks_per_node is not None and not 1 <= args.tasks_per_node <= args.gpus_per_node: + raise ValueError("--tasks-per-node must be between 1 and --gpus-per-node.") if args.cpus_per_task is not None and args.cpus_per_task < 1: raise ValueError("--cpus-per-task must be at least 1.") if any(not value.strip() for value in args.srun_args): @@ -201,7 +208,7 @@ def _build_executor(args: argparse.Namespace, env_names: list[str], mounts: list account=args.account, partition=args.partition, nodes=args.nodes, - ntasks_per_node=args.gpus_per_node, + ntasks_per_node=args.tasks_per_node or args.gpus_per_node, cpus_per_task=args.cpus_per_task, mem=args.mem, exclusive=True if args.exclusive else None, diff --git a/scripts/training/run_recipe.py b/scripts/training/run_recipe.py index e3bd313e2a..559f86de22 100755 --- a/scripts/training/run_recipe.py +++ b/scripts/training/run_recipe.py @@ -132,6 +132,9 @@ TrainMode = Literal["pretrain", "finetune"] +IMPORT_TIME_RECIPE_ENV_VARS = frozenset({"NVTE_CPU_OFFLOAD_V1"}) + + COMMON_OVERRIDE_FIELDS = ( ("max_steps", "train.train_iters"), ("global_batch_size", "train.global_batch_size"), @@ -420,6 +423,12 @@ def _selected_recipe_name(args: argparse.Namespace) -> str: return args.recipe or f"{args.model}_{recipe_task(args.mode)}_config" +def _requires_recipe_environment_bootstrap(recipe: ConfigContainer) -> bool: + """Return whether a recipe declares an unset import-time environment variable.""" + recipe_env = getattr(recipe, "env_vars", None) or {} + return any(name in recipe_env and name not in os.environ for name in IMPORT_TIME_RECIPE_ENV_VARS) + + def _load_selected_recipe(args: argparse.Namespace) -> ConfigContainer: """Load the requested recipe by its complete name or model-derived library name.""" peft_scheme = args.mode if args.mode in {"lora", "dora"} else None @@ -496,17 +505,19 @@ def main(argv: list[str] | None = None) -> None: world_size=benchmark_world_size, ) configuration_mode = _train_mode(args.mode) - - if benchmark_metadata is not None: + if benchmark_metadata is not None or _requires_recipe_environment_bootstrap(recipe): recipe = bootstrap_recipe_environment( recipe, script_path=str(Path(__file__).resolve()), argv=list(argv) if argv is not None else sys.argv[1:], ) + else: + recipe = apply_runtime_environment(recipe) + + if benchmark_metadata is not None: 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/create-model-verification-card/scripts/verify_hf_inference.py b/skills/create-model-verification-card/scripts/verify_hf_inference.py index d462eb47b6..40fd88f2a6 100644 --- a/skills/create-model-verification-card/scripts/verify_hf_inference.py +++ b/skills/create-model-verification-card/scripts/verify_hf_inference.py @@ -49,12 +49,27 @@ def _parse_args() -> argparse.Namespace: help="Pass enable_thinking=False to the tokenizer chat template.", ) parser.add_argument("--device", default="cuda", help="Torch device used for inference.") + parser.add_argument( + "--device-map", + choices=("auto", "balanced", "balanced_low_0", "sequential"), + help="Optional Hugging Face device-map strategy for sharded model loading.", + ) parser.add_argument( "--dtype", choices=("bfloat16", "float16", "float32"), default="bfloat16", help="Model loading dtype.", ) + parser.add_argument( + "--autocast", + action="store_true", + help="Run generation under autocast using the model loading dtype.", + ) + parser.add_argument( + "--require-gpu-only", + action="store_true", + help="Fail if any model shard is placed on CPU or disk.", + ) args = parser.parse_args() if args.max_new_tokens <= 0: parser.error("--max-new-tokens must be positive") @@ -122,6 +137,23 @@ def _validate_loading_info(loading_info: dict[str, Any]) -> None: raise RuntimeError(f"Exported Hugging Face checkpoint did not reload strictly: {details}") +def _validate_gpu_only_placement(model: Any) -> None: + """Require every model shard to be resident on CUDA devices.""" + device_map = getattr(model, "hf_device_map", None) + if device_map is None: + if getattr(model.device, "type", str(model.device).split(":", 1)[0]) != "cuda": + raise RuntimeError(f"Model is not GPU-only: device={model.device}") + return + + non_gpu = { + name: str(device) + for name, device in device_map.items() + if not isinstance(device, int) and str(device).split(":", 1)[0] != "cuda" + } + if non_gpu: + raise RuntimeError(f"Model has non-GPU placements: {non_gpu}") + + def _load_runtime(args: argparse.Namespace) -> tuple[Any, Any, Any]: """Load torch, the selected HF auto-model, and its tokenizer or processor.""" import torch @@ -137,27 +169,53 @@ def _load_runtime(args: argparse.Namespace) -> tuple[Any, Any, Any]: processor = AutoTokenizer.from_pretrained(args.hf_model, trust_remote_code=args.trust_remote_code) model_cls = AutoModelForCausalLM - model, loading_info = model_cls.from_pretrained( - args.hf_model, - dtype=dtype, - trust_remote_code=args.trust_remote_code, - output_loading_info=True, - ) + model_kwargs = { + "dtype": dtype, + "trust_remote_code": args.trust_remote_code, + "output_loading_info": True, + } + if args.device_map: + model_kwargs["device_map"] = args.device_map + model, loading_info = model_cls.from_pretrained(args.hf_model, **model_kwargs) _validate_loading_info(loading_info) - model = model.to(args.device).eval() + if not args.device_map: + model = model.to(args.device) + model = model.eval() + if args.require_gpu_only: + _validate_gpu_only_placement(model) + LOG.info("Strict HF reload complete (%d model modules)", sum(1 for _ in model.modules())) return torch, model, processor +def _model_input_device(model: Any) -> Any: + """Return the device holding the model's input embeddings.""" + get_input_embeddings = getattr(model, "get_input_embeddings", None) + if get_input_embeddings is not None: + embeddings = get_input_embeddings() + weight = getattr(embeddings, "weight", None) + if weight is not None and getattr(weight.device, "type", None) != "meta": + return weight.device + return model.device + + def main() -> int: """Run one bounded greedy generation and print its completion.""" args = _parse_args() torch, model, processor = _load_runtime(args) tokenizer = getattr(processor, "tokenizer", processor) - inputs = _prepare_inputs(processor, args).to(model.device) + input_device = _model_input_device(model) + inputs = _prepare_inputs(processor, args).to(input_device) prompt_length = inputs["input_ids"].shape[1] pad_token_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id - with torch.inference_mode(): + with ( + torch.inference_mode(), + torch.autocast( + device_type=torch.device(input_device).type, + dtype=getattr(torch, args.dtype), + enabled=args.autocast, + ), + ): output = model.generate( **inputs, do_sample=False, diff --git a/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md b/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md index 5b58640d63..dd2352f5c8 100644 --- a/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md +++ b/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md @@ -74,6 +74,10 @@ 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_gb200_fp8mx_library_config` | Pretrain | 1 | 4 | 16 | 64 (GB200; VPP4; natural unlimited capacity) | +| `deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config` | SFT | 1 | 4 | 8 | 32 (GB200; offline packed) | +| `deepseek_v4_flash_peft_openmath_thinking_packed_config` | PEFT | 1 | 4 | 8 | 32 (offline packed) | +| `deepseek_v4_flash_peft_openmath_thinking_packed_gb200_config` | PEFT | 1 | 4 | 8 | 32 (GB200; offline packed) | ### GLM-4.5 diff --git a/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py b/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py index cdccfa4a44..598135a201 100644 --- a/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py +++ b/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py @@ -944,11 +944,10 @@ def maybe_modify_converted_hf_weight( ) -> Dict[str, torch.Tensor]: """Recreate DSv4 quantized weight/scale pairs expected by the source shard index. + Legacy indexer scorer names are restored before selecting the export dtype. When ``task.weight_dtype`` is set, skip requantization and return the weights unchanged — the generic export path casts the dtype. """ - if task.weight_dtype is not None: - return converted_weights_dict native_scorer_key = next( ( key @@ -962,6 +961,8 @@ def maybe_modify_converted_hf_weight( if legacy_key in hf_state_dict: converted_weights_dict = dict(converted_weights_dict) converted_weights_dict[legacy_key] = converted_weights_dict.pop(native_scorer_key) + if task.weight_dtype is not None: + return converted_weights_dict return quantization_utils.requantize_hf_weight_scale_pairs( converted_weights_dict, hf_state_dict, diff --git a/src/megatron/bridge/recipes/deepseek/__init__.py b/src/megatron/bridge/recipes/deepseek/__init__.py index 1cf2ac77dd..b22e03b3be 100644 --- a/src/megatron/bridge/recipes/deepseek/__init__.py +++ b/src/megatron/bridge/recipes/deepseek/__init__.py @@ -34,6 +34,7 @@ # DeepSeek V4 from .deepseek_v4 import ( deepseek_v4_flash_no_mtp_sft_config, + deepseek_v4_flash_peft_openmath_thinking_packed_config, deepseek_v4_flash_pretrain_config, deepseek_v4_flash_pretrain_gb200_config, deepseek_v4_flash_pretrain_muon_config, @@ -47,9 +48,12 @@ set_deepseek_v4_pipeline_model_parallel_layout, ) from .gb200.deepseek_v4 import ( + deepseek_v4_flash_peft_openmath_thinking_packed_gb200_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_64gpu_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, @@ -66,6 +70,8 @@ "deepseek_v3_pretrain_config_32nodes", "set_deepseek_v3_pipeline_model_parallel_layout", # DeepSeek V4 + "deepseek_v4_flash_peft_openmath_thinking_packed_config", + "deepseek_v4_flash_peft_openmath_thinking_packed_gb200_config", "deepseek_v4_flash_pretrain_config", "deepseek_v4_flash_pretrain_gb200_config", "deepseek_v4_flash_pretrain_muon_config", @@ -80,6 +86,8 @@ "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_64gpu_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", "set_deepseek_v4_pipeline_model_parallel_layout", diff --git a/src/megatron/bridge/recipes/deepseek/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/deepseek_v4.py index 5ad648049b..fdbbb42d43 100644 --- a/src/megatron/bridge/recipes/deepseek/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/deepseek_v4.py @@ -12,13 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. # ruff: noqa: F401 -"""Compatibility aliases for legacy recipe names.""" +"""Hardware-agnostic DeepSeek V4 recipes and compatibility aliases.""" from __future__ import annotations from megatron.bridge.models.deepseek.deepseek_v4_bridge import ( set_deepseek_v4_pipeline_model_parallel_layout, ) +from megatron.bridge.peft.lora import LoRA from megatron.bridge.recipes.deepseek.gb200.deepseek_v4 import ( deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config, deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config, @@ -66,6 +67,7 @@ __all__ = [ "deepseek_v4_flash_no_mtp_sft_config", + "deepseek_v4_flash_peft_openmath_thinking_packed_config", "deepseek_v4_flash_pretrain_config", "deepseek_v4_flash_pretrain_muon_config", "deepseek_v4_flash_pretrain_mxfp8_config", @@ -85,6 +87,34 @@ ] +_DEEPSEEK_V4_LORA_TARGET_MODULES = [ + "linear_q_down_proj", + "linear_q_up_proj", + "linear_kv_proj", + "linear_proj", + "linear_fc1", + "linear_fc2", +] + + +def _apply_deepseek_v4_lora(cfg: ConfigContainer) -> None: + """Apply the DeepSeek V4 LoRA convergence contract to an SFT config.""" + cfg.peft = LoRA( + target_modules=list(_DEEPSEEK_V4_LORA_TARGET_MODULES), + dim=32, + alpha=32, + dropout=0.0, + share_expert_adapters=False, + ) + cfg.optimizer.lr = 1.0e-4 + cfg.optimizer.min_lr = 0.0 + # MCore's unified recompute path cannot replay the frozen grouped-expert graph. + cfg.model.recompute_granularity = None + cfg.model.recompute_modules = None + cfg.model.recompute_method = None + cfg.model.recompute_num_layers = None + + def deepseek_v4_flash_sft_openmath_thinking_packed_config() -> ConfigContainer: """DSv4 Flash SFT on OpenMathInstruct-2 with thinking channel and offline-packed sequences. @@ -93,6 +123,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; @@ -104,3 +137,15 @@ def deepseek_v4_flash_sft_openmath_thinking_packed_config() -> ConfigContainer: pad_seq_to_mult=2 * cfg.model.context_parallel_size, ) return cfg + + +def deepseek_v4_flash_peft_openmath_thinking_packed_config() -> ConfigContainer: + """DSv4 Flash LoRA on packed OpenMathInstruct-2 thinking data. + + The attention targets follow the DeepSeek MLA projection layout. Both shared + and routed MLP projections are adapted; grouped routed experts use one adapter + per local expert to match the verl training layout. + """ + cfg = deepseek_v4_flash_sft_openmath_thinking_packed_config() + _apply_deepseek_v4_lora(cfg) + return cfg diff --git a/src/megatron/bridge/recipes/deepseek/gb200/__init__.py b/src/megatron/bridge/recipes/deepseek/gb200/__init__.py index 696f6cae8b..8a46121edd 100644 --- a/src/megatron/bridge/recipes/deepseek/gb200/__init__.py +++ b/src/megatron/bridge/recipes/deepseek/gb200/__init__.py @@ -14,14 +14,20 @@ """GB200 recipes for DeepSeek V4 Flash.""" from megatron.bridge.recipes.deepseek.gb200.deepseek_v4 import ( + deepseek_v4_flash_peft_openmath_thinking_packed_gb200_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_64gpu_gb200_fp8mx_library_config, + deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config, ) __all__ = [ + "deepseek_v4_flash_peft_openmath_thinking_packed_gb200_config", + "deepseek_v4_flash_pretrain_64gpu_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 677dabbe01..c6adda2d41 100644 --- a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py @@ -11,10 +11,7 @@ # 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. -"""GB200 recipes for DeepSeek V4 Flash. - -Validated on OCI GB200 NVL72 (64 GPUs, PP=8, EP=8). -""" +"""GB200 pretraining, SFT, and PEFT recipes for DeepSeek V4 Flash.""" import torch @@ -36,14 +33,16 @@ from megatron.bridge.training.comm_overlap import CommOverlapConfig from megatron.bridge.training.config import ConfigContainer from megatron.bridge.training.mixed_precision import bf16_mixed, bf16_with_mxfp8_mixed +from megatron.bridge.utils.cuda_graph import set_cuda_graph_modules def deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config() -> ConfigContainer: """DeepSeek-V4-Flash BF16 pre-training on 64-GPU GB200 (PP=8, EP=8). - All features validated in 1k-step convergence runs on OCI GB200 NVL72: - fused DSA indexer, DSA indexer loss, HybridEP dispatcher, GroupedGEMM, - permute fusion, selective recompute with mla_up_proj. + The core 64-GPU topology completed 1k-step convergence runs with fused DSA + indexer and loss, HybridEP, grouped GEMM, permute fusion, and selective + recompute. The current router fusion was also exercised by the derived + library recipe's 64-GPU 100-step run. Note: set checkpoint.save_optim=False when using HybridEP. """ @@ -71,7 +70,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 @@ -81,11 +80,12 @@ def deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config() -> ConfigContainer: cfg.model.moe_token_dispatcher_type = "flex" cfg.model.moe_flex_dispatcher_backend = "hybridep" cfg.model.moe_flex_dispatcher_num_sms = 16 - cfg.model.moe_hybridep_num_sms = None 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" @@ -93,8 +93,9 @@ def deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config() -> ConfigContainer: 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 @@ -128,6 +129,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 @@ -135,6 +137,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 @@ -146,10 +149,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, @@ -189,6 +191,64 @@ def deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config() -> ConfigContainer: return cfg +def deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_library_config() -> ConfigContainer: + """Return the scalable real-training DeepSeek V4 Flash GB200 config. + + This 64-GPU variant uses PP4/VPP4, EP16, and selective recompute without + activation offload. Expert capacity, paged stash, CUDA graphs, and forced + load balancing 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 = 4 + cfg.model.virtual_pipeline_model_parallel_size = 4 + cfg.model.context_parallel_size = 1 + cfg.model.expert_model_parallel_size = 16 + cfg.model.expert_tensor_parallel_size = 1 + cfg.model.sequence_parallel = False + cfg.model.pipeline_model_parallel_layout = "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" + 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_flex_dispatcher_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 = None + cfg.model.fine_grained_offloading_max_inflight_offloads = None + cfg.model.moe_pad_experts_for_cuda_graph_inference = False + cfg.model.cuda_graph_impl = "none" + set_cuda_graph_modules(cfg.model, []) + cfg.model.use_te_rng_tracker = False + cfg.rng.te_rng_tracker = False + 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": 16, + "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, + "NVLINK_DOMAIN_SIZE": 72, + "USE_MNNVL": 1, + "NVTE_CPU_OFFLOAD_V1": 0, + "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() @@ -229,3 +289,78 @@ 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, fused DSA + execution, static packed-sequence shapes, and grouped expert GEMMs to the + hardware-agnostic packed SFT recipe. + """ + 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.moe_token_dispatcher_type = "flex" + cfg.model.moe_flex_dispatcher_backend = "hybridep" + cfg.model.moe_flex_dispatcher_num_sms = 16 + cfg.model.moe_shared_expert_overlap = False + cfg.model.moe_hybridep_pad_uneven_dispatch_inputs = True + cfg.model.moe_grouped_gemm = True + cfg.model.moe_permute_fusion = True + cfg.model.moe_router_fusion = True + + cfg.model.recompute_granularity = "selective" + cfg.model.recompute_modules = ["moe", "mhc", "mla_up_proj", "layernorm"] + cfg.model.recompute_method = None + cfg.model.recompute_num_layers = None + cfg.model.calculate_per_token_loss = True + 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.dataset.offline_packing_specs.pad_seq_to_mult = 4 + cfg.dataset.offline_packing_specs.pad_cu_seqlens = True + cfg.dataset.dataset_kwargs = { + **(cfg.dataset.dataset_kwargs or {}), + "pad_to_max_length": True, + } + + cfg.model.cross_entropy_fusion_impl = "te" + cfg.env_vars = { + **cfg.env_vars, + "NVTE_CPU_OFFLOAD_V1": 1, + } + return cfg + + +def deepseek_v4_flash_peft_openmath_thinking_packed_gb200_config() -> ConfigContainer: + """Return the GB200-optimized packed OpenMath LoRA config.""" + from megatron.bridge.recipes.deepseek.deepseek_v4 import _apply_deepseek_v4_lora + + cfg = deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config() + _apply_deepseek_v4_lora(cfg) + # PEFT's smaller training-state footprint also lets this graph avoid activation offload. + cfg.model.fine_grained_activation_offloading = False + cfg.model.offload_modules = None + cfg.env_vars.pop("NVTE_CPU_OFFLOAD_V1", None) + + cfg.model.moe_flex_dispatcher_num_sms = 32 + cfg.model.cuda_graph_impl = "transformer_engine" + set_cuda_graph_modules(cfg.model, ["moe_router", "moe_preprocess"]) + cfg.model.cuda_graph_warmup_steps = 3 + cfg.model.use_te_rng_tracker = True + cfg.rng.te_rng_tracker = True + cfg.env_vars = { + **cfg.env_vars, + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 8, + "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, + "NVLINK_DOMAIN_SIZE": 72, + "USE_MNNVL": 1, + } + return cfg diff --git a/src/megatron/bridge/training/utils/flop_utils.py b/src/megatron/bridge/training/utils/flop_utils.py index f400664389..7e9d86ddd0 100644 --- a/src/megatron/bridge/training/utils/flop_utils.py +++ b/src/megatron/bridge/training/utils/flop_utils.py @@ -12,10 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +from collections.abc import Callable from dataclasses import dataclass from pathlib import Path +from typing import cast import torch +from megatron.core.models.gpt import experimental_attention_variant_module_specs from megatron.core.models.hybrid.hybrid_layer_allocation import ( Symbols, get_hybrid_layer_counts, @@ -30,6 +33,17 @@ _lora_seq_stats_cache: dict = {} +_mcore_is_gated_delta_net_variant = cast( + Callable[[str | None], bool] | None, + getattr(experimental_attention_variant_module_specs, "is_gated_delta_net_variant", None), +) + + +def _is_gated_delta_net_variant(experimental_attention_variant: str | None) -> bool: + """Recognize GDN variants across current main and older MCore dev branches.""" + if _mcore_is_gated_delta_net_variant is not None: + return _mcore_is_gated_delta_net_variant(experimental_attention_variant) + return experimental_attention_variant in {"gated_delta_net", "gdn", "gdn2"} @dataclass(frozen=True) @@ -1406,7 +1420,7 @@ def count_indexer_layers(layer_count: int) -> int: # Handle GDN (Gated DeltaNet) hybrid attention variants. MCore normalizes # the deprecated "gated_delta_net" alias to "gdn" during config finalization. - if experimental_attention_variant in {"gated_delta_net", "gdn", "gdn2"}: + if _is_gated_delta_net_variant(experimental_attention_variant): linear_attention_freq = cfg.model.linear_attention_freq decoder_num_layers = cfg.model.num_layers if linear_attention_freq is None: diff --git a/tests/functional_tests/test_groups/models/deepseek/test_deepseek_v4_conversion.py b/tests/functional_tests/test_groups/models/deepseek/test_deepseek_v4_conversion.py index f9d0a39213..adb5023637 100644 --- a/tests/functional_tests/test_groups/models/deepseek/test_deepseek_v4_conversion.py +++ b/tests/functional_tests/test_groups/models/deepseek/test_deepseek_v4_conversion.py @@ -225,6 +225,7 @@ def deepseek_v4_toy_model_path(self, tmp_path_factory): def test_deepseek_v4_roundtrip_ep(self, deepseek_v4_toy_model_path, tmp_path): test_output_dir = tmp_path / "deepseek_v4_ep" test_output_dir.mkdir(exist_ok=True) + megatron_checkpoint_dir = test_output_dir / "megatron" cmd = [ "python", @@ -249,6 +250,8 @@ def test_deepseek_v4_roundtrip_ep(self, deepseek_v4_toy_model_path, tmp_path): "1", "--ep", "2", + "--megatron-save-path", + str(megatron_checkpoint_dir), ] result = subprocess.run( @@ -260,6 +263,21 @@ def test_deepseek_v4_roundtrip_ep(self, deepseek_v4_toy_model_path, tmp_path): print(f"STDERR: {result.stderr}") assert result.returncode == 0, f"DeepSeek V4 conversion failed with {result.returncode}" + reload_cmd = [ + *cmd, + "--megatron-load-path", + str(megatron_checkpoint_dir), + "--skip-save", + ] + reload_result = subprocess.run( + reload_cmd, capture_output=True, text=True, cwd=Path(__file__).parent.parent.parent.parent.parent.parent + ) + + assert reload_result.returncode == 0, ( + f"DeepSeek V4 checkpoint reload conversion failed with {reload_result.returncode}\n" + f"STDOUT:\n{reload_result.stdout}\nSTDERR:\n{reload_result.stderr}" + ) + converted_dir = test_output_dir / Path(deepseek_v4_toy_model_path).name assert (converted_dir / "config.json").exists() assert list(converted_dir.glob("*.safetensors")) diff --git a/tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py b/tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py index f2223f9c1d..06870cceb2 100644 --- a/tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py +++ b/tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py @@ -56,10 +56,10 @@ def _by_megatron(registry): return {m.megatron_param: m for m in registry.mappings} -def _dummy_task(): +def _dummy_task(*, weight_dtype=None): from megatron.bridge.models.conversion.model_bridge import WeightConversionTask - return WeightConversionTask(param_name="", global_param_name="", mapping=None) + return WeightConversionTask(param_name="", global_param_name="", mapping=None, weight_dtype=weight_dtype) def _deepseek_v4_hf_config(): @@ -151,13 +151,16 @@ def test_import_accepts_legacy_flat_indexer_weight_name(self): assert result is weight - def test_export_uses_legacy_flat_indexer_weight_name_when_source_requires_it(self): + @pytest.mark.parametrize("weight_dtype", [None, torch.bfloat16]) + def test_export_uses_legacy_flat_indexer_weight_name_when_source_requires_it(self, weight_dtype): bridge = DeepSeekV4Bridge() scorer_name = "layers.1.attn.indexer.scorer.weights_proj.weight" flat_name = "layers.1.attn.indexer.weights_proj.weight" weight = torch.randn(4, 4, dtype=torch.bfloat16) - result = bridge.maybe_modify_converted_hf_weight(_dummy_task(), {scorer_name: weight}, {flat_name: weight}) + result = bridge.maybe_modify_converted_hf_weight( + _dummy_task(weight_dtype=weight_dtype), {scorer_name: weight}, {flat_name: weight} + ) assert scorer_name not in result assert result[flat_name] is weight diff --git a/tests/unit_tests/recipes/test_deepseek_recipes.py b/tests/unit_tests/recipes/test_deepseek_recipes.py index 6981105be0..0d571a2a34 100644 --- a/tests/unit_tests/recipes/test_deepseek_recipes.py +++ b/tests/unit_tests/recipes/test_deepseek_recipes.py @@ -46,6 +46,7 @@ { "deepseek_v3_pretrain_config", "deepseek_v3_pretrain_config_32nodes", + "deepseek_v4_flash_peft_openmath_thinking_packed_config", "deepseek_v4_flash_pretrain_config", "deepseek_v4_flash_pretrain_mxfp8_config", "deepseek_v4_flash_pretrain_muon_config", @@ -385,6 +386,15 @@ def _build_deepseek_v4_recipe(name: str, monkeypatch: pytest.MonkeyPatch): return getattr(mod, name)() +def test_deepseek_v4_portable_peft_recipe_disables_recompute(monkeypatch: pytest.MonkeyPatch): + cfg = _build_deepseek_v4_recipe("deepseek_v4_flash_peft_openmath_thinking_packed_config", monkeypatch) + + assert cfg.model.recompute_granularity is None + assert cfg.model.recompute_modules is None + assert cfg.model.recompute_method is None + assert cfg.model.recompute_num_layers is None + + def test_deepseek_v4_adam_mxfp8_recipe_uses_validated_optimizer_defaults(monkeypatch: pytest.MonkeyPatch): cfg = _build_deepseek_v4_recipe("deepseek_v4_flash_pretrain_mxfp8_config", monkeypatch) @@ -482,7 +492,7 @@ def test_deepseek_v4_recipes_keep_hardware_qualified_dispatcher( if expected_dispatcher == "flex": assert cfg.model.moe_flex_dispatcher_backend == "hybridep" assert cfg.model.moe_flex_dispatcher_num_sms == 16 - assert cfg.model.moe_hybridep_num_sms is None + assert getattr(cfg.model, "moe_hybridep_num_sms", None) is None @pytest.mark.parametrize( 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..0e2d287a17 --- /dev/null +++ b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py @@ -0,0 +1,244 @@ +# 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 + +import megatron.bridge.recipes as recipes +from megatron.bridge.recipes.deepseek.gb200.deepseek_v4 import ( + deepseek_v4_flash_peft_openmath_thinking_packed_gb200_config as flash_packed_peft_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_64gpu_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_fp8mx_config as pro_fp8_config, +) +from megatron.bridge.training.config import ConfigContainer +from megatron.bridge.utils.cuda_graph import cuda_graph_module_names +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) + + +def test_flash_base_recipe_enables_precision_independent_fusions() -> None: + cfg = flash_bf16_base_config() + + 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 == "te" + 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 + + +@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 == 4 + assert cfg.dataset.offline_packing_specs.pad_cu_seqlens is True + assert cfg.dataset.dataset_kwargs == {"pad_to_max_length": True} + assert cfg.model.apply_dsa_kernel_fusion is True + assert cfg.model.dsa_indexer_loss_coeff == 0.0 + assert cfg.model.dsa_indexer_use_sparse_loss is False + assert cfg.model.moe_token_dispatcher_type == "flex" + assert cfg.model.moe_flex_dispatcher_backend == "hybridep" + assert cfg.model.moe_flex_dispatcher_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 True + assert cfg.model.cross_entropy_fusion_impl == "te" + assert cfg.model.recompute_granularity == "selective" + assert cfg.model.recompute_modules == ["moe", "mhc", "mla_up_proj", "layernorm"] + assert cfg.model.recompute_method is None + assert cfg.model.recompute_num_layers is None + assert cfg.model.calculate_per_token_loss is True + 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_peft_recipe_adapts_mla_and_unshared_experts() -> None: + cfg = flash_packed_peft_config() + + assert cfg.peft.target_modules == [ + "linear_q_down_proj", + "linear_q_up_proj", + "linear_kv_proj", + "linear_proj", + "linear_fc1", + "linear_fc2", + ] + assert cfg.peft.dim == 32 + assert cfg.peft.alpha == 32 + assert cfg.peft.dropout == 0.0 + assert cfg.peft.share_expert_adapters is False + assert cfg.optimizer.lr == 1.0e-4 + assert cfg.optimizer.min_lr == 0.0 + assert cfg.model.recompute_granularity is None + assert cfg.model.recompute_modules is None + assert cfg.model.recompute_method is None + assert cfg.model.recompute_num_layers is None + assert cfg.model.fine_grained_activation_offloading is False + assert cfg.model.offload_modules is None + assert cfg.model.moe_flex_dispatcher_num_sms == 32 + assert cfg.model.cuda_graph_impl == "transformer_engine" + assert cuda_graph_module_names(cfg.model) == ["moe_router", "moe_preprocess"] + assert cfg.model.cuda_graph_warmup_steps == 3 + assert cfg.model.use_te_rng_tracker is True + assert cfg.rng.te_rng_tracker is True + assert cfg.env_vars["NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN"] == 8 + assert cfg.env_vars["NUM_OF_TOKENS_PER_CHUNK_COMBINE_API"] == 128 + assert cfg.env_vars["NVLINK_DOMAIN_SIZE"] == 72 + assert cfg.env_vars["USE_MNNVL"] == 1 + assert "NVTE_CPU_OFFLOAD_V1" not in cfg.env_vars + + sft_cfg = flash_packed_sft_config() + assert cfg.dataset == sft_cfg.dataset + assert cfg.train == sft_cfg.train + assert cfg.rng.seed == sft_cfg.rng.seed + assert cfg.rng.inference_rng_tracker == sft_cfg.rng.inference_rng_tracker + assert cfg.rng.data_parallel_random_init == sft_cfg.rng.data_parallel_random_init + assert cfg.model.seq_length == sft_cfg.model.seq_length + assert cfg.model.pipeline_model_parallel_size == sft_cfg.model.pipeline_model_parallel_size + assert cfg.model.expert_model_parallel_size == sft_cfg.model.expert_model_parallel_size + assert cfg.model.moe_token_dispatcher_type == sft_cfg.model.moe_token_dispatcher_type + + +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 cuda_graph_module_names(cfg.model) == [] + assert cfg.model.use_te_rng_tracker is False + assert cfg.rng.te_rng_tracker is False + assert cfg.model.use_transformer_engine_op_fuser is True + assert cfg.model.moe_mlp_glu_interleave_size == 32 + assert cfg.model.moe_flex_dispatcher_num_sms == 32 + assert getattr(cfg.model, "moe_deepep_num_sms", None) is None + assert getattr(cfg.model, "moe_hybridep_num_sms", None) is None + assert cfg.checkpoint.pretrained_checkpoint is None + assert cfg.model.pipeline_model_parallel_size == 4 + assert cfg.model.virtual_pipeline_model_parallel_size == 4 + assert cfg.model.expert_model_parallel_size == 16 + assert cfg.model.pipeline_model_parallel_layout == ( + "Et*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*2|t*2|t*2|t*2|t*2mL" + ) + assert cfg.model.recompute_modules == ["mhc", "mla_up_proj"] + assert cfg.model.fine_grained_activation_offloading is False + assert cfg.model.offload_modules is None + 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 cfg.env_vars["NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN"] == 16 + assert cfg.env_vars["NVTE_CPU_OFFLOAD_V1"] == 0 + + +def test_high_scale_deepseek_v4_recipes_are_exported() -> None: + assert recipes.deepseek_v4_flash_peft_openmath_thinking_packed_gb200_config is flash_packed_peft_config + assert recipes.deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_library_config is flash_library_config + assert recipes.deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config is flash_packed_sft_config diff --git a/tests/unit_tests/scripts/inference/test_setup_inference.py b/tests/unit_tests/scripts/inference/test_setup_inference.py index 8e440f8ca6..84ee969c29 100644 --- a/tests/unit_tests/scripts/inference/test_setup_inference.py +++ b/tests/unit_tests/scripts/inference/test_setup_inference.py @@ -182,6 +182,8 @@ def test_submission_dry_run_aliases_are_consumed(submission_option): [ (("--nodes", "0"), "--nodes"), (("--gpus-per-node", "0"), "--gpus-per-node"), + (("--gpus-per-node", "4", "--tasks-per-node", "0"), "--tasks-per-node"), + (("--gpus-per-node", "4", "--tasks-per-node", "5"), "--tasks-per-node"), (("--cpus-per-task", "0"), "--cpus-per-task"), (("--srun-arg=",), "--srun-arg"), ], @@ -256,6 +258,8 @@ def __init__(self, **kwargs): _launcher_args( "--gpus-per-node", "4", + "--tasks-per-node", + "1", "--srun-arg=--mpi=pmix", "--srun-arg=--container-writable", ) @@ -263,7 +267,7 @@ def __init__(self, **kwargs): executor = module._build_executor(args, ["HF_TOKEN"], ["/host:/container"]) - assert executor.kwargs["ntasks_per_node"] == 4 + assert executor.kwargs["ntasks_per_node"] == 1 assert executor.kwargs["gpus_per_node"] == 4 assert executor.kwargs["exclusive"] is None assert "launcher" not in executor.kwargs @@ -275,6 +279,25 @@ def __init__(self, **kwargs): assert executor.env_vars == {} +def test_slurm_executor_defaults_to_one_task_per_gpu(tmp_path, monkeypatch): + module = _load_setup_inference_module() + + class _SlurmExecutor: + def __init__(self, **kwargs): + self.kwargs = kwargs + + module.run.Packager = lambda: "packager" + module.run.LocalTunnel = lambda **kwargs: types.SimpleNamespace(**kwargs) + module.run.SlurmExecutor = _SlurmExecutor + monkeypatch.setattr(module, "get_nemorun_home", lambda: str(tmp_path)) + args, _ = module.parse_args(_launcher_args("--gpus-per-node", "4")) + + executor = module._build_executor(args, [], []) + + assert executor.kwargs["ntasks_per_node"] == 4 + assert executor.kwargs["gpus_per_node"] == 4 + + def test_slurm_executor_can_request_exclusive_nodes(tmp_path, monkeypatch): module = _load_setup_inference_module() diff --git a/tests/unit_tests/scripts/training/test_run_recipe.py b/tests/unit_tests/scripts/training/test_run_recipe.py index 23aede2023..3cbe7e8e6b 100644 --- a/tests/unit_tests/scripts/training/test_run_recipe.py +++ b/tests/unit_tests/scripts/training/test_run_recipe.py @@ -476,7 +476,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_applies_runtime_environment_without_restart(): module, handles = _load_module() handles.recipe_runner.load_recipe.return_value = SimpleNamespace() @@ -489,10 +489,40 @@ def test_library_only_canonical_name_does_not_enable_benchmark_runtime(): ] ) + handles.recipe_runner.apply_runtime_environment.assert_called_once_with( + handles.recipe_runner.load_recipe.return_value + ) handles.recipe_runner.bootstrap_recipe_environment.assert_not_called() handles.recipe_runner.load_forward_step.assert_called_once_with("llm_step", mode="pretrain") +def test_library_recipe_bootstraps_unset_import_time_environment(monkeypatch): + module, handles = _load_module() + monkeypatch.delenv("NVTE_CPU_OFFLOAD_V1", raising=False) + config = SimpleNamespace(env_vars={"NVTE_CPU_OFFLOAD_V1": 1}) + handles.recipe_runner.load_recipe.return_value = config + + module.main(["--recipe", "deepseek_v4_flash_pretrain_config", "--mode", "pretrain"]) + + handles.recipe_runner.bootstrap_recipe_environment.assert_called_once() + bootstrap_call = handles.recipe_runner.bootstrap_recipe_environment.call_args + assert bootstrap_call.args == (config,) + assert bootstrap_call.kwargs["script_path"].endswith("scripts/training/run_recipe.py") + handles.recipe_runner.apply_runtime_environment.assert_not_called() + + +def test_library_recipe_preserves_explicit_import_time_environment(monkeypatch): + module, handles = _load_module() + monkeypatch.setenv("NVTE_CPU_OFFLOAD_V1", "0") + config = SimpleNamespace(env_vars={"NVTE_CPU_OFFLOAD_V1": 1}) + handles.recipe_runner.load_recipe.return_value = config + + module.main(["--recipe", "deepseek_v4_flash_pretrain_config", "--mode", "pretrain"]) + + handles.recipe_runner.apply_runtime_environment.assert_called_once_with(config) + handles.recipe_runner.bootstrap_recipe_environment.assert_not_called() + + def test_benchmark_dry_run_accepts_config_overrides(monkeypatch): module, handles = _load_module() monkeypatch.delenv("WORLD_SIZE", raising=False) @@ -1447,3 +1477,4 @@ def test_config_container_overrides_are_forwarded_directly(): ["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_not_called() diff --git a/tests/unit_tests/skills/create_model_verification_card/test_validate_card.py b/tests/unit_tests/skills/create_model_verification_card/test_validate_card.py index efc9463a47..85784a930f 100644 --- a/tests/unit_tests/skills/create_model_verification_card/test_validate_card.py +++ b/tests/unit_tests/skills/create_model_verification_card/test_validate_card.py @@ -22,9 +22,10 @@ ("deepseek-v3", "pretrain_performance", "GB200"): (4096, 4096, 256), ("deepseek-v3", "pretrain_performance", "GB300"): (4096, 4096, 256), ("deepseek-v4-flash", "pretrain", "GB200"): (4096, 256, 64), - ("deepseek-v4-flash", "sft", "GB200"): (1024, 128, 64), - ("deepseek-v4-flash", "sft_long_context", "GB200"): (1024, 128, 64), ("deepseek-v4-flash", "checkpoint_resume", "GB200"): (4096, 256, 64), + ("deepseek-v4-flash", "sft", "GB200"): (1024, 128, 32), + ("deepseek-v4-flash", "peft", "GB200"): (1024, 128, 32), + ("deepseek-v4-flash", "sft_long_context", "GB200"): (1024, 128, 64), ("deepseek-v4-flash", "pretrain_performance", "GB200"): (4096, 2048, 128), ("deepseek-v4-flash", "pretrain_performance", "GB300"): (4096, 2048, 128), ("gemma-4-26b-a4b-it", "sft", "H100"): (4096, 32, 8), diff --git a/tests/unit_tests/skills/create_model_verification_card/test_verify_hf_inference.py b/tests/unit_tests/skills/create_model_verification_card/test_verify_hf_inference.py index f98a69b174..49f4ddae26 100644 --- a/tests/unit_tests/skills/create_model_verification_card/test_verify_hf_inference.py +++ b/tests/unit_tests/skills/create_model_verification_card/test_verify_hf_inference.py @@ -18,6 +18,7 @@ import importlib.util import sys +import types from pathlib import Path import pytest @@ -85,6 +86,78 @@ def generate(self, **kwargs): return torch.tensor([[1, 2, 3, 4]]) +def test_sharded_load_uses_device_map_without_moving_model(monkeypatch): + module = _load_module() + calls = [] + + class _LoadedModel: + hf_device_map = {"model.embed_tokens": 0, "model.layers": 1} + + @classmethod + def from_pretrained(cls, model_name, **kwargs): + calls.append((model_name, kwargs)) + return cls(), {key: [] for key in module._LOADING_ISSUE_KEYS} + + def to(self, device): + raise AssertionError(f"sharded model must not be moved to {device}") + + def eval(self): + return self + + def modules(self): + return iter((self,)) + + transformers = types.ModuleType("transformers") + transformers.AutoModelForCausalLM = _LoadedModel + transformers.AutoTokenizer = types.SimpleNamespace(from_pretrained=lambda *args, **kwargs: _Tokenizer()) + monkeypatch.setitem(sys.modules, "transformers", transformers) + args = types.SimpleNamespace( + image=None, + hf_model="model", + trust_remote_code=True, + dtype="bfloat16", + device_map="auto", + device="cuda", + require_gpu_only=True, + ) + + _, model, _ = module._load_runtime(args) + + assert isinstance(model, _LoadedModel) + assert calls == [ + ( + "model", + { + "dtype": torch.bfloat16, + "trust_remote_code": True, + "output_loading_info": True, + "device_map": "auto", + }, + ) + ] + + +def test_model_input_device_prefers_input_embeddings(): + module = _load_module() + model = types.SimpleNamespace( + device=torch.device("cuda:1"), + get_input_embeddings=lambda: types.SimpleNamespace(weight=torch.empty(1, device="meta")), + ) + + assert module._model_input_device(model) == torch.device("cuda:1") + + model.get_input_embeddings = lambda: types.SimpleNamespace(weight=torch.empty(1)) + assert module._model_input_device(model) == torch.device("cpu") + + +def test_gpu_only_placement_rejects_cpu_or_disk_shards(): + module = _load_module() + model = types.SimpleNamespace(hf_device_map={"layers.0": 0, "layers.1": "cpu", "layers.2": "disk"}) + + with pytest.raises(RuntimeError, match="non-GPU placements"): + module._validate_gpu_only_placement(model) + + def test_image_content_uses_processor_native_location_keys(): module = _load_module() @@ -220,3 +293,36 @@ def test_text_main_keeps_the_legacy_tokenizer_path(monkeypatch): assert tokenizer.prompt == "formatted text prompt" assert len(model.calls) == 1 assert "pixel_values" not in model.calls[0] + + +def test_text_main_enables_requested_autocast(monkeypatch): + module = _load_module() + tokenizer = _Tokenizer() + + class _AutocastModel(_Model): + def generate(self, **kwargs): + assert torch.is_autocast_enabled("cpu") + assert torch.get_autocast_dtype("cpu") == torch.bfloat16 + return super().generate(**kwargs) + + model = _AutocastModel() + monkeypatch.setattr(module, "_load_runtime", lambda args: (torch, model, tokenizer)) + monkeypatch.setattr( + sys, + "argv", + [ + "verify_hf_inference.py", + "--hf-model", + "model", + "--prompt", + "Describe Paris.", + "--max-new-tokens", + "2", + "--dtype", + "bfloat16", + "--autocast", + ], + ) + + assert module.main() == 0 + assert len(model.calls) == 1 diff --git a/tests/unit_tests/training/utils/test_flop_utils.py b/tests/unit_tests/training/utils/test_flop_utils.py index 6af0cb17ee..6dc69077db 100644 --- a/tests/unit_tests/training/utils/test_flop_utils.py +++ b/tests/unit_tests/training/utils/test_flop_utils.py @@ -14,6 +14,8 @@ """Unit tests for flop_utils module.""" +import subprocess +import sys from dataclasses import dataclass, field from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -22,6 +24,7 @@ import torch from megatron.bridge.peft.lora import LoRA +from megatron.bridge.training.utils import flop_utils from megatron.bridge.training.utils.flop_utils import ( GlobalFlopsRuntimeStats, _lora_seq_stats_cache, @@ -711,6 +714,39 @@ def _qwen35_27b_config(self, **overrides): defaults.update(overrides) return MockModelConfig(**defaults) + def test_flop_utils_imports_without_mcore_gdn_helper(self) -> None: + code = """ +import importlib +import sys + +from megatron.core.models.gpt import experimental_attention_variant_module_specs +from megatron.bridge.training import utils + +import megatron.bridge.training.utils.flop_utils + +experimental_attention_variant_module_specs.__dict__.pop("is_gated_delta_net_variant", None) +del sys.modules["megatron.bridge.training.utils.flop_utils"] +del utils.flop_utils +flop_utils = importlib.import_module("megatron.bridge.training.utils.flop_utils") + +assert flop_utils._mcore_is_gated_delta_net_variant is None +assert flop_utils._is_gated_delta_net_variant("gdn") is True +""" + result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, check=False) + + assert result.returncode == 0, result.stderr + + @pytest.mark.parametrize( + ("variant", "expected"), + [("gated_delta_net", True), ("gdn", True), ("gdn2", True), ("mamba", False), (None, False)], + ) + def test_gdn_variant_fallback_supports_older_mcore_dev( + self, monkeypatch: pytest.MonkeyPatch, variant: str | None, expected: bool + ) -> None: + monkeypatch.setattr(flop_utils, "_mcore_is_gated_delta_net_variant", None) + + assert flop_utils._is_gated_delta_net_variant(variant) is expected + def test_gdn_flops_differ_from_pure_attention(self): """GDN-enabled config should produce different FLOPs than pure-attention baseline.""" batch_size = 1