Skip to content

[distributed] Context parallelism: pad the sequences to a multiple of cp_size - #50515

Closed
3outeille wants to merge 8 commits into
cp-9-ep-x-cpfrom
cp-7-auto-pad
Closed

[distributed] Context parallelism: pad the sequences to a multiple of cp_size#50515
3outeille wants to merge 8 commits into
cp-9-ep-x-cpfrom
cp-7-auto-pad

Conversation

@3outeille

@3outeille 3outeille commented Nov 9, 2026 •
edited
Loading

Copy link
Copy Markdown
Member

What does this PR do?

Stack: #50493 → #50496 → #50511 → #50514 → this PR → #50499 → #50521 → #50522

With context parallelism, a sequence length that cp_size does not divide used to raise. Now ModelContextParallel, the style of the model itself (#50493), pads the full sequences at the end, before the model builds its masks:

  • input_ids get pad_token_id (or 0), and inputs_embeds get zeros. No real token attends to the padding, because it comes after them.
  • position_ids, if given, continue those of the last document.
  • A 2D attention_mask, if given, gets zeros.
  • labels and shift_labels get -100, so num_items_in_batch still counts the real tokens only.
  • outputs.logits is this rank's slice of the padded sequence.

Tests

tests/context_parallel, CPU (gloo):

  • length 7 with cp=2 (plain, packed, inputs_embeds, with a padding mask) and length 5 with cp=4;
  • the loss, every gradient, and the logits at the real positions match an unsharded run on the unpadded sequences;
  • a single-process test of what ModelContextParallel pads.

This was referenced Nov 9, 2026
Closed
Merged
Draft
Draft
Draft
Draft
Draft
@3outeille
3outeille added this pull request to stack #50517 November 9, 2026 14:29
@3outeille
3outeille removed this pull request from stack #50517 November 9, 2026 14:30
@3outeille
3outeille added this pull request to stack #50518 November 9, 2026 14:30
@HuggingFaceDocBuilderDev

HuggingFaceDocBuilderDev commented Nov 9, 2026

Copy link
Copy Markdown

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@3outeille
3outeille removed this pull request from stack #50518 November 9, 2026 14:36
@3outeille
3outeille added this pull request to stack #50520 November 9, 2026 14:37
This was referenced Nov 9, 2026
Draft
Draft
@3outeille
3outeille removed this pull request from stack #50520 November 10, 2026 02:50
@3outeille
3outeille added this pull request to stack #50546 November 10, 2026 02:51
@3outeille
3outeille force-pushed the cp-7-auto-pad branch 2 times, most recently from 811cfea to dca9b1c Compare November 12, 2026 06:19
@3outeille
[distributed] Add context parallelism: DistributedConfig(cp_size=N)
8442714 
`cp_size` splits every sequence across the ranks of a `cp` mesh, which get the same batch and
keep a contiguous slice each. The plan follows the other parallelisms: a `base_model_cp_plan`
default in `PreTrainedConfig`, a `cp_plan` override merged and expanded like `tp_plan`, and
`apply_context_parallelism` installing the plan's styles after TP/EP and before FSDP2.

The model keeps taking whole sequences: the first decoder layer cuts `hidden_states`, every
layer cuts the rotary embeddings, the model-built mask and positions stay whole, and the loss is
the one of the whole sequences. FSDP2 shards over `fsdp` and `cp` together (`fsdp_cp`), and
`dp_size` counts the cp ranks. `tp_size`, `pp_size` and `ep_size` are refused with it for now,
and so is the `Trainer`, whose dataloader would split the batch across a cp group.
@3outeille
[distributed] Test context parallelism end to end
e2b848d 
Each case compares a 2-rank run against the same weights on the full sequences: loss, this
rank's slice of the logits and every gradient, for plain, packed and `inputs_embeds` batches,
with `fsdp_size=2` and `dp_replicate_size=2`, and with left and right padding through a 2D
`attention_mask`. Plus a sliding-window Mistral and a Gemma3 mix of sliding and full layers with
`sdpa`, a custom and an unknown `cp_implementation`, the refusals, and the save / distributed
checkpoint round trips.
@3outeille
[distributed] Context parallelism with tensor parallelism
3022544 
`fsdp_cp` is now an axis of a dense root mesh `(pp, dp_replicate, fsdp_cp, tp)`, and the
five dense axes are an unflattened view of it: FSDP2 composes the mesh it shards over with the
mesh the parameters are already sharded on, and torch only composes meshes built from the same
root. Requires `torch>=2.10` (`DeviceMesh._unflatten`), refused in `DistributedConfig`.

Under TP this rank holds `num_heads / tp_size` heads and Ulysses splits those across the cp
group, so `cp_size * tp_size` must divide the attention and key/value heads.
@3outeille
[distributed] Test tensor parallelism with context parallelism
29467a6 
`tp_size=2 x cp_size=2` on four ranks, plain and packed, against the same weights on the full
sequences, plus the gathered save, the head divisibility check under TP, and a test that the two
plans compose: the attention module whose call CP routes is the one whose projections TP sharded.
@3outeille
[distributed] Context parallelism with expert parallelism
9c88067 
Each rank of a cp group dispatches the tokens of its own slice only, so no token reaches an expert
twice and `efsdp * ep` still covers `fsdp * cp * tp`. The expert gradient divide factor already
counts the cp ranks through `dp_size`, so only the explanatory comment and the mesh docstring
change; `cp_size` with `ep_size` is no longer refused.
@3outeille
[distributed] Test expert parallelism with context parallelism
aba53d9 
The tiny Qwen3-MoE against the same weights on the full sequences, with token dispatch on
`cp2 x ep2`, `cp2 x fsdp2 x ep4`, `cp2 x fsdp2 x ep2` (efsdp), `cp2 x tp2 x ep4` and
`cp2 x hsdp2 x ep2`, plus the all-reduce plan with `ep_size=tp_size=2` where the ep group holds
one cp slice. Every case stays within four ranks.
@3outeille
[distributed] Context parallelism: pad the sequences to a multiple of…
847de83 
… cp_size

The root pre-hook pads every sequence-shaped input at its end instead of raising: `input_ids` with
the model's `pad_token_id` (0 without one), `inputs_embeds` with zeros, `position_ids` with the
positions that continue each row's last document, a 2D `attention_mask` with zeros, and `labels`
or a caller's `shift_labels` with -100. The padding happens before the model builds its mask and
its rotary embeddings, so nothing on the padding is attended to or trained on, and the loss
wrapper counts `num_items_in_batch` on the real tokens only.
@3outeille
[distributed] Test the padding of a sequence that cp_size does not di…
d5bdb5c 
…vide

A length of 7 over 2 ranks (plain, packed, `inputs_embeds` and through a 2D padding mask) and a
length of 5 over 4 ranks, compared on the real positions of the logits only, plus a single-process
test of the hook's padding of every input, of a positional `input_ids` and of a config without a
`pad_token_id`.
3outeille added a commit that referenced this pull request Nov 12, 2026
@3outeille
[distributed] Context parallelism: the loss on its own seam, the chec…
7c2d87a 
…ks in `apply_context_parallelism`, pad the sequences

- `context_parallel_loss` is a module-level function that `apply_context_parallelism` sets as
  `model.loss_function`, the seam transformers has for a model's loss; `ModelContextParallel` is
  about the forward only.
- The model checks open `apply_context_parallelism` instead of a `check` method.
- A sequence length that `cp_size` does not divide is padded at its end to a multiple of it, in
  `ModelContextParallel.transform_inputs_pre_forward`, before the model builds its masks:
  `input_ids` get `pad_token_id` (or 0), `inputs_embeds` zeros, `attention_mask` zeros, `labels`
  and `shift_labels` -100, and `position_ids` continue those of the last document. The logits are
  this rank's slice of the padded sequence; the loss counts the real tokens only. Was #50515.
@3outeille
3outeille removed this pull request from stack #50546 November 12, 2026 10:09
@3outeille

3outeille commented Nov 12, 2026

Copy link
Copy Markdown
Member Author

Folded into #50493 (the padding in ModelContextParallel.transform_inputs_pre_forward, with its single-process tests) and #50496 (the end-to-end tests of padded sequences). The stack is now #50493 → #50496 → #50511 → #50514 → #50499 → #50521 → #50522.

@3outeille 3outeille closed this Nov 12, 2026
@3outeille
3outeille deleted the cp-7-auto-pad branch November 12, 2026 10:10
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@3outeille @HuggingFaceDocBuilderDev