[3outeille]

What does this PR do?

Stack: #50493 → #50496 → #50497 → this PR

With context parallelism, a sequence length that is not a multiple of cp_size used to raise. Now _shard_inputs pads the full sequences at the end before splitting them:

• input_ids get pad_token_id (or 0), inputs_embeds get zeros. No real token attends to the padding, because it comes after them.
• position_ids continue those of the last document, so packed sequences and flash attention see no new document.
• The labels of the padding are -100. shift_labels and num_items_in_batch are built from the real tokens first, so the loss does not change.
• outputs.logits is this rank's slice of the padded sequence. On the last rank, the padding positions at the end are not meaningful (documented).

Tests

tests/context_parallel: length 7 with cp=2 (unpacked, packed, inputs_embeds) 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.

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

Pages changed in this PR:

• en/context_parallelism

[3outeille]

This was marked merged into cp-5-docs by the restack (that branch now contains its commits); it was not merged into main. Continued in #50515.