Repository navigation
Conversation
|
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. Pages changed in this PR: |
… model `ContextParallelTesterMixin` joins the bases of `CausalLMModelTest`, next to the FSDP and tensor parallel mixins, so every causal language model test class runs its tester's tiny config on 2 ranks: `cp_size=2` over the whole batch, with the loss, this rank's slice of the logits and every gradient compared against the unsharded run. The packed variant restarts `position_ids` mid-sequence, which the mask of every model then has to rebuild. A tester config is only changed when the context parallel checks refuse it and the change is harmless: the sliding-window layers, whose mask `sdpa` cannot rebuild for the whole sequence. The change is verified by building the model from the config a checkpoint reloads, which is the config the workers get, and a config class that validates its own fields can refuse it. A model the checks still refuse is loaded with `cp_size=2` on both ranks and has to raise `ValueError`; the test then skips with the reason. The allow-list starts with 7 model types that run today. Set `CONTEXT_PARALLEL_TEST_MODEL_TYPES = None` to run every model.
This was referenced Nov 9, 2026
FSDP2 composes the mesh it shards a parameter over with the mesh the parameter is already sharded on, and torch only composes meshes built from the same root mesh. With context parallelism the sharding axis is fsdp x cp, so it is one axis of the dense root mesh and fsdp and cp are its unflattened view: the ranks, their order and the groups are unchanged, and a tp or ep mesh now composes with it.
The tp plan keeps num_heads / tp_size heads on each rank and Ulysses splits those across the cp group, so the heads and key/value heads must be divisible by cp_size * tp_size. Only pipeline and expert parallelism are still refused with context parallelism.
tp_size x cp_size on the tiny Llama: loss, this rank slice of the logits and every gradient match the unsharded run, with and without FSDP, packed and unpacked, and the gathered save reloads the reference weights. The workers now also drop the DTensor dispatch cache between tests: it is keyed by mesh, and a mesh rebuilt from the same ranks compares equal to the destroyed one.
Under cp the ranks of an ep group hold the slices of several batches, which token dispatch already handles: each rank sends the tokens of its own slice. The expert gradient divides by dp_size, which includes cp_size, so the reduced expert gradient is still the gradient of the full-sequence loss.
A tiny Qwen3 MoE against the same weights on full sequences: loss, this rank's logits slice and every gradient, with token dispatch (cp2 x ep2, cp2 x fsdp2 x ep4, cp2 x fsdp2 x ep2, cp2 x tp2 x ep4, cp2 x hsdp2 x ep2) and with the all-reduce plan (tp2 x cp2 x ep2).
Each rank has to hold the same number of positions, so the pre-hook pads the full sequences at their end: input_ids with the padding token, inputs_embeds with zeros, position_ids continuing after the last one of each row (the padding joins the last document) and labels with -100. num_items_in_batch keeps counting the real tokens.
…Ulysses context parallelism The all-to-all gathers the sequence before every attention call, so the mask of the whole sequence can be built there from transformers' own mask primitives: causal, packed documents, padding, sliding window and chunked attention. - `_shard_inputs` keeps the full 2D padding mask on the implementation and passes `attention_mask=None` to the model. - `UlyssesCP.attention` rebuilds the `sdpa` mask for the whole sequence, or returns `None` when it is plain causal so that the kernel relies on `is_causal`. - flash kernels get the full 2D mask, which unpads and rebuilds the sequence boundaries. - `UlyssesCP.check` only refuses sliding-window layers whose window is not visible at the attention call, and `chunked_attention` layers with a flash kernel.
…h context parallelism - the mask rebuilt over the whole sequence is compared with the one the model builds on an unsharded run, for causal, packed, padded, sliding-window and chunked attention; - right and left padding batches, a sliding-window model, a mix of sliding and full layers, and a chunked model are compared with the same weights on full sequences (loss, this rank's logits, every gradient); - the flash padding test runs on 2 GPUs with `require_flash_attn`; - the checks now accept sliding windows and chunked attention with `sdpa`, and refuse what they cannot rebuild.
Member
Author
|
Restacked: the docs page moves to the end of the stack, after Trainer support. It will be reopened as a new PR on top of it (GitHub does not allow changing the base of this one). |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Stack: #50493 → #50496 → this PR (tests on every causal LM: #50499)
Adds a "Context parallelism" page under Distributed training, next to FSDP2: how to train with
DistributedConfig(cp_size=N), how to feed the same batch to the ranks of a cp group, thecp_implementationvalues and how to register your own, what is not supported, and how the names relate tosp_size/cp_sizeelsewhere and to thesequence_parallelTP style.