Skip to content

[distributed] Run the context parallel tests on every causal LM - #50499

Draft
3outeille wants to merge 1 commit into
cp-9-ep-x-cpfrom
cp-4-tester-mixin
Draft

3outeille wants to merge 1 commit into
cp-9-ep-x-cpfrom
cp-4-tester-mixin

Conversation

@3outeille

@3outeille 3outeille commented Nov 9, 2026 •

Copy link
Copy Markdown
Member

What does this PR do?

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

Adds ContextParallelTesterMixin to CausalLMModelTest, so every causal LM gets two tests, test_cp_matches_unsharded and test_cp_matches_unsharded_packed. Each test runs on 2 CPU ranks with cp_size=2 and checks that the loss, the logits of each rank's slice, and every gradient match an unsharded run of the same weights.

  • The tester configs run unchanged. The only change is output_router_logits=False, because the MoE auxiliary loss is not a sum over tokens.
  • If CP refuses a model, the test checks that loading with cp_size=2 raises, then skips with the reason.
  • An allow list (CONTEXT_PARALLEL_TEST_MODEL_TYPES) limits CI to 7 model types, like the FSDP mixin. The tests carry is_fsdp_test.

Tests

With the allow list off, over the 95 causal LM test classes:

result classes
pass (plain and packed) 54
refused with a clear error 41
fail 0

The refusals: attention outside the attention interface, eager-only models, linear attention, recurrent or conv layers, sparse indexers, a single key/value head, layers the default plan does not name (GPT-2, jetmoe, ...), and the models that set base_model_cp_plan = None in #50496 (gemma3n, gemma4 and gemma4_unified, which inherits it, beam, ministral3).

@HuggingFaceDocBuilderDev

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 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
@3outeille
3outeille removed this pull request from stack #50518 November 9, 2026 14:36
@3outeille
3outeille changed the base branch from cp-3-cp-size to cp-8-masks November 9, 2026 14:37
@3outeille
3outeille added this pull request to stack #50520 November 9, 2026 14:37
@3outeille
3outeille removed this pull request from stack #50520 November 10, 2026 02:50
@3outeille
3outeille changed the base branch from cp-8-masks to cp-7-auto-pad 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-4-tester-mixin branch 2 times, most recently from 6df787f to ebf1bd8 Compare November 12, 2026 06:19
@3outeille
3outeille force-pushed the cp-4-tester-mixin branch 2 times, most recently from f9e3c95 to 1ded7ce Compare November 12, 2026 10:09
@3outeille
3outeille removed this pull request from stack #50546 November 12, 2026 10:09
@3outeille
3outeille changed the base branch from cp-7-auto-pad to cp-9-ep-x-cp November 12, 2026 10:09
@3outeille
3outeille added this pull request to stack #50648 November 12, 2026 10:10
… 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 auxiliary load-balancing loss of the MoE models is turned off. The change is
verified by building the model from the config a checkpoint reloads, which is the config the
workers get. 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 causal language model.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants