Add base_model_tp_plan to OlmoeConfig

[dacorvo]

Fixes #44677

Summary

• Add base_model_tp_plan to OlmoeConfig, enabling from_pretrained(tp_plan="auto") for OLMoE models
• Add TensorParallelTesterMixin to OLMoE tests for TP validation coverage
• Uses "colwise" for q_norm and k_norm because OLMoE applies these norms after the q/k projections — norm weight dimensions must match the sharded projection output

Design note

Qwen3-MoE uses "replicated_with_grad_allreduce" for its q/k norms, but that only works because Qwen3 applies norms before the projections (on the full hidden state). OLMoE's architecture applies norms after projections, so the norm weights must be sharded the same way as the projection output — hence "colwise".

Test plan

• python -m pytest tests/models/olmoe/test_modeling_olmoe.py::OlmoeModelTest::test_tp_plan_matches_params -xvs — passes

AI disclosure

This PR was developed with AI assistance (Claude). All changes reviewed and validated by a human contributor.

🤖 Generated with Claude Code

[Copilot]

Pull request overview

Adds a default tensor-parallel (TP) sharding plan to OlmoeConfig so OLMoE models can be loaded with from_pretrained(tp_plan="auto"), and wires OLMoE’s modeling tests into the shared TP test mixin.

Changes:

• Define base_model_tp_plan on OlmoeConfig.
• Add TensorParallelTesterMixin to OLMoE’s model test class.
• Document why q/k norms need special handling for TP.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
src/transformers/models/olmoe/configuration_olmoe.py	Introduces a default TP plan mapping for OLMoE modules.
tests/models/olmoe/test_modeling_olmoe.py	Adds TP test mixin to validate TP behavior for OLMoE.

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

[Cyrilvallez]

Hey @dacorvo! YOu'll need a small rebase as we recently migrated configs to dataclasses! Apart from that, did you test the tp plan on real checkpoints to ensure correctness by any chance?

[dacorvo]

@Cyrilvallez I rebased the branch.
I tested two models with this TP Plan against device_map='auto', with three different prompts.

Model	TP=2	TP=4	TP=8
OLMoE-1B-7B-0924	MATCH	MATCH	MATCH
OLMoE-1B-7B-0125	MATCH	MATCH	MATCH

As a sanity, I also tested the wrong initial plan I submitted, and verified it failed for every prompt 😅 (divergence after a few tokens).

[Cyrilvallez]

@bot /style

[github-actions]

Style fix bot fixed some files and pushed the changes.

[Cyrilvallez]

Yes, that's why I remembered from olmo haha, the norm on the full last dim prevents nice TP 🥲
LGTM now!

[Cyrilvallez]

Actually, did not see this but let's not do that 😅 Let's put it directly in the class above!

[github-actions]

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

run-slow: olmoe