[Feature-request] Context Parallel support for DSAttention
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 3h
- Merged PRs (30d)
- 272
Description
Megatron-Core currently does not support Context Parallelism (CP) with `DSAttention`. This blocks long-context training and post-training for models that use the experimental DSA path, such as GLM-style sparse attention models. This also affects Megatron Bridge/ Nemo-RL usage, since the Megatron Bridge GLM recipe points to `experimental_attention_variant == "dsa"`.
### Current Behavior
When `experimental_attention_variant == "dsa"`, [transformer_config.py](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/transformer/transformer_config.py#L2452) asserts that `context_parallel_size == 1`:
```python
assert self.context_parallel_size == 1, "Currently context parallelism is not supported by DSAttention!"
```
### Possible Solution / Reference
One possible reference is the Slime RL PR that added GLM 5.1 support:
https://github.com/THUDM/slime/pull/1599
From a preliminary read, Slime does not use mcore DSAttention directly for that recipe. Instead, it builds mcore GPT/MLA layers but replaces the attention path with a custom GLM5 DSA implementation. They also add a Slime-specific [--allgather-cp](https://github.com/THUDM/slime/blob/main/slime/backends/megatron_utils/data.py#L78) flag to support Context Parallelism
This is only a possible design reference, not necessarily a request to implement the Slime approach directly.
Contributor guide
Assessment
This issue has not been assessed yet.