[IlyasMoutawwakil]

What does this PR do?

Gives embeddings their own tensor-parallel style, shards every model's input embeddings by default, and lets layers shard widths that don't divide by the TP size. ColwiseParallel and RowwiseParallel now only handle linear layers.

Warning

Breaking: custom TP plans must use embedding_rowwise (vocabulary sharding) or embedding_colwise (hidden-dimension sharding) for embeddings. Linear styles (colwise* / rowwise*) now raise on embeddings. Eligible models also shard embed_tokens by default, including untied models; explicit plan entries take precedence.

Embeddings get their own style

ColwiseParallel and RowwiseParallel used to branch on isinstance(module, nn.Embedding) in their parameter sharding, input handling and validation. Those branches move to a dedicated EmbeddingParallel(shard_dim), which keeps the token ids replicated:

Plan entry	Style	Shards	Output
embedding_rowwise	EmbeddingParallel(shard_dim=0)	the vocabulary (vocab-parallel, the usual choice: it shares its layout with a tied, vocab-sharded lm_head)	partial lookups are all-reduced
embedding_colwise (new)	EmbeddingParallel(shard_dim=1)	the hidden dim (for large untied tables, e.g. Qwen4-Exp's ngram_embedding)	slices are all-gathered

ColwiseParallel and RowwiseParallel are now linear-only: applying them to a module without out_features / in_features raises a clear error instead of silently taking the embedding path.

Plans that targeted embeddings with linear styles move to the embedding ones:

• Doge: mlp.down_embed / mlp.up_embed → embedding_rowwise (was rowwise_split_input).
• Qwen4-Exp: ple.ple_embedding.ngram_embedding → embedding_colwise (was colwise_gather_output).
• Gemma4 vision: the plan targets the .linear inside the clippable-linear wrappers, which have no in_features/out_features themselves.

Input embeddings are sharded by default

Almost every model replicated its input embeddings while sharding its lm_head. A replicated embedding costs vocab x hidden on every rank (~1 GiB for Llama-3-8B in bf16, ~2 GiB for Llama-3-70B), and a tied model's vocab-sharded lm_head doesn't match the replicated embedding it shares a weight with.

The config used to add "embed_tokens": "embedding_rowwise" only for tied models. That default moves to init_parallel_plans, without the tie_word_embeddings condition: a base model with a TP plan and an nn.Embedding named embed_tokens shards it on the vocabulary, unless its plan already names it. Untied models (most large checkpoints) benefit the most. Composite models (VLMs, encoder-decoders) pick the entry up from their text sub-model's plan, e.g. model.language_model.embed_tokens. Vision towers and models without an embed_tokens embedding (GPT-NeoX's embed_in) get nothing, so no plan carries a key that matches no module.

No model plan is edited for this, and Inkling's and MuseGlimmer's explicit entries are removed as redundant. Models whose only plan is a class-level lm_head entry (AFMoE, DBRX, GPT-OSS, HunYuan, LFM2, PhiMoE) are unchanged: they need full plans first, in a follow-up.

Divisibility

Real vocab sizes often don't divide by the TP size (GPT-2's 50257, Whisper's 51865, GotOcr2's), and colwise_gather_output used to reject them outright. distribute_tensor already splits a parameter unevenly (the last rank holds a shorter slice), and DTensor ops track the real shape. The one gap was the local fast path (inference and quantized kernels): DTensor.from_local assumes every rank holds an equal slice. A small TensorParallelLayer._from_local_last_dim_shard helper wraps the local tensor with its real width, taken from the module's out_features / in_features (or embedding_dim).

Style	Width that doesn't divide by the number of ranks
colwise*, rowwise*	supported
embedding_rowwise, embedding_colwise	supported
packed_colwise, packed_rowwise (MoE experts' fused gate_up_proj)	supported: the strided split and the checkpoint loader both give each rank matching ceil-sized slices of gate and up. Verified with a TP forward and backward of a tiny Qwen3-MoE whose expert intermediate size splits 8/8/8/8/4 over 5 ranks.
expert parallelism	rejected, as before (num_experts must divide by ep_size)

Cost. Divisible sizes run exactly as before. An uneven width only costs anything where it is gathered (a colwise_gather_output lm_head, embedding_colwise), where DTensor pads and unpads around the all-gather; layers that reduce a full-size Partial (attention, MLP, packed experts, vocab-parallel embeddings) are unaffected. Gathering bf16 logits over 2 GPUs (NCCL), vocab 128257 vs 128256:

Logits gathered	Even	Uneven
1 token (a decode step)	530 µs	676 µs
2048 tokens	82.6 ms	101.3 ms

generate only gathers the last token's logits, so that's about 0.15 ms per decode step.

Tests

• New test_tp_plan_shards_input_embeddings in the TP mixin: a static check that a model whose TP plan covers more than the lm_head shards its nn.Embedding input embeddings. With test_tp_plan_matches_params, it passes for every model (258 passed).
• The TP mixin no longer rounds the tester's vocab up to a multiple of the world size: vocab_size=99 exercises the uneven lm_head in every TP test. The backward test slices the reference gradients (plain and packed) with ceil-sized shards to match.
• Unit tests cover the linear-only check and uneven local shapes for both embedding shardings.
• CPU TP/EP mixin tests, with the allowlist forced off: Llama, Qwen3, Gemma3, DeepSeek-V3, Qwen3-MoE pass. Inkling still fails: its embedding module applies a norm to the partial lookup before it's reduced, which main already hits with its explicit embedding_rowwise entry (Inkling isn't in the CI allowlist).

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

[IlyasMoutawwakil]

I added a cost section to discuss whether we should warn the user, bcz it might make sense to drop some tok ids from the vocab for faster inference/training.

[3outeille]

no need to edit every tp plan, we have it by default in configuration_utils.py

   if getattr(self, "tie_word_embeddings", False) and self.base_model_tp_plan is not None:
            self.base_model_tp_plan = {
                **self.base_model_tp_plan,
                "embed_tokens": "embedding_rowwise",
            }

[IlyasMoutawwakil]

nice !

[IlyasMoutawwakil]

can i drop the tie_word embedding condition ? it's not necessary for sharding the embedding right ? actually checkpoints that aren't tied benefit the most bcz the tied ones might already have the lm_head sharded.

[IlyasMoutawwakil]

also maybe better self.base_model_tp_plan = {"embed_tokens": "embedding_rowwise", **self.base_model_tp_plan} is better to let the model's tp plan take precedence (no model overrides it for now tho)

[IlyasMoutawwakil]

updated !

[IlyasMoutawwakil]

okay there's a problem in this, setting self.base_model_tp_plan changes it from a class attr to instance attr, changing the nature of the attr and breaking some tests where we patch the class attr.

[IlyasMoutawwakil]

adding it in the config to models that don't necessarily tie weights also broke some vlms that do have embed_tokens in the text decoder, but don't have it in the vision encoder so test_tp_plan_matches_params complained.

[IlyasMoutawwakil]

i moved it to the distributed mixin

            # Shard the input embeddings on the vocabulary unless the plan already names them.
            if self._tp_plan and isinstance(getattr(self, "embed_tokens", None), torch.nn.Embedding):
                self._tp_plan = {"embed_tokens": "embedding_rowwise", **self._tp_plan}

which you might not like 🥲 lmk what u think

[3outeille]

it would be best inside resolve_parallel_plans once I clean this PR #49527

[3outeille]

this is ready

[ArthurZucker]

BTW this IS breaking we need a small alert! And prob ping @hmellor

[ArthurZucker]

arf yeah would be last resort!
I expect us to move more default plan for defaults so finding a solution is welcome 👁️

[IlyasMoutawwakil]

the explicit solution is to name it in every plan 😆 see conv in #49570 (comment)

[ArthurZucker]

I Know !

[ArthurZucker]

❤️ exactly what we needed for al ong time ty!

[ArthurZucker]

ow what?

[IlyasMoutawwakil]

yep 🥲 q_proj is a Gemma4ClippableLinear wrapper; the actual nn.Linear and weight live at q_proj.linear

[IlyasMoutawwakil]

@vasqu

[vasqu]

was not deeply involved in gemma4 but wow indeedd.. super weird 😢

[ArthurZucker]

yeah I am surprised it even worked before

[3outeille]

we will be able to catch those errors/regression soon with the distributed CI dashboard

[IlyasMoutawwakil]

yeah I am surprised it even worked before

yeh normally we have a test that checks for plans matching params/modules, will check how it got past it. the encoder.layers.*.self_attn.q_proj is a module but now with the distinct linear/embedding colwise/rowwise plans we can check each plan is used with the class it's meant for for example.

[IlyasMoutawwakil]

@ArthurZucker I added a warning with what's breaking + the 🚨

[hmellor]

Thanks for the ping, I don't think this will cause me problems in vLLM.

I surgically replace all nn.Embeddings with vLLM's VocabParallelEmbedding/ParallelLMHead which have their TP behaviour built in, so I don't need the TP plan for these layers.

[ArthurZucker]

small nit as the func is not really readable / understandable but ok

[ArthurZucker]

yeah I am surprised it even worked before

[github-actions]

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

run-slow: doge, gemma4, inkling, muse_glimmer, nemotron_h, qwen4_exp

[github-actions]

---

CI recap

Dashboard: View test results in Grafana
Latest run: 39218537734:1
Result: success | Jobs: 15 | Tests: 206,374 | Failures: 0 | Duration: 13h 6m

[BenjaminBossan]

Just FYI, this broke some TP tests in PEFT. The fix seems to be simple enough though: huggingface/peft#4050.

If one of the TP experts (@3outeille or @michaelbenayoun) could confirm that this is valid, it would be great. Not sure if we need updates for LoRA embedding layers too.