Support context parallel training with ring-flash-attention
Author: zhuzilin
Pull request: no
Labels: name: Feature request

name: Good Difficult Issue

name: Flash Attention

### Feature request

Hi, I'm the author of [zhuzilin/ring-flash-attention](https://github.com/zhuzilin/ring-flash-attention).

I wonder if you are interested in integrating context parallel with [zhuzilin/ring-flash-attention](https://github.com/zhuzilin/ring-flash-attention), so that user can train llm with long data more efficiently.

### Motivation

As openai o1 released, it will probably be common for people to train model with really long cot data. And it will be nice if most model within the transformers library can support training with long context efficiently with certain type of context parallel, i.e. the context length scale linearly with the number of GPUs.

The 3 existing context parallel methods are the deepspeed ulysses, ring attention and the one proposed in the [llama3 tech report](https://arxiv.org/abs/2407.21783). The deepspeed ulysses will be limited by the number of kv heads (the maximum context length can be `num_head_kv * seq_length_per_gpu`), which makes it a little unfriendly to GQA models. So it will be great if the transformers library could support the one or both of the other 2 context parallel methods.

And both ring attention and the llama3 strategy are supported with flash attention in [zhuzilin/ring-flash-attention](https://github.com/zhuzilin/ring-flash-attention), whose correctness has been proved by [jzhang38/EasyContext](https://github.com/jzhang38/EasyContext). The library basically has the same api as flash attention, and hides the communication required from its user to make it a easy substitution from any origin flash attention api callsite.

Therefore, I believe it will be easy to support the context parallel with [zhuzilin/ring-flash-attention](https://github.com/zhuzilin/ring-flash-attention). For example, we could have different branch in `modeling_flash_attention_utils._flash_attention_forward`.

### Your contribution

I'd love to help if you have interests :)

Comments:
cc @ArthurZucker 

Hey hey! I think similarly to the recent integration of liger kernels this would make sense for sure to have a soft dependency, and monkey patch the `flash_attn_forward` 🤗 happy to review a pr like this!

Hi, that's very nice :)   Has this feature been integrated into the Huggingface library now?

Hmm... After a brief look through, it seems hard to give a clean PR to transformers repo for ring flash attn... There are main obstacles:
1. For really long data, different rank may need to share the same data, e.g. 4 gpu sharing the same batch with only 1 sample within. In that case, we need to customize the `num_replicas` and `rank` passed to samplers, which is currently managed in [accelerate](https://github.com/huggingface/accelerate). And we could not pass config to accelerate to change the behavior at the moment.
2. The ring attention need to manage a new process group for the GPUs that share the same batch, so that they can communicate with each other during the attention calculation. There is no clear location where to put that. (we can set a global variable for the process group, but I'm not sure if that will broke some config branch of the `Trainer`....)
3. We cannot only change the accerlate repo, because we need a new way to manage the data in `Trainer`, something like `DataCollatorWithFlattening`.

I'm afraid we have to postpone this feature after some prerequisites are ready.

Just found out that the deepspeed team is adding ulysses. We can wait until they land the feature, as ring attention and deepspeed ulysses will share many prerequisites~

DeepSpeed added Ulysses. Any plans to integrate this (highly needed for reasoning models) @ArthurZucker @zhuzilin 

https://www.deepspeed.ai/tutorials/ds-sequence/

> DeepSpeed added Ulysses.

@casper-hansen That's great! Could you point to the PR that integrate ulysses into huggingface transformers/accelerator?

@zhuzilin it's not integrated into huggingface. it's only in deepspeed

Pr (#35301) was unfortunately closed! Anyone can work on this!