Question on context parallelism

[mces89]

Hi team, I saw recently axolotl start supporting context parallelism with accelerate, from the accelerate doc: https://github.com/huggingface/accelerate/blob/main/docs/source/concept_guides/context_parallelism.md
it mentions something like "Context parallelism works only with SDPA and only with no mask or causal mask. We can't properly detect this for you, so it's your responsibility to ensure that you are using SDPA with no mask or causal mask. If you use any other attention implementation, it will raise an error."
Does this mean axolotl's cp can only use with sdpa and pre-training style sft like no causal mask on the prompt? Thanks.

[NanoCode012]

Nope, Context Parallelism works with sample packing and SFT. We use ring-flash-attn (and FA) to make this work.

I don't think ours work with SDPA. Edit: Correction see below.

@djsaunde , does SP work with pretraining?

[djsaunde]

Ours works with both SDPA and flash attention, the former being thanks to our integration with accelerate.

Both should work in pretraining and fine-tuning contexts.

[mces89]

Thanks @djsaunde for the confirmation. Do you have any example using fsdp+context parallelism with SFT in both single node and multiple nodes?

[NanoCode012]

We have an example of FSDP+TP+CP here https://github.com/axolotl-ai-cloud/axolotl/blob/main/examples/distributed-parallel/qwen3-8b-fsdp-tp-cp.yaml

For multi-node, you can see https://github.com/axolotl-ai-cloud/axolotl/tree/main/examples/distributed-parallel#multi-node for more understanding.