Context Parallel for Qwen3

[unavailableun]

Thanks for supporting Qwen3 models!

CP is not supported currently because of RoPE embedding implementation details.

Any plan to support CP + EP for Qwen3 MoE models?
If no plan in short time, can you help guide how can I implement it myself?

[tianyu-l]

cc @fegin

[unavailableun]

I saw alert only for CP & FlexAttention, does that mean CP works with SDPA?

[tianyu-l]

cc @wwwjn @shuhuayu

[wwwjn]

Thanks for asking! Now Qwen3 is not officially supported for Qwen3, and we haven't implemented and tested CP on Qwen3. This is because of the RoPE embedding differences (In Qwen3, the RoPE field is named "rope_cache" while in train.py, it's looking for freqs_cis field.
Qwen3 README: https://github.com/pytorch/torchtitan/tree/refs/heads/main/torchtitan/models/qwen3
Train.py: https://github.com/pytorch/torchtitan/blob/refs/heads/main/torchtitan/train.py#L480

So it won't work with SDPA, but it should be a simple change