[3outeille]

What does this PR do?

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

[Trainer] can now train a model loaded with DistributedConfig(cp_size=N), also combined with fsdp_size and tp_size.

• Data: the ranks of a cp group (and of a tp group) must get the same batch. Trainer gives Accelerate's dataloader a (dp_shard, cp, tp) device mesh, so the batches are split over the data parallel ranks only. Accelerate's own context parallelism is not used: the model shards its own inputs.
• Counting: num_items_in_batch, the loss scale, num_input_tokens_seen and the FLOPs count each batch once per data parallel group, not once per cp rank.
• Evaluation: the loss works. predict and metrics on logits raise, because each rank only holds the logits of its own slice.
• For a cp-sharded model, Trainer does not build the ParallelismConfig it builds for TP models, because that model's dp_size counts the cp ranks.

Tests

test_context_parallel in tests/trainer/distributed/test_trainer_distributed_transformers.py runs the same samples twice, once with CP and once in a single process. It compares the per-step loss, the grad norm, the eval loss and num_input_tokens_seen:

setup	ranks
cp=2	2
cp=2, gradient accumulation 2	2
fsdp=2 × cp=2, gradient accumulation 2	4
tp=2 × cp=2	4
cp=2, average_tokens_across_devices=False	2

All cases pass on 4× H100 (NCCL) and on CPU (gloo). The existing FSDP / HSDP Trainer tests pass too.

[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.