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.

References:

• Technical report: https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro/blob/main/DeepSeek_V4.pdf
• Hugging Face implementation PR: Add DeepSeek V4 huggingface/transformers#45616
• TileKernels mHC reference: https://github.com/deepseek-ai/TileKernels/tree/main
• Emerging Optimizers Muon reference: https://github.com/NVIDIA-NeMo/Emerging-Optimizers

Last updated: 2026-06-10

Model Architecture

• Status: Partial Support
• Requested Features:
  • Megatron-LM configs for DeepSeek-V4-Flash and DeepSeek-V4-Pro training
  • CSA/HCA layer schedule wired into the model spec
  • Hash routing initial MoE layer support
  • mHC and MTP support
• Overall Status:
  • Hybrid CSA/HCA Attention
  • Hash routing MoE and ClampedSwiGLU
  • mHC and MTP support

Hybrid CSA/HCA Attention

• Status: Merged in Dev
• Implementation PR:
  • [dev] [DeepSeek-v4] Part 1: Hybrid Attention with CSA and HCA #4458
  • [dev] [DeepSeek-v4] Part 4: Fusion Kernels for DSv4 Hybrid Attention #4894
• Requested Features:
  • Compressed Sparse Attention (CSA)
  • Heavily Compressed Attention (HCA)
  • Integration with Megatron attention/module specs
• Related: Feature Request: Adding NSA and other sparse attention mechanisms #4252

Hash routing MoE and ClampedSwiGLU

• Status: Merged in Dev
• Implementation PR: [dev] [DeepSeek-v4] Part 2: Hash MoE and SwiGLU clamp #4481
• Requested Features:
  • Hash routing MoE
  • ClampedSwiGLU

mHC Integration

• Status: Merged in Dev
• Implementation PRs:
  • [dev] feat(mHC): Add basic pytorch implementation of manifold hyper connection(mHC). #2943
  • [dev] fix: Support mHC with cuda graph and activation offloading #4190
  • [dev] [DeepSeek-v4] Part 3: MTP support with mHC and new mHC contract #4518
• Tracking Issues:
  • Feature Request: mHC (Manifold-Constrained Hyper-Connections) #2890
  • Design Proposal: mHC (Manifold-Constrained Hyper-Connections) #2919

Sequence and Context Support

Packed Sequence

• Status: In Progress (ETA: May 29)
• Related PRs:
  • [Dev] add support for deepep/hybridep dispatcher under thd format training #4816
  • varlendataset for thd e2e and benchmark #4832
  • [Dev] add cuda graph support for thd format training. #4359
  • [Dev] Add chunk-wise (whole-block) CUDA graph support for THD training #5258
  • [dev] [DeepSeek-v4] Packed Sequence (THD) support for DSv4 Hybrid Attention #5011
• Requested Features:
  • Packed sequence support for DSv4 Hybrid Attention
  • End-to-end packed sequence support
  • CUDA Graphs support for DSv4 THD

Long-Context Training

• Status: In Progress (ETA: Jun 19)
• Related PRs:
  • [dev] [DeepSeek-v4] Context Parallel support #5087
• Requested Features:
  • Context parallel support for DSv4 Hybrid Attention
  • 64K to 1M context training curriculum validation

Optimizer and Training Recipe

Muon Optimizer

• Status: In Progress
• Related PR:
  • bump emerging-optimizer package for deepseek v4 coefficients type #4523
  • Route non-Muon params through DistributedOptimizer #4771
  • Integrate LayerWiseDistributedOptimizer with DDP buffer infrastructure #4509
  • feat(optim): support --fp8-param-gather for muon + mxfp8 in LayerWise #4987
• Requested Features:
  • DeepSeek-V4 Muon/AdamW training recipe with Emerging Optimizers

FP4 QAT Recipe

• Status: Partially supported / needs DeepSeek-V4 specific validation
• Requested Features:
  • MXFP4 QAT for routed expert weights
  • Simulated FP4-to-FP8 training path with FP32 master weights
  • FP4 QK path for CSA indexer
  • BF16 index-score path for top-k selector

Performance Optimization

mHC Fusion Kernels

• Status: Merged in Dev
• Related PR:
  • [dev] mHC kernel fusion  #3828
  • [dev]: faster implementation of mHC fused kernels #4624
  • [TransformerEngine] [Common, PyTorch] Add triton mHC kernels & pytorch APIs TransformerEngine#2790
• TileLang kernels from DeepSeek https://github.com/deepseek-ai/TileKernels/tree/main

DSv4 Hybrid Attention Fusion Kernels

• Status: Merged in Dev
• The fusion kernels are released in https://github.com/NVIDIA/cudnn-frontend release 1.24.0
• Related PR:
  • [dev] [DeepSeek-v4] Part 4: Fusion Kernels for DSv4 Hybrid Attention #4894

FP8 Indexer

• Status: In Progress
• Requested Features:
  • FP8 indexer kernels
  • Integration in MCore
• Related PR:

FP8 parameter + Muon + MXFP8

• Status: In Progress
• Related PR: feat(optim): support --fp8-param-gather for muon + mxfp8 in LayerWise #4987

mHC support with EP Overlapping

• Status: Planning
• Related PR:

Bug Fixes

• Related PR:
  • [dev] [fix] [DeepSeek-v4] fix dense loss and rope type in DSv4 Hybrid Attention #5018

Megatron Bridge Examples

• Status: Needs implementation
• Requested Features:
  • DeepSeek-V4-Flash proxy pretraining recipe
  • DeepSeek-V4-Pro config/provider support
  • Hugging Face ↔ Megatron checkpoint conversion coverage

[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.