Add support for context parallelism

[lewtun]

Feature request

Long context models like 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

[jquesnelle]

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

[dgiofre]

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

[ydshieh]

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

[ydshieh]

Hi, just some updates:

• Personally, I am trying using torch's 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)

[sfc-gh-sbekman]

@dgiofre, I'm working on porting Ulysses sequence parallelism into HF Transformers in ArcticTraining snowflakedb/ArcticTraining#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 ;)

[ydshieh]

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 (update loss computation in modeling code #37993 and this comment).
• loss.backward has to be done within context_parallel if no_restore_buffers is not specified (or not containing all tensors in buffer)

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()

[stas00]

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
linkedin/Liger-Kernel#683

so with liger you just do:

loss = model(...).loss

and shift_labels will be used if passed.

[Uxito-Ada]

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?

[ydshieh]

@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?).