Repository navigation
Key the tied-embedding TP rule on the model's input embedding attribute - #48779
Closed
Siddhant231xyz wants to merge 3 commits into
Closed
Siddhant231xyz wants to merge 3 commits into
Siddhant231xyz wants to merge 3 commits into
Conversation
PreTrainedConfig injected a hard-coded {'embed_tokens': 'embedding_rowwise'} into base_model_tp_plan whenever tie_word_embeddings is set, which produces an invalid plan entry for models whose embedding attribute has another name (GPT-2 style 'wte', GPT-NeoX 'embed_in', ...), as noted in the TODO.
Move the rule into ParallelInterface mixin plan initialization, keyed on the model's _input_embed_layer (default 'embed_tokens'), so the config stays declarative and models with differently named embeddings get vocab-parallel tied embeddings too.
No behaviour change for existing models: tests/tensor_parallel + test_tensor_parallel_mixin pass (16), tp_plan tests for llama/qwen3_moe pass, and a tied-embedding Llama on a 2-rank CPU mesh still shards 15/20 params with identical logits.
Related: huggingface#34789
Siddhant231xyz
force-pushed
the
tp-tied-embed-layer
branch
from
September 13, 2026 23:45
aea2ccd to
37370ad
Compare
This was referenced Sep 13, 2026
Siddhant231xyz
force-pushed
the
tp-tied-embed-layer
branch
from
September 13, 2026 23:52
37370ad to
1043552
Compare
Member
|
cc @3outeille indeed probably cleaner |
|
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. |
Member
|
lgtm |
3outeille
approved these changes
Sep 17, 2026
Contributor
CI recapDashboard: View test results in Grafana |
Collaborator
|
Thanks! #49570 just reworked this: the tied-embedding injection was removed from |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
PreTrainedConfig.__init__injects{"embed_tokens": "embedding_rowwise"}intobase_model_tp_planwhenevertie_word_embeddingsis set. For models whose embedding attribute has another name (GPT-2 stylewte, GPT-NeoXembed_in) that entry matches nothing, which blocks adding TP plans to those models (see the TODO above that block inconfiguration_utils.py).This moves the rule into
ParallelInterface.init_parallel_plans, keyed on the model's_input_embed_layer(default"embed_tokens", so existing models are unchanged). The config stays declarative.Part of #34789. Follow-up PRs adding TP plans for GPT-Neo and GPT-J depend on this.
Verification
tests/tensor_parallel+tests/test_tensor_parallel_mixin.py: 16 passed.test_tp_plan_matches_paramsfor llama and qwen3_moe pass.DistributedConfig(tp_size=2)) still shards 15/20 params, embedding included, with logits identical to the single-process run.The only CI failure is
OlmoModelTest::test_generate_with_static_cacheintests_generateshard 4, which is unrelated to this change and flaky: it also fails intermittently on unmodifiedmainlocally (1 of 3 runs).AI disclosure
I used an AI coding assistant (Claude Code) to draft this change and the verification script. I have read and understand every changed line, ran the tests listed above locally, and take responsibility for the code.
Before submitting
Tensor Parallelsupport for ALL models #34789 ("just open a PR and link it to this issue").Who can review?
@ArthurZucker @Cyrilvallez