DeepSeek-V4 training support

[sbhavani]

This issue tracks DeepSeek-V4 training support in Megatron Core.

DeepSeek-V4 extends DeepSeek-V3/V3.2 with hybrid compressed attention, mHC, updated routing, Muon-based training recipes, FP4 QAT, and million-token context training support.

Overall status: Core model architecture, packed-sequence training, context parallelism, Muon, and the primary fusion paths are available in dev. Current work is focused on CUDA Graph coverage, long-context and memory optimization, low-precision modules, and other performance optimizations.

Status: ✅ available / merged / resolved · 🚧 in progress / open / draft · 📋 planned / pending validation

Status represents end-to-end capability readiness. A merged supporting PR does not necessarily mean the entire capability is complete.

Last updated: 2026-08-14

Status at a glance

Core functionality

Capability	Status	Summary
Model architecture	✅	Hybrid CSA/HCA attention, hash routing, ClampedSwiGLU, mHC, MTP, and HybridModel support
Packed sequence / THD	✅	End-to-end variable-length training and HybridEP/DeepEP dispatcher support
Context parallelism	✅	DSv4 hybrid-attention context-parallel support
Muon optimizer	✅	Muon/AdamW routing, distributed parameter layouts, and FP8 primary-weight gather
Recipes and examples	🚧	Existing recipes and examples are available; broader platform validation remains active

Optimization work

Area	Status	Current focus
Long-context optimization	🚧	Communication overlap, load balancing, CP layouts, and packed-CP efficiency
Memory optimization	🚧	Optimizer-state offload, module-level memory optimization, and long-context memory optimization
Low-precision modules	🚧	BF16/MXFP8 BatchedGEMM, FP8 indexer integration, and FP4 QAT
CUDA Graphs	🚧	Whole-block THD capture
Fusion and other performance work	🚧	Primary fusion paths are available; additional fusion and integration work continues

1. Core Functionality

1.1 Model Architecture

Status: ✅ Available in dev

Coverage

• Megatron-LM configurations for DeepSeek-V4-Flash and DeepSeek-V4-Pro training
• CSA/HCA layer schedule wired into the model spec
• Hash-routing initial MoE layers
• ClampedSwiGLU
• mHC and MTP support
• HybridModel integration

Implementation PRs

Capability	Status	PR
Hybrid CSA/HCA attention	✅	#4458
Hash routing and ClampedSwiGLU	✅	#4481
Base mHC implementation	✅	#2943
MTP support with mHC and the updated mHC contract	✅	#4518
DeepSeek-V4 HybridModel support	✅	#4949, #5042, #5485, #5762

Related tracking:

• Sparse-attention feature request: #4252
• mHC feature request: #2890
• mHC design proposal: #2919

1.2 Packed Sequence / THD

Status: ✅ Available in dev

Coverage

• Packed-sequence support for DSv4 hybrid attention
• End-to-end variable-length THD training
• HybridEP/DeepEP dispatcher support under THD

Implementation PRs

Capability	Status	PR
HybridEP/DeepEP dispatcher support under THD	✅	#4816
Variable-length dataset and THD end-to-end benchmark	✅	#4832
DSv4 hybrid-attention packed-sequence support	✅	#5011

1.3 Context Parallelism and Long-Context Training

Status: ✅ Base capability available in dev

Coverage

• Context-parallel support for DSv4 hybrid attention
• 64K-to-1M-context training curriculum validation

Implementation PRs

Capability	Status	PR
Context-parallel support for DSv4 hybrid attention	✅	#5087

1.4 Muon Optimizer

Status: ✅ Available in dev

Target coverage: DeepSeek-V4 Muon/AdamW training recipes using NVIDIA Emerging Optimizers.

Implementation PRs

Capability	Status	PR
Emerging Optimizers integration	✅	#5320
LayerWise optimizer integration with DDP buffers	✅	#4509
Route non-Muon parameters through DistributedOptimizer	✅	#4771
Decoupled compact LayerWise DDP layout	✅	#5388
Muon FP8 primary weight	✅	#5470

Reference: NVIDIA Emerging Optimizers

2. Performance Optimization

2.1 Long-Context Optimization

Status: 🚧 In progress

The base context-parallel capability is available. This section tracks performance and scalability improvements on top of that support.

Implementation PRs

Optimization	Status	PR
Overlap CP communication for the indexer and compressed KV	🚧	#5691
Load-balanced context-parallel DSA indexer	🚧	#6058
Optimize MTP contiguous packed-CP rolls	✅	#6246

2.2 Memory Optimization

Status: 🚧 In progress

Implementation PRs

Optimization	Status	PR
Chunked optimizer-state and master-weight offload	🚧	#6244
Fuse DSv4 output inverse RoPE into CSA sparse-attention functions	🚧	#6185

Planned work:

• 📋 Indexer memory reduction
• 📋 Q-chunk CSA

2.3 Low-Precision Modules

Status: 🚧 In progress

2.3.1 BF16/MXFP8 BatchedGEMM

Component	Status	PR
Transformer Engine strided BatchedGEMM and BatchedLinear	🚧	TransformerEngine#3160
MCore DSv4HybridAttention integration	🚧	#5722

2.3.2 FP8/MXFP8 DSA Indexer

Current coverage:

• FP8/MXFP8 indexer kernels
• Compressed Top-K indexer-forward path in cuDNN Frontend
• Complete MCore integration and end-to-end validation

Component	Status	PR
cuDNN Frontend FP8/MXFP8 and compressed Top-K paths	✅	cudnn-frontend#370
Compact BF16/MXFP8 DSA indexer integration in MCore	🚧	#5992

2.3.3 FP4 QAT Recipe

Status: 📋 DeepSeek-V4-specific implementation and validation required

• MXFP4 QAT for routed-expert weights
• Simulated FP4-to-FP8 training path with FP32 master weights
• FP4 QK path for the CSA indexer
• BF16 index-score path for the top-k selector

No complete MCore implementation PR is currently tracked.

2.4 CUDA Graph Optimization

Status: 🚧 In progress

Implementation PRs

Capability	Status	PR
mHC CUDA Graph and activation-offload compatibility	✅	#4190
Partial/per-layer CUDA Graph support for THD training	✅	#4359
Chunk-wise whole-block CUDA Graph support for THD training	🚧	#5258

2.5 Fusion and Other Performance Optimization

2.5.1 mHC Fusion Kernels

Status: ✅ Available in dev

Component	Status	PR
MCore cuTile mHC fusion	✅	#3828
Faster MCore fused mHC implementation	✅	#4624

Reference implementation: DeepSeek TileKernels mHC, Transformer Engine Triton mHC kernels and PyTorch APIs #2790

2.5.2 DSv4 Hybrid-Attention Fusion Kernels

Status: ✅ Available in dev

The fused CSA/HCA kernels are available through cuDNN Frontend 1.24.0.

Component	Status	PR
DSv4 hybrid-attention fusion integration	✅	#4894
THD CSA compressor gated-pooling dispatch to cuDNN Frontend	✅	#5984

2.5.3 Fused GroupedGEMM with ClampedSwiGLU

Status: ✅ Available in dev

Component	Status	PR
ClampedSwiGLU support in the MoE MLP fuser and hash-routing force balance	✅	#5130

2.5.4 mHC with EP Overlap

Status: 🚧 Further integration work in progress

Capability	Status	PR
Correct mHC boundaries in the EP-overlap schedule	✅	#5471
mHC selective recompute with CUDA Graphs under EP A2A overlap	🚧	#5841

2.5.5 MegaMoE Kernel and Integration

Status: 🚧 In progress

Implementation work is in progress. No public PR is currently tracked.

3. Correctness, Stability, and Maintenance

3.1 Merged Bug Fixes

Fix	Status	PR
Fix dense loss and RoPE type in DSv4 hybrid attention	✅	#5018
Correct DSv4 hybrid Q-up FLOPs calculation	✅	#5142
Restore DSv4 TFLOPS calculation, including packed sequence	✅	#5358
Keep the CSA compressor and indexer in high precision during FP8 training	✅	#5308
Preserve DSA output across fused inverse RoPE	✅	#5526
Fix MTP with contiguous CP partition mode	✅	#5706
Fix zero-loss indexer gradients under DSv4 context parallelism	✅	#5809
Fix distributed optimizer state save/load for mixed-dtype parameter groups	✅	#5835
Use the full CSA denominator for unfused indexer loss	✅	#5960
Keep mHC mapping computation in FP32 on the fused cuTile path	✅	#6172
Fix fused CSA indexer-loss normalization and compact attention indices	✅	#6349

3.2 Merged Refactors and Supporting Changes

Change	Status	PR
Migrate the main-first DSA implementation to dev	✅	#6020
Move CSA implementation helpers into csa_utils	✅	#6372

4. Recipes and Ecosystem Integration

4.1 Megatron and Megatron Bridge

Status: 🚧 Broader validation in progress

Implementation PRs

Component	Status	PR
MCore DeepSeek-V4-Flash training recipe	✅	#5266

Megatron Bridge resources:

• Examples: examples/models/deepseek_v4
• Validation roadmap: NVIDIA-NeMo/Megatron-Bridge#4633

Current example coverage includes:

• DeepSeek-V4-Flash proxy pretraining recipe
• DeepSeek-V4-Pro configuration/provider support
• Hugging Face ↔ Megatron checkpoint conversion
• Inference and SFT launch examples

4.2 NeMo AutoModel

The following support is maintained outside Megatron Core:

Capability	Status	PR
DeepSeek-V4-Flash SFT and full finetuning with FSDP2	✅	NVIDIA-NeMo/Automodel#2039
Optimized TileLang kernels and FSDP integration	✅	NVIDIA-NeMo/Automodel#2076

References

• DeepSeek-V4 technical report
• Hugging Face Transformers implementation
• DeepSeek TileKernels
• NVIDIA Emerging Optimizers
• NVIDIA cuDNN Frontend

[Minami-su]

1

[yiakwy-xpu-ml-framework-team]

NeMo AutoModel has supported SFT + FSDP2:

1. sft_fsdp : feat: DeepSeek V4 Flash support NVIDIA-NeMo/Automodel#2039
2. tilelang kernel + fsdp : feat: add optimized DeepSeek V4 kernels NVIDIA-NeMo/Automodel#2076

[sbhavani]

@YJMSTR please create a new issue with your bug

[Meirtz]

Originally reported here as a suspected fused MLA yarn-rope kernel NaN in DeepSeek-V4-Flash SFT.

Resolved — root cause was on the Megatron-Bridge side, not the kernel. See the final comment below and the fix in NVIDIA-NeMo/Megatron-Bridge#4271. (Interim debug comments removed to avoid misleading readers.)

[Meirtz]

DeepSeek-V4-Flash SFT NaN with apply_rope_fusion=True — root cause + fix (consolidated record; resolved)

• Symptom: SFT NaNs at iteration 2 ("found NaN in local forward loss calculation") with apply_rope_fusion=True; same config unfused trains clean. Reproduced on H100 and GB300, via verl and standalone Megatron-Bridge.
• Root cause (Megatron-Bridge, not this repo): the generic partial_rotary_factor → rotary_percent mapping double-applies HF's 0.125 (relative to head_dim=512) on top of qk_pos_emb_head_dim=64, which already encodes the rope split. The rope cos/sin cache shrinks to 8 dims; fused_mla_rope_inplace is then called with emb_dim=64 and reads cos/sin out of bounds (in-situ capture: healthy inputs → 1e14–1e36 garbage, nondeterministic because the OOB memory is arbitrary). The unfused path masks the same misconfig by silently rotating only 8/64 dims.
• Fix: NVIDIA-NeMo/Megatron-Bridge#4271 (rotary_percent=1.0 in the DSv4 bridge). Validated on GB300: fused-rope SFT now trains clean and matches the unfused control to 4 decimals.
• Suggestion for this repo: fused_mla_rope_inplace could assert cos.shape[-1] == emb_dim — the OOB read surfaces as nondeterministic downstream NaNs, which is expensive to localize. Happy to send that one-liner if useful.