NVIDIA / NVIDIA/Megatron-LM

DeepSeek-V4 training support

Open
#4,468 5 comments 26 reactions 1 assignee Claimed by @hxbai View on GitHub
enhancement
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

> 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-08

## 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

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

#### Implementation PRs

| Capability | Status | PR |
|---|:---:|---|
| Hybrid CSA/HCA attention | βœ… | [#4458](https://github.com/NVIDIA/Megatron-LM/pull/4458) |
| Hash routing and ClampedSwiGLU | βœ… | [#4481](https://github.com/NVIDIA/Megatron-LM/pull/4481) |
| Base mHC implementation | βœ… | [#2943](https://github.com/NVIDIA/Megatron-LM/pull/2943) |
| MTP support with mHC and the updated mHC contract | βœ… | [#4518](https://github.com/NVIDIA/Megatron-LM/pull/4518) |
| DeepSeek-V4 HybridModel support | βœ… | [#4949](https://github.com/NVIDIA/Megatron-LM/pull/4949), [#5042](https://github.com/NVIDIA/Megatron-LM/pull/5042), [#5485](https://github.com/NVIDIA/Megatron-LM/pull/5485), [#5762](https://github.com/NVIDIA/Megatron-LM/pull/5762) |

Related tracking:

- Sparse-attention feature request: [#4252](https://github.com/NVIDIA/Megatron-LM/issues/4252)
- mHC feature request: [#2890](https://github.com/NVIDIA/Megatron-LM/issues/2890)
- mHC design proposal: [#2919](https://github.com/NVIDIA/Megatron-LM/issues/2919)

### 1.2 Packed Sequence / THD

**Status:** βœ… Available in `dev`

#### Coverage

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

#### Implementation PRs

| Capability | Status | PR |
|---|:---:|---|
| HybridEP/DeepEP dispatcher support under THD | βœ… | [#4816](https://github.com/NVIDIA/Megatron-LM/pull/4816) |
| Variable-length dataset and THD end-to-end benchmark | βœ… | [#4832](https://github.com/NVIDIA/Megatron-LM/pull/4832) |
| DSv4 hybrid-attention packed-sequence support | βœ… | [#5011](https://github.com/NVIDIA/Megatron-LM/pull/5011) |

### 1.3 Context Parallelism and Long-Context Training

**Status:** βœ… Base capability available in `dev`

#### Coverage

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

#### Implementation PRs

| Capability | Status | PR |
|---|:---:|---|
| Context-parallel support for DSv4 hybrid attention | βœ… | [#5087](https://github.com/NVIDIA/Megatron-LM/pull/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](https://github.com/NVIDIA/Megatron-LM/pull/5320) |
| LayerWise optimizer integration with DDP buffers | βœ… | [#4509](https://github.com/NVIDIA/Megatron-LM/pull/4509) |
| Route non-Muon parameters through DistributedOptimizer | βœ… | [#4771](https://github.com/NVIDIA/Megatron-LM/pull/4771) |
| Decoupled compact LayerWise DDP layout | βœ… | [#5388](https://github.com/NVIDIA/Megatron-LM/pull/5388) |
| Muon FP8 primary weight | βœ… | [#5470](https://github.com/NVIDIA/Megatron-LM/pull/5470) |

Reference: [NVIDIA Emerging Optimizers](https://github.com/NVIDIA-NeMo/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](https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash-Vision-Exp)

#### Implementation PRs

| Capability | Status | PR |
|---|:---:|---|
| Add DeepSeek-V4-Flash-Vision support to multimodal_dev | 🚧 | [#7022](https://github.com/NVIDIA/Megatron-LM/pull/7022) |

## 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](https://github.com/NVIDIA/Megatron-LM/pull/5691) |
| Load-balanced context-parallel DSA indexer | 🚧 | [#6058](https://github.com/NVIDIA/Megatron-LM/pull/6058) |
| Optimize MTP contiguous packed-CP rolls | βœ… | [#6246](https://github.com/NVIDIA/Megatron-LM/pull/6246) |

### 2.2 Memory Optimization

**Status:** 🚧 In progress

#### Implementation PRs

| Optimization | Status | PR |
|---|:---:|---|
| Chunked optimizer-state and master-weight offload | βœ… | [#6244](https://github.com/NVIDIA/Megatron-LM/pull/6244) |
| Fuse DSv4 output inverse RoPE into CSA sparse-attention functions | 🚧 | [#7036](https://github.com/NVIDIA/Megatron-LM/pull/7036) |

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](https://github.com/NVIDIA/TransformerEngine/pull/3160) |
| MCore `DSv4HybridAttention` integration | 🚧 | [#5722](https://github.com/NVIDIA/Megatron-LM/pull/5722) |

#### 2.3.2 FP8/MXFP8 DSA Indexer

Current coverage:

- [x] FP8/MXFP8 indexer kernels
- [x] 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](https://github.com/NVIDIA/cudnn-frontend/pull/370) |
| Compact BF16/MXFP8 DSA indexer integration in MCore | βœ… | [#5992](https://github.com/NVIDIA/Megatron-LM/pull/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](https://github.com/NVIDIA/Megatron-LM/pull/4190) |
| Partial/per-layer CUDA Graph support for THD training | βœ… | [#4359](https://github.com/NVIDIA/Megatron-LM/pull/4359) |
| Chunk-wise whole-block CUDA Graph support for THD training | 🚧 | [#5258](https://github.com/NVIDIA/Megatron-LM/pull/5258) |
| Preserve HybridEP capacity across CUDA graph replay | βœ… | [#6999](https://github.com/NVIDIA/Megatron-LM/pull/6999) |

### 2.5 Fusions

#### 2.5.1 mHC Fusion Kernels

**Status:** βœ… Available in `dev`

| Component | Status | PR |
|---|:---:|---|
| MCore cuTile mHC fusion | βœ… | [#3828](https://github.com/NVIDIA/Megatron-LM/pull/3828) |
| Faster MCore fused mHC implementation | βœ… | [#4624](https://github.com/NVIDIA/Megatron-LM/pull/4624) |

Reference implementation: [DeepSeek TileKernels mHC](https://github.com/deepseek-ai/TileKernels/tree/main), [Transformer Engine Triton mHC kernels and PyTorch APIs #2790](https://github.com/NVIDIA/TransformerEngine/pull/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](https://github.com/NVIDIA/Megatron-LM/pull/4894) |
| THD CSA compressor gated-pooling dispatch to cuDNN Frontend | βœ… | [#5984](https://github.com/NVIDIA/Megatron-LM/pull/5984) |
| DeepSeek-V4 CSA optimization: fuse some THD-related ops | βœ… | [#7064](https://github.com/NVIDIA/Megatron-LM/pull/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](https://github.com/NVIDIA/Megatron-LM/pull/5130) |

#### 2.5.4 MegaMoE Kernel and Integration

**Status:** 🚧 MOK available in `dev`

| Capability | Status | PR |
|---|:---:|---|
| mok megakernel integration | βœ… | [#6572](https://github.com/NVIDIA/Megatron-LM/pull/6572) |
| MegaMOE Integration | 🚧 | [#6810](https://github.com/NVIDIA/Megatron-LM/pull/6810) |

### 2.6 Parallelism

#### 2.6.1 Parallelism Optimizations

**Status:** 🚧 In progress

| Capability | Status | PR |
|---|:---:|---|
| enable parallel PP prewarm | 🚧 | [#6850](https://github.com/NVIDIA/Megatron-LM/pull/6850) |
| enable TP attention and CSA sequence-parallel paths | 🚧 | [#6887](https://github.com/NVIDIA/Megatron-LM/pull/6887) |
| Skip shape exchange for fixed packed sequences | 🚧 | [#7000](https://github.com/NVIDIA/Megatron-LM/pull/7000) |

#### 2.6.2 mHC with EP Overlap

**Status:** 🚧 Further integration work in progress

| Capability | Status | PR |
|---|:---:|---|
| Correct mHC boundaries in the EP-overlap schedule | βœ… | [#5471](https://github.com/NVIDIA/Megatron-LM/pull/5471) |
| mHC selective recompute with CUDA Graphs under EP A2A overlap | βœ… | [#5841](https://github.com/NVIDIA/Megatron-LM/pull/5841) |

## 3. Correctness, Stability, and Maintenance

### 3.1 Merged Bug Fixes

| Fix | Status | PR |
|---|:---:|---|
| Fix dense loss and RoPE type in DSv4 hybrid attention | βœ… | [#5018](https://github.com/NVIDIA/Megatron-LM/pull/5018) |
| Correct DSv4 hybrid Q-up FLOPs calculation | βœ… | [#5142](https://github.com/NVIDIA/Megatron-LM/pull/5142) |
| Restore DSv4 TFLOPS calculation, including packed sequence | βœ… | [#5358](https://github.com/NVIDIA/Megatron-LM/pull/5358) |
| Keep the CSA compressor and indexer in high precision during FP8 training | βœ… | [#5308](https://github.com/NVIDIA/Megatron-LM/pull/5308) |
| Preserve DSA output across fused inverse RoPE | βœ… | [#5526](https://github.com/NVIDIA/Megatron-LM/pull/5526) |
| Fix MTP with contiguous CP partition mode | βœ… | [#5706](https://github.com/NVIDIA/Megatron-LM/pull/5706) |
| Fix zero-loss indexer gradients under DSv4 context parallelism | βœ… | [#5809](https://github.com/NVIDIA/Megatron-LM/pull/5809) |
| Fix distributed optimizer state save/load for mixed-dtype parameter groups | βœ… | [#5835](https://github.com/NVIDIA/Megatron-LM/pull/5835) |
| Use the full CSA denominator for unfused indexer loss | βœ… | [#5960](https://github.com/NVIDIA/Megatron-LM/pull/5960) |
| Keep mHC mapping computation in FP32 on the fused cuTile path | βœ… | [#6172](https://github.com/NVIDIA/Megatron-LM/pull/6172) |
| Fix fused CSA indexer-loss normalization and compact attention indices | βœ… | [#6349](https://github.com/NVIDIA/Megatron-LM/pull/6349) |
| Fix hybrid MTP partial CUDA graph capture | βœ… | [#6583](https://github.com/NVIDIA/Megatron-LM/pull/6583) |
| Fix mock data, hybrid recompute, CUDA graph reuse, and Hash MoE selection | βœ… | [#6704](https://github.com/NVIDIA/Megatron-LM/pull/6704) |
| Preserve offload events in Hybrid mHC CUDA graphs | 🚧 | [#7014](https://github.com/NVIDIA/Megatron-LM/pull/7014) |

### 3.2 Merged Refactors and Supporting Changes

| Change | Status | PR |
|---|:---:|---|
| Migrate the main-first DSA implementation to `dev` | βœ… | [#6020](https://github.com/NVIDIA/Megatron-LM/pull/6020) |
| Move CSA implementation helpers into `csa_utils` | βœ… | [#6372](https://github.com/NVIDIA/Megatron-LM/pull/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](https://github.com/NVIDIA/Megatron-LM/pull/5266) |
| add DSv4 Flash THD 4K and 64K recipe | βœ… | [#6911](https://github.com/NVIDIA/Megatron-LM/pull/6911) |

Megatron Bridge resources:

- Examples: [`examples/models/deepseek_v4`](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples/models/deepseek_v4)
- Validation roadmap: [NVIDIA-NeMo/Megatron-Bridge#4633](https://github.com/NVIDIA-NeMo/Megatron-Bridge/issues/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](https://github.com/NVIDIA-NeMo/Automodel/pull/2039) |
| Optimized TileLang kernels and FSDP integration | βœ… | [NVIDIA-NeMo/Automodel#2076](https://github.com/NVIDIA-NeMo/Automodel/pull/2076) |

## References

- [DeepSeek-V4 technical report](https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro/blob/main/DeepSeek_V4.pdf)
- [Hugging Face Transformers implementation](https://github.com/huggingface/transformers/pull/45616)
- [DeepSeek TileKernels](https://github.com/deepseek-ai/TileKernels/tree/main)
- [NVIDIA Emerging Optimizers](https://github.com/NVIDIA-NeMo/Emerging-Optimizers)
- [NVIDIA cuDNN Frontend](https://github.com/NVIDIA/cudnn-frontend)

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.