Add support for context parallelism
Author: lewtun
Pull request: no
Labels: name: Feature request

### Feature request

Long context models like [Qwen/Qwen2.5-7B-Instruct-1M](https://huggingface.co/Qwen/Qwen2.5-7B-Instruct-1M) have support for up to 1M tokens. However, fine-tuning such models in `transformers` leads to OOM errors and special methods like Ring Attention are needed. A similar issue arises during inference, where generating on a 1M prefill gives OOM.

It would be very exciting to have support for _context parallelism_ where in each layer we split the KQV computation across GPUs. 

As far as an API goes, having something like `attn_implemention="ring"` in `from_pretrained()` would likely be the simplest way to support this feature.

Links to papers and code:

* Ring Attention: https://arxiv.org/abs/2310.01889
* Reference code: https://github.com/zhuzilin/ring-flash-attention/blob/main/ring_flash_attn/adapters/hf_adapter.py
* Picotron code: https://github.com/huggingface/picotron/blob/main/picotron/context_parallel/context_parallel.py

### Motivation

The main motivation is two-fold: to support fine-tuning large context models and to enable online RL methods like GRPO to scale better in TRL (where we generate potentially large CoTs during training)

### Your contribution

Happy to review / test the feature on TRL side

Comments:
FSDPv2 (which is supported officially in Torch 2.5.1) has native support for Llama 3-style context parallelism. It's not "true" ring attention but works better under varlen (i.e. packed/masked SFT) training, so that could be an alternative avenue to consider. See their recent blog about it here: https://discuss.pytorch.org/t/distributed-w-torchtitan-breaking-barriers-training-long-context-llms-with-1m-sequence-length-in-pytorch-using-context-parallel/215082

Hi @lewtun 
What about deepspeed ulysse integration in transformers for context parallelism?
https://www.deepspeed.ai/tutorials/ds-sequence/#2-how-to-use-deepspeed-ulysses-in-your-application

It would be nice to fit it with current Deepspeed integration in Transformers, and it can work also with flash-attention-2 (+triton)
https://www.deepspeed.ai/tutorials/ds-sequence/#2-how-to-use-deepspeed-ulysses-in-your-application


reference code
https://github.com/deepspeedai/DeepSpeedExamples/blob/uly-hf/post_training/sequence_parallelism/test_ulysses.py

blog
https://github.com/deepspeedai/DeepSpeed/blob/master/blogs/deepspeed-ulysses/README.md

paper
https://arxiv.org/pdf/2309.14509


I am still talking to internal members to get more familiar with this topic 🙏 


Hi, just some updates:

- Personally, I am trying using [torch's context_parallel](https://pytorch.org/docs/stable/distributed.tensor.html#torch.distributed.tensor.experimental.context_parallel), and below is a minimal example of using it (a simple `forward`). 

- There is another person that is also working on CP with DeepSpeed.


### My minimal code to show how to use CP
```
from torch.distributed.tensor.experimental._attention import _cp_options
_cp_options.enable_load_balance = True

from torch.distributed.device_mesh import init_device_mesh
cp_mesh = init_device_mesh("cuda", (4,))

from torch.distributed.tensor.experimental import context_parallel


import torch.distributed as dist
import torch
from torch.nn.attention import sdpa_kernel, SDPBackend
from transformers import AutoModelForCausalLM




repo = "Qwen/Qwen2.5-Coder-32B-Instruct"
repo = "Qwen/Qwen2.5-Coder-3B-Instruct"

token = ""

model_class = AutoModelForCausalLM
sdpa_backend = SDPBackend.FLASH_ATTENTION
dtype = torch.float32
device = "cuda"

model = model_class.from_pretrained(repo, token=token)
model = model.to(device, dtype=dtype)
model.eval()

batch_size = 1
world_size = 4
seq_len = 32

input_ids = torch.randint(
    low=8,
    high=16,
    size=(batch_size, world_size * seq_len),
    device=device ,
    dtype=torch.int32,
)

position_ids = torch.arange(
    start=0,
    end=world_size * seq_len,
    device=device ,
    dtype=torch.int32,
)[None, :]


with sdpa_kernel(sdpa_backend):
    with context_parallel(
        cp_mesh, buffers=(input_ids, position_ids), buffer_seq_dims=(1, 1)
    ):
        cp_out = model(input_ids, position_ids=position_ids)
```

@dgiofre, I'm working on porting Ulysses sequence parallelism into HF Transformers in ArcticTraining https://github.com/snowflakedb/ArcticTraining/pull/45. Once this is done it should be easy to integrate into HF Trainer / Accelerate. It's working - just polishing rough edges and then will make it into a nice simple to use API. At the moment it's a huge pile up, not ready for others ;)

Hi, just an update: below is a minimal example of using torch's `context_parallel` to enable CP. The example gives you all the necessary information about how to use it. In short:

- We need to pass `position_ids` and `shift_labels` explicitly.
- loss computation has to be done manually instead of just getting it from `model(...).outputs`, see (#37993 and [this comment](https://github.com/huggingface/transformers/pull/36607#issuecomment-2781494656)).
- `loss.backward` has to be done within `context_parallel` ~~if `no_restore_buffers` is not specified (or not containing all tensors in `buffer`)~~


```python
import os

import torch
import torch.distributed as dist

from torch.distributed.tensor.experimental import context_parallel
from torch.distributed.device_mesh import init_device_mesh

from torch.nn.attention import sdpa_kernel, SDPBackend
from torch.nn.parallel import DistributedDataParallel as DDP

from transformers import AutoModelForCausalLM
from transformers.loss.loss_utils import ForCausalLMLoss


world_size = int(os.environ.get("WORLD_SIZE", "1"))
cp_mesh = init_device_mesh("cuda", (world_size,))
rank = torch.distributed.get_node_local_rank()

device = "cuda"
dtype = torch.bfloat16
sdpa_backend = SDPBackend.FLASH_ATTENTION

# prepare inputs
batch_size = 1
seq_len = 128

input_ids = torch.randint(low=8, high=64, size=(batch_size, seq_len), device=device)

ignore_index = -100
# When using CP, we need to use `shift_labels`
shift_labels = torch.nn.functional.pad(input_ids, (0, 1), value=ignore_index)
shift_labels = shift_labels[..., 1:].contiguous()

position_ids = torch.cumsum(torch.ones(size=input_ids.size(), dtype=input_ids.dtype, device=input_ids.device), dim=1) - 1

# sync input as they are created randomly
dist.broadcast(input_ids, src=0)
dist.broadcast(shift_labels, src=0)
dist.broadcast(position_ids, src=0)

# model and optimizer
repo_id = "Qwen/Qwen2.5-Coder-0.5B-Instruct"
model = AutoModelForCausalLM.from_pretrained(repo_id, torch_dtype=dtype, device_map=device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)

model.train()
model.zero_grad()
optimizer.zero_grad()

# For loss
vocab_size = model.config.vocab_size

# so training could be synced
model = DDP(model, device_ids=[rank])

# prepare for CP
buffers = (input_ids, shift_labels, position_ids)
buffer_seq_dims = (1, 1, 1)
# `no_restore_buffers=set(buffers)` is required if `loss.backward` is outside `context_parallel`.
# no_restore_buffers = set(buffers)
no_restore_buffers = None

# run with CP
with sdpa_kernel(sdpa_backend):
    with context_parallel(
        cp_mesh, buffers=buffers, buffer_seq_dims=buffer_seq_dims, no_restore_buffers=no_restore_buffers,
    ):
        outputs = model(input_ids, shift_labels=shift_labels, position_ids=position_ids)
        print(outputs.logits.shape)

        # So far we need to compute `loss` outside `model.forward` when using `shift_labels`
        # loss = outputs.loss
        loss = ForCausalLMLoss(logits=outputs.logits, labels=None, shift_labels=shift_labels, vocab_size=vocab_size)

        # ⚠️Might give incorrect outputs, see https://github.com/pytorch/pytorch/issues/149306#issuecomment-2862058203
        # (~~Don't trust this part: This could be outside `context_parallel` context if `no_restore_buffers` is specified~~)
        loss.backward()
        optimizer.step()

```

You rock, @ydshieh - thank you!

Do you guys want to fix this situation with `shift_labels` in HF Transformers so that `model(**batch)` will compute the loss with `shift_labels` automatically? 

I have already asked the Liger-Kernel team to do that and they did it in the latest version 0.5.9
https://github.com/linkedin/Liger-Kernel/pull/683

so with liger you just do:

```
loss = model(...).loss
```
and `shift_labels` will be used if passed. 

> Hi, just an update: below is a minimal example of using torch's `context_parallel` to enable CP. The example gives you all the necessary information about how to use it. In short:
> 
> * We need to pass `position_ids` and `shift_labels` explicitly.

Hi @ydshieh , amazing work! tks for your sharing!

I'd like to know whether CP can cache kv on each GPU locally?

@Uxito-Ada I haven't yet dive into this part yet. I am currently working on integrate `CP` to `transformers`, mostly for training purpose (where KV cache is irrelevant, right?).



HI @ydshieh, I have tried to adapt your example to 2D mesh: one dimension for context parallel an another one for data parallel and it fails. And looking into DDP code, I [found](https://github.com/pytorch/pytorch/blob/9968c854b6a38857d05fb1da7bea797fbf23c880/torch/nn/parallel/distributed.py#L682) that it can only handle 1D device mesh. 

On the other hand, I know that DeepSpeed handles different types of parallelism with [MPU](https://github.com/deepspeedai/DeepSpeed/blob/05818e90d916734d45fb4a23a85d17574161d0e2/deepspeed/runtime/engine.py#L204) argument (it can be created based on device mesh). 

FSDP handle with SequenceParallel strategy, I guess, but haven't dived into it in much details


I would like to give a hand with this feature, please let me know, if I can help in any way

@pavelgein In accelerate we're planning to support CP + FSDP within the next release (next week ish) with this [PR](https://github.com/huggingface/accelerate/pull/3604), we can further adapt this to also support raw DDP, but meanwhile you can take a look at it and try to use that.
As for your problem with 1D device mesh, you can probably recreate that using a mesh of `(dp, dp_shard, cp) = (dp_size, 1, cp_size)`, then creating a flat `dp_shard_cp` submesh, passing `mesh[(dp, dp_shard_cp)]` into `fully_shard` and passing `mesh[cp]` into `context_parallel`, I don't see a reason why you wouldn't want to shard the model across the context_parallel dimension, as this should get overlapped pretty much always and incurs no additional cost really.

In the image below you can see it very clearly. As we do a lot of kernel launches for *large* sequence lengths (Stream 7), in this time we're entirely bound by the compute of a single GPU. While this is happening, we can launch the prefetch for the next layer of FSDP (Stream 23), you can see that the gpu becomes "idle" of sorts only in short windows in stream 7 (compute stream) when we're waiting for another all-gather from context parallel to conclude (stream 31), this corresponds to the expected outcome, where we have a bit of idle time while we're gathering the shards from other ranks.

This idle time could be gone if it worked properly with all-to-all sharding (but it just doesn't bc of lengthy list of issues I'm not gonna go into)

<img width="1800" alt="Image" src="https://github.com/user-attachments/assets/a9502042-926d-4319-8118-42987534fe84" />
profiler images are from 8x H100 llama 8b 128k seq length 8 CP-size

TLDR: Always use CP + FSDP as you get the model sharding for free.

@S1ro1 Thanks a lot for the link and the explanation, they are really really useful. 

I don't think that anyone including me would really use CP without any model sharding (DeepSpeed or FSDP).

In #37877

https://github.com/huggingface/transformers/pull/37877/files#diff-707bda7ac3ef6a32e4b6b3460e1cf7d04df7a4465411812d71a473b01c6ca17c

there is an example script.

I tried (one month ago)

> CP_SIZE=2 TP_SIZE=2 torchrun --nproc_per_node=4 examples/3D_parallel.py

This doesn't use DDP and with some custom sync

Let me know if you have issues running this. (Back then, I got some strange issues with wandb, and also changed a few lines of code myself)

This can probably be closed, no?

Since this has been merged: https://github.com/huggingface/transformers/pull/40205

I'm now adding ALST/UlyssesSP sequence parallelism:
- https://github.com/huggingface/accelerate/pull/3817
- https://github.com/huggingface/transformers/pull/41832