DeepSeek v3.2 support

[yueming-yuan]

Major changes

Megatron side

Update Megatron version to dev branch

dev branch megatron provides support for deepseek sparse attention (DSA) (ref: NVIDIA/Megatron-LM#2154) with torch native attention (sbhd format, without CP). Our implementation starts from this version and extends with CP and thd conversion.

better DSA support (support CP and thd format):

1. Implement CP in DSA
   - Index compute: because k only has one head, use one single all-gather on k and compute the entire index of size [q_local, k_global] in one single GEMM. (location: megatron.patch:289-316)
   - core attention: For all TE backends, indexer style attention (namely, supporting arbitrary attention bias) + thd format is not supported. For quick iteration, we use the eagle attention with CP in Megatron dev branch (ref: [Community][Dev] feat(moe): Adding context parallel support to eager attention implementation NVIDIA/Megatron-LM#1859), with some minor modifications for DSA. (location: megatron.patch:322-408)
   - index loss compute: In this part, we need to re-compute the attention score (softmax((q@k)*scale), and sum up on n_head dimension. We implement the CP logic with an all-gather backend, parallel on the n_heads dimension. We use double buffers for k and slice along n_heads dimension, and sum each head up. (Code index: location:120-283)

2. Support THD format with ad hoc converter
   Currently, Megatron's DSA implementation does not support thd format, which is required in Miles. As a quick solution, we implemented a converter, which gathers activation from zigzag-ed thd format and reshards it to zigzag-ed bsnd format on each rank, and reversely for attention output. (Code index: megatron.patch:415-526)

Miles side

Support Deepseek-v32 in miles

(Node: I use mbridge rather than megatron-bridge because DeepSeek-v3.2 is not supported by transformers..)

1. mbrigde:
   location: miles_plugins/mbridge/deepseekv32.py, miles_plugins/mbridge/__init__.py
2. corresponding weight update logic:
   location: miles/backends/megatron_utils/megatron_to_hf/deepseekv32.py, miles/backends/megatron_utils/megatron_to_hf/__init__.py

New Docker for deepseek-v32

1. Added Megatron patch (see above)
2. Added transformers patch, since DeepSeek-v3.2 is not officially supported in transformers
3. install fast-hadamard-transform
4. update transformer_engine version
   location: docker/deepseekv32

Fix for Megatron version update

1. return num_tokens as torch.tensor(1) rather than int(1)
   Newer Megatron require receiving num_tokens as a torch tensor for later processing
   location: miles/backends/megatron_utils/loss.py
2. add explicit argument name for get_gpt_layer_with_transformer_engine_spec because of API change
   location: miles/backends/megatron_utils/model_provider.py

Fix quantization scale tensor syncing for MoE with fused backend

Need to send float32 quant scale tensor when fused backend is applied (when deepep is not enabled, will fall back to fused backend)
location: miles/backends/megatron_utils/megatron_to_hf/processors/quantizer.py

chat_template hacking for DeepSeek-v32

DeepSeek-v32 cannot load chat_template automatically. So we need to do some hacking here. (Reference: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/entrypoints/openai/serving_chat.py#L288)
location: miles/utils/data.py

training script for debugging

location: scripts/run_deepseek_v3.2_5layer.py

Test result

1. Megatron side, unit test:

• DSAttention: equivalent result when use CP=1 and CP=2
• MultiLatentAttention(containing DSAttention): equivalent result when use (1) sbhd format, CP=1 (2) thd format, CP=1, (3)thd format, CP=2

2. Miles side, training test
   successful run with debugging mode (reduce model layers to 5), with TP = 8, CP = 8, seq_len = 8192, batch = 256.

[jeromeku]

@yueming-yuan @xiuhu17

Are there plans for a fused FWD / BWD DSA kernel?

IIUC, the current implementation is using Megatron's unfused_dsa which is essentially a full attention kernel with a large "sparse" [bs, sq, skv] mask. This is inefficient from a computational and memory perspective and essentially defeats the purpose of the sparse formulation.

Currently, FlashMLA currently only has fwd prefill for sparse attention for sm_{90,100}. Are you looking into implementing a fused bwd pass?

[yueming-yuan]

@yueming-yuan @xiuhu17

Are there plans for a fused FWD / BWD DSA kernel?

IIUC, the current implementation is using Megatron's unfused_dsa which is essentially a full attention kernel with a large "sparse" [bs, sq, skv] mask. This is inefficient from a computational and memory perspective and essentially defeats the purpose of the sparse formulation.

Currently, FlashMLA currently only has fwd prefill for sparse attention for sm_{90,100}. Are you looking into implementing a fused bwd pass?

Hi @jeromeku , in the current PR, we're using tilelang's fused implementation instead. It should be much faster than the torch native one. Because in most time the training part is not the bottleneck in RL, the tilelang kernel should already be good.