[3outeille]

What does this PR do?

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

Wires the context parallel (CP) module of #50493 into transformers.distributed, with the same plan API as TP, EP and FSDP. Each rank of a cp group keeps a contiguous slice of every sequence, so activation memory per GPU shrinks with cp_size. The training loop does not change:

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-8B",
    attn_implementation="flash_attention_2",
    distributed_config=DistributedConfig(fsdp_size=4, cp_size=2),  # cp_plan optional
)
loss = model(**batch).loss  # full sequences, the same batch on the ranks of a cp group
loss.backward()

What this PR brings

• DistributedConfig(cp_size, cp_plan). dp_size now includes cp_size: the ranks of a cp group share a batch but process different tokens, so FSDP2 reduces gradients over them too.
• PreTrainedConfig.base_model_cp_plan defaults to {"layers.*": "cp_split_sequence", "layers.*.self_attn": "cp_ulysses"}, like base_model_fsdp_plan.
  • The plan works like the others: model.cp_plan, overrides with cp_plan (another attention algorithm is cp_plan={"model.layers.*.self_attn": "cp_ring"}, a registered style), and the "no cp_plan" error, which is also raised when the plan matches nothing or lacks a split or an attention style.
  • 4 models whose layers take other full-sequence inputs set the plan to None: gemma3n, gemma4, beam, ministral3 (gemma4_unified and diffusion_gemma inherit it from gemma4). longcat_flash declares its own plan.
• maybe_distribute_model calls apply_context_parallelism ([distributed] Context parallel module, with Ulysses #50493) with the resolved plan, before FSDP2, like the other parallelisms.
• Mesh (pp, dp_replicate, fsdp, cp, tp). FSDP2 shards over fsdp × cp, an axis of the root mesh (pp, dp_replicate, fsdp_cp, tp) that fsdp and cp are an _unflatten view of, so that it composes with tp later. CP needs torch>=2.10.
• Refused for now:
  • pp, tp or ep together with cp (TP and EP come next in the stack);
  • Trainer.

What #50493 brings

The module itself, context_parallel.py: ContextParallelLayer and the styles (cp_split_sequence, the attention styles with cp_ulysses, ModelContextParallel with the checks, the cache refusal and the loss) and apply_context_parallelism. This PR adds no code to that file.

Tests

tests/context_parallel, CPU (gloo), through from_pretrained(distributed_config=...). The loss, this rank's logits and every gradient match an unsharded run of the same weights:

setup	ranks
cp=2: plain, packed, inputs_embeds, left and right padding	2
cp=2: Mistral with a sliding window (plain, packed), Gemma3 with sliding and full layers	2
cp=2 × fsdp=2, and dp_replicate=2 × cp=2	4

Also covered: a custom attention style named in cp_plan, the config refusals, the plan resolution and its errors, and save_pretrained (gathered and DCP).

GPU check (2× H100, Qwen3-0.6B, bf16)

With sdpa and the FA3 hub kernel, CP and an unsharded run give the same loss and logits. This holds for packed batches, padded batches and sliding windows. The gradients match within bf16 rounding (at most 1.5% relative).

Peak memory per GPU and step time (forward + backward + AdamW), FA3:

tokens per sequence	cp=1 (1 GPU)	cp=2 (2 GPUs)
16k	62.0 GiB, 502 ms	32.5 GiB, 319 ms
32k	out of memory	63.2 GiB, 737 ms

[HuggingFaceDocBuilderDev]

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.

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: beam, diffusion_gemma, gemma3n, gemma4, gemma4_unified, longcat_flash, ministral3