Skip to content

Commit 0265640

Browse files
committed
[docs] Add a context parallelism page
`docs/source/en/context_parallelism.md` documents the plan API: the `base_model_cp_plan` default and the two styles, the `cp_plan` override, `cp_implementation` and how to register your own, the model's own masks (packing, padding, sliding windows, chunked attention, custom mask functions), the cost, tensor and expert parallelism, the automatic padding, `Trainer`, and what is refused. The page joins the performance section of the table of contents after FSDP2.
1 parent 128e5d7 commit 0265640

2 files changed

Lines changed: 136 additions & 0 deletions

File tree

‎docs/source/en/_toctree.yml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -231,6 +231,8 @@
231231
title: DDP
232232
- local: fsdp
233233
title: FSDP2
234+
- local: context_parallelism
235+
title: Context parallelism
234236
- local: deepspeed
235237
title: DeepSpeed ZeRO
236238
- local: deepspeed_alst
Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
1+
<!--Copyright 2026 The HuggingFace Team. All rights reserved.
2+
3+
Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
4+
the License. You may obtain a copy of the License at
5+
6+
http://www.apache.org/licenses/LICENSE-2.0
7+
8+
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
9+
an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
10+
specific language governing permissions and limitations under the License.
11+
12+
⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
13+
rendered properly in your Markdown viewer.
14+
15+
-->
16+
17+
# Context parallelism
18+
19+
Context parallelism (CP) splits every sequence across the ranks of a cp group: each rank keeps one contiguous slice of it. Activation memory per GPU shrinks with `cp_size`, so longer sequences fit on the same GPUs. Attention still needs the whole sequence, which the attention style of the plan gives it back.
20+
21+
```text
22+
full sequence (8 tokens) cp_size = 2
23+
┌───┬───┬───┬───┬───┬───┬───┬───┐ ┌───┬───┬───┬───┐
24+
│ 0 │ 1 │ 2 │ 3 │ 4 │ 5 │ 6 │ 7 │ ───► │ 0 │ 1 │ 2 │ 3 │ rank 0
25+
└───┴───┴───┴───┴───┴───┴───┴───┘ ├───┼───┼───┼───┤
26+
│ 4 │ 5 │ 6 │ 7 │ rank 1
27+
└───┴───┴───┴───┘
28+
```
29+
30+
## Usage
31+
32+
Set `cp_size` to the number of ranks that split a sequence. Every rank of a cp group receives the same full batch, and the model splits it. The `world_size // cp_size` groups run different batches.
33+
34+
```py
35+
import os
36+
37+
import torch
38+
from torch.utils.data import DataLoader, DistributedSampler
39+
40+
from transformers import AutoModelForCausalLM
41+
from transformers.distributed import DistributedConfig
42+
43+
CP_SIZE = 4
44+
rank, world_size = int(os.environ["RANK"]), int(os.environ["WORLD_SIZE"])
45+
46+
distributed_config = DistributedConfig(fsdp_size=world_size // CP_SIZE, cp_size=CP_SIZE)
47+
model = AutoModelForCausalLM.from_pretrained(
48+
"Qwen/Qwen3-8B",
49+
distributed_config=distributed_config,
50+
attn_implementation="flash_attention_2",
51+
dtype=torch.bfloat16,
52+
)
53+
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5)
54+
55+
# the ranks of a cp group take the same rows, so the sampler has one replica per group
56+
sampler = DistributedSampler(dataset, num_replicas=world_size // CP_SIZE, rank=rank // CP_SIZE)
57+
loader = DataLoader(dataset, batch_size=1, sampler=sampler)
58+
59+
for batch in loader:
60+
batch = {name: value.to(model.device) for name, value in batch.items()}
61+
loss = model(**batch).loss # the loss of the full sequences, the same on every rank of the group
62+
loss.backward()
63+
optimizer.step()
64+
optimizer.zero_grad()
65+
```
66+
67+
```shell
68+
torchrun --nproc-per-node 8 train_cp.py
69+
```
70+
71+
- `loss` is the loss of the full sequences. The `logits` are only this rank's slice.
72+
- A sequence whose length `cp_size` does not divide is padded at its end. The padding gets labels `-100`, and no real token attends to it.
73+
- [FSDP2](./fsdp) shards the parameters over `fsdp_size * cp_size` ranks.
74+
- Training only: `past_key_values` raise, so `generate` does not work. Requires `torch>=2.10`.
75+
76+
## The context parallel plan
77+
78+
`PreTrainedConfig.base_model_cp_plan` holds the default plan, and `DistributedConfig(cp_plan=...)` overrides its rules, like `tp_plan` does. The default plan pairs two kinds of styles: one on the decoder layers, one on their attention modules.
79+
80+
| style | module | what it does |
81+
| :--- | :--- | :--- |
82+
| `"cp_split_sequence"` | the decoder layers (`layers.*`) | run on this rank's slice: the first layer cuts `hidden_states`, and every layer cuts `position_embeddings` |
83+
| `"cp_ulysses"` | their attention modules (`layers.*.self_attn`) | run the attention function over the whole sequence, with Ulysses |
84+
85+
The model builds the embeddings, the rotary tables and the attention mask on the whole sequence, before the first decoder layer cuts it. So every mask the model supports works unchanged: packing with `position_ids`, a padding `attention_mask`, sliding windows and chunked attention.
86+
87+
A model whose layers have other names declares its own `base_model_cp_plan`. A model that context parallelism cannot run sets it to `None`.
88+
89+
## Attention styles
90+
91+
The attention style says how the attention modules see the whole sequence. `DistributedConfig(cp_plan={"model.layers.*.self_attn": ...})` picks another one.
92+
93+
| style | how it works | limits |
94+
| :--- | :--- | :--- |
95+
| `"cp_ulysses"` (default) | an all-to-all gives each rank the whole sequence for `num_heads / cp_size` heads, and a second one goes back | `cp_size * tp_size` must divide the attention heads and the key/value heads; `eager` and `flex_attention` are refused |
96+
97+
To add one, subclass `ContextParallelAttention` and register it in `ALL_PARALLEL_STYLES`, like a tensor parallel style. `check` raises a `ValueError` for a module it cannot run. `attention` runs `attention_function` on this rank's slice of `query`, `key` and `value`, with the mask the model built for the whole sequence, and returns this rank's output and the attention weights.
98+
99+
```py
100+
from transformers.distributed import DistributedConfig
101+
from transformers.distributed.context_parallel import ContextParallelAttention
102+
from transformers.distributed.parallel_utils import ALL_PARALLEL_STYLES
103+
104+
105+
class MyAttention(ContextParallelAttention):
106+
def check(self, module, cp_mesh):
107+
...
108+
109+
def attention(self, cp_mesh, attention_function, module, query, key, value, attention_mask, **kwargs):
110+
...
111+
112+
113+
ALL_PARALLEL_STYLES.register("cp_mine", MyAttention())
114+
distributed_config = DistributedConfig(cp_size=4, cp_plan={"model.layers.*.self_attn": "cp_mine"})
115+
```
116+
117+
## Trainer
118+
119+
[`Trainer`] trains a model loaded with `DistributedConfig(cp_size=N)` like a model on a single device: its dataloader gives the same batch to every rank of a cp group, and the loss and the gradients match a single-device run with the same total batch size. Evaluation only supports the loss (`prediction_loss_only=True`), because a rank only holds the logits of its slice.
120+
121+
## Combining with the other parallelisms
122+
123+
- [Tensor parallelism](./tensor_parallelism): each rank runs Ulysses on its `num_heads / tp_size` heads. The ranks of a tp group hold the same batch and the same slice.
124+
- [Expert parallelism](./expert_parallelism) with token dispatch (`"ep_dispatch_experts"`): each rank dispatches the tokens of its own slice. `ep_size` must divide `fsdp_size * cp_size * tp_size`.
125+
- Pipeline parallelism is not supported.
126+
127+
## What is not supported
128+
129+
- Attention that does not go through the [attention interface](./attention_interface) (`_supports_attention_backend = False`).
130+
- Layers that mix tokens outside softmax attention: linear attention, recurrent layers, convolutions, sparse indexers.
131+
- Heads other than a causal LM.
132+
- Layers that take another full-sequence input than `hidden_states` and `position_embeddings`.
133+
134+
Ring attention is not implemented. The `"sequence_parallel"` rule of a tensor parallel plan is a different thing: it splits the activations of the normalization layers between the TP ranks.

0 commit comments

Comments
 (0)