From 6ca781d6fc857e2c7de753ec88270cbcf26ad054 Mon Sep 17 00:00:00 2001 From: Ferdinand Mom Date: Tue, 10 Nov 2026 11:29:13 +0900 Subject: [PATCH] [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. --- docs/source/en/_toctree.yml | 2 + docs/source/en/context_parallelism.md | 133 ++++++++++++++++++++++++++ 2 files changed, 135 insertions(+) create mode 100644 docs/source/en/context_parallelism.md diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 2928fb8e301b..7b575753e631 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -231,6 +231,8 @@ title: DDP - local: fsdp title: FSDP2 + - local: context_parallelism + title: Context parallelism - local: deepspeed title: DeepSpeed ZeRO - local: deepspeed_alst diff --git a/docs/source/en/context_parallelism.md b/docs/source/en/context_parallelism.md new file mode 100644 index 000000000000..6704b937771c --- /dev/null +++ b/docs/source/en/context_parallelism.md @@ -0,0 +1,133 @@ + + +# Context parallelism + +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. + +```text + full sequence (8 tokens) cp_size = 2 +┌───┬───┬───┬───┬───┬───┬───┬───┐ ┌───┬───┬───┬───┐ +│ 0 │ 1 │ 2 │ 3 │ 4 │ 5 │ 6 │ 7 │ ───► │ 0 │ 1 │ 2 │ 3 │ rank 0 +└───┴───┴───┴───┴───┴───┴───┴───┘ ├───┼───┼───┼───┤ + │ 4 │ 5 │ 6 │ 7 │ rank 1 + └───┴───┴───┴───┘ +``` + +## Usage + +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. + +```py +import os + +import torch +from torch.utils.data import DataLoader, DistributedSampler + +from transformers import AutoModelForCausalLM +from transformers.distributed import DistributedConfig + +CP_SIZE = 4 +rank, world_size = int(os.environ["RANK"]), int(os.environ["WORLD_SIZE"]) + +distributed_config = DistributedConfig(fsdp_size=world_size // CP_SIZE, cp_size=CP_SIZE) +model = AutoModelForCausalLM.from_pretrained( + "Qwen/Qwen3-8B", + distributed_config=distributed_config, + attn_implementation="flash_attention_2", + dtype=torch.bfloat16, +) +optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5) + +# the ranks of a cp group take the same rows, so the sampler has one replica per group +sampler = DistributedSampler(dataset, num_replicas=world_size // CP_SIZE, rank=rank // CP_SIZE) +loader = DataLoader(dataset, batch_size=1, sampler=sampler) + +for batch in loader: + batch = {name: value.to(model.device) for name, value in batch.items()} + loss = model(**batch).loss # the loss of the full sequences, the same on every rank of the group + loss.backward() + optimizer.step() + optimizer.zero_grad() +``` + +```shell +torchrun --nproc-per-node 8 train_cp.py +``` + +- `loss` is the loss of the full sequences. The `logits` are only this rank's slice. +- 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. +- [FSDP2](./fsdp) shards the parameters over `fsdp_size * cp_size` ranks. +- Training only: `past_key_values` raise, so `generate` does not work. Requires `torch>=2.10`. + +## The context parallel plan + +`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. + +| style | module | what it does | +| :--- | :--- | :--- | +| `"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` | +| `"cp_ulysses"` | their attention modules (`layers.*.self_attn`) | run the attention function over the whole sequence, with Ulysses | + +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. + +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`. + +## Attention styles + +The attention style says how the attention modules see the whole sequence. `DistributedConfig(cp_plan={"model.layers.*.self_attn": ...})` picks another one. + +| style | how it works | limits | +| :--- | :--- | :--- | +| `"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 | + +To add one, subclass `ContextParallelAttention` and add it to `ALL_CP_STYLES` under its name. `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. + +```py +from transformers.distributed import DistributedConfig +from transformers.distributed.context_parallel import ALL_CP_STYLES, ContextParallelAttention + + +class MyAttention(ContextParallelAttention): + def check(self, module, cp_mesh): + ... + + def attention(self, cp_mesh, attention_function, module, query, key, value, attention_mask, **kwargs): + ... + + +ALL_CP_STYLES["cp_mine"] = MyAttention() +distributed_config = DistributedConfig(cp_size=4, cp_plan={"model.layers.*.self_attn": "cp_mine"}) +``` + +## Trainer + +[`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. + +## Combining with the other parallelisms + +- [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. +- [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`. +- Pipeline parallelism is not supported. + +## What is not supported + +- Attention that does not go through the [attention interface](./attention_interface) (`_supports_attention_backend = False`). +- Layers that mix tokens outside softmax attention: linear attention, recurrent layers, convolutions, sparse indexers. +- Heads other than a causal LM. +- Layers that take another full-sequence input than `hidden_states` and `position_embeddings`. + +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.