Repository navigation
Add tensor parallel support for GPT-J - #48783
Closed
Siddhant231xyz wants to merge 2 commits into
Closed
Siddhant231xyz wants to merge 2 commits into
Siddhant231xyz wants to merge 2 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
- Make GPTJAttention head-count agnostic (split/merge heads with -1) so column-sharded q/k/v projections work. - Add base_model_tp_plan to GPTJConfig (attn q/k/v colwise, out_proj rowwise, mlp fc_in colwise, fc_out rowwise) and a colwise_gather_output plan for lm_head. - Declare _input_embed_layer = 'wte' on GPTJModel so the tied input embedding is vocab-sharded. Verified on a 2-rank CPU/gloo mesh: sharded forward matches single-process forward (max abs diff 2e-7), 18/25 params sharded. tests/models/gptj: 112 passed. Related: huggingface#34789
1 of 5 tasks
Contributor
|
[For maintainers] Suggested jobs to run (before merge) run-slow: gptj |
Siddhant231xyz
force-pushed
the
tp-gptj
branch
from
September 13, 2026 23:52
c4cd5f2 to
41ef36e
Compare
Contributor
CI recapDashboard: View test results in Grafana |
Collaborator
|
Thanks for this! Main has moved under you: embedding sharding for TP was reworked in #49570 ( |
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?
GPTJAttentionsplit/merge heads with-1so column-sharded q/k/v projections work.GPTJConfig.base_model_tp_plan(attn q/k/v colwise,out_projrowwise, mlpfc_incolwise,fc_outrowwise) +lm_headplan.GPTJModel._input_embed_layer = "wte"so the tied input embedding is vocab-sharded.Part of #34789. Depends on #48779; this branch includes that commit and will be rebased once it lands.
Verification
tests/models/gptj: 112 passed.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").base_model_tp_planon the config.test_tp_plan_matches_paramsand the model test suite; 2-rank numerical check described above.Who can review?
@ArthurZucker @Cyrilvallez