DeepSeek-V4 training support
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. DeepSeek-V4 also includes DeepSeek-V4-Flash-Vision-Exp, an experimental multimodal variant that extends the V4-Flash architecture with visual modules and visual-understanding capabilities.
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. The newly released DeepSeek-V4-Flash-Vision-Exp is not yet supported in Megatron Core.
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-09-18
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 |
| Multimodal / vision | DeepSeek-V4-Flash-Vision-Exp has been released; Megatron Core training support is not yet available | |
| 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:
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: ✅ 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
1.5 Multimodal / Vision Training Support
Status: In progress
DeepSeek-V4-Flash-Vision-Exp is the first experimental multimodal model in the DeepSeek-V4 family. It extends DeepSeek-V4-Flash with visual modules and continued multimodal training while retaining comparable text-only agent performance.
Target coverage
- Multimodal model support
- Multimodal training recipe
- End-to-end training validation
Reference: deepseek-ai/DeepSeek-V4-Flash-Vision-Exp
Implementation PRs
| Capability | Status | PR |
|---|---|---|
| Add DeepSeek-V4-Flash-Vision support to multimodal_dev | #7022 |
2. Performance Optimization
2.1 Long-Context Optimization
Status: ✅ Available in dev
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 | ✅ | #7036 |
| Share rotary position embeddings across multi-latent-attention layers | #7111 | |
| Compact CP CSA compressor projections instead of hidden states | #7388 |
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 | |
| support thd full iter graph under fixed GA | #5807 | |
| Share CUDA graph slots across dynamic CP variants | #6479 | |
| Preserve HybridEP capacity across CUDA graph replay | ✅ | #6999 |
| Enable mHC attention CUDA graph split for HybridStack and packed THD | ✅ | #7385 |
2.5 Fusions
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 |
| DeepSeek-V4 CSA optimization: fuse some THD-related ops | ✅ | #7064 |
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 MegaMoE Kernel and Integration
Status: MOK available in dev
| Capability | Status | PR |
|---|---|---|
| mok megakernel integration | ✅ | #6572 |
| MegaMOE Integration | #6810 |
2.6 Parallelism
2.6.1 Parallelism Optimizations
Status: In progress
| Capability | Status | PR |
|---|---|---|
| enable parallel PP prewarm | #6850 | |
| enable TP attention and CSA sequence-parallel paths | #6887 | |
| Skip shape exchange for fixed packed sequences | ✅ | #7000 |
2.6.2 mHC with EP Overlap
Status: ✅ Available in dev
| Capability | Status | PR |
|---|---|---|
| Correct mHC boundaries in the EP-overlap schedule | ✅ | #5471 |
| mHC selective recompute with CUDA Graphs under EP A2A overlap | ✅ | #5841 |
3. Correctness, Stability, and Maintenance
3.1 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 |
| Fix hybrid MTP partial CUDA graph capture | ✅ | #6583 |
| Fix mock data, hybrid recompute, CUDA graph reuse, and Hash MoE selection | ✅ | #6704 |
| Preserve offload events in Hybrid mHC CUDA graphs | ✅ | #7014 |
| Fix regression in CSA for packed THD sequence training on SM100 | #7343 | |
| Make DSv4 bit-reproducible under --deterministic-mode | #7441 |
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 |
| add DSv4 Flash THD 4K and 64K recipe | ✅ | #6911 |
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
Source: NVIDIA/Megatron-LM