🚨 Pass tp_plan from lm_heads by molbap · Pull Request #47253 · huggingface/transformers

molbap · 2026-07-10T14:43:00Z

What does this PR do?

Noticed currently the lm_head-specific tp plans are wiped in PreTrainedModel post-init (we just ignore them). This should fix it

HuggingFaceDocBuilderDev · 2026-07-10T14:56:23Z

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.

molbap · 2026-07-10T15:49:44Z

+        config = self.model_tester.get_config()
+        text_config = config.get_text_config()
+        remainder = text_config.vocab_size % self.tensor_parallel_size
+        if remainder:
+            text_config.vocab_size += self.tensor_parallel_size - remainder
+        return config


this is because we have typically vocab size 99 on tiny tests, which can't be sharded on 2 ranks obviously

I'm thinking wheter we should automatically extend an embedding if we notice this? So in from pretrained, if we notice

TP

Embeddings that get shared we resize the embeddings before

ah auto-pad you mean? hmm we could ig, I feel like we are already doing it in a couple other places IIRC

Yep, because users will complain I'm sure 😅 Maybe a warning but having users to think leads to problems

but we should pad just before the gather and unpad right after, no?

other idea (sorry for the noise): we should just validate, feels less magic. like if TP is provided with a wrong vocab size just raise from the get-go. Should be in another PR though because it could break existing setups

I think this is already breaking in itself no? Users could have arbitrary vocab sizes before this

But yea, I can see the validation path - less magic and add how to properly do it

ah yes, you're right... before it was silently replicating so of course this will raise. I'll add the validation logic then

IMO, if it does not cost much and we can resize for the user, would be nice. People probably expect us to do so, and we are the ones setting the plan to default to colwise rep.
The main concern is for the gather to not take into account the padding, which is important

vasqu

Imo, this looks good but let's wait for @ArthurZucker and @3outeille for a final look 🫡

github-actions · 2026-07-10T16:50:08Z

CI recap

Dashboard: View test results in Grafana
Latest run: 29107923987:1
Result: failure | Jobs: 15 | Tests: 170,565 | Failures: 1 | Duration: 15h 14m

ArthurZucker

ty, not sure we need generic, and as pointed let's maybe pad by default?
we can isolate the 2 changes. let's also trace the blame on PR that regressed the init, it was I am certain working before

ArthurZucker · 2026-07-13T08:15:12Z

+    def validate_module(self, module: nn.Module, device_mesh, layer_name: str = ""):
+        out_features = getattr(module, "out_features", None)
+        if self.gather_output and out_features is not None and out_features % device_mesh.size() != 0:
+            raise ValueError(
+                f"`{layer_name}` ({type(module).__name__} with out_features={out_features}) is sharded with "
+                f"'colwise_gather_output', which requires out_features to be divisible by the number of ranks "
+                f"({device_mesh.size()}) to all-gather equal-size shards. Resize the weight (e.g. "
+                f"`model.resize_token_embeddings` for LM heads) or override this module's entry in the tp_plan."
+            )
+


since only this layer uses it, we can not make it generic for now

yeah it's a bit awkward. Wdy suggest? move it the the layer?

ArthurZucker · 2026-07-13T08:15:54Z

        # easily available.
-        self._tp_plan, self._ep_plan, self._pp_plan, self._fsdp_plan = {}, {}, {}, {}
+        # Start from the class-level plans (e.g. `{"lm_head": "colwise_rep"}` on `...ForCausalLM` classes), copying
+        # them as they are mutated below and would otherwise contaminate the class attribute shared by all instances


absolutely, this was already the case.... do you know where the regression comes from?

yeah wrote that mostly for review. Yes, it seems to be coming from a big Revert, likely an oversight #46246

it's actually in #36677 and then we keep the pattern

Are you sure @3outeille ? it doesn't seem to be there that we set
self._tp_plan, self._ep_plan, self._pp_plan, self._fsdp_plan = {}, {}, {}, {}

ArthurZucker · 2026-07-13T08:16:10Z

    # models, this attribute is currently defined in respective model code. For base models, it comes from
    # `config.base_model_pp_plan` during `post_init`.
    _pp_plan: dict[str, tuple[str, str]] = None
+    # An expert parallel plan used instead of `_tp_plan` when expert parallelism is enabled. For base models, it comes


can you elaborate? 😁

ArthurZucker · 2026-07-13T08:17:57Z

+        config = self.model_tester.get_config()
+        text_config = config.get_text_config()
+        remainder = text_config.vocab_size % self.tensor_parallel_size
+        if remainder:
+            text_config.vocab_size += self.tensor_parallel_size - remainder
+        return config


IMO, if it does not cost much and we can resize for the user, would be nice. People probably expect us to do so, and we are the ones setting the plan to default to colwise rep.
The main concern is for the gather to not take into account the padding, which is important

ArthurZucker · 2026-07-15T18:16:45Z

Merging we need it for release

* pass tp_plan from lm_heads * add a small test to avoid regressiosn on TP * world size must divide vocab size * early raise for incompatible vocabs * test our test

molbap added 2 commits July 10, 2026 16:40

pass tp_plan from lm_heads

273c2e1

add a small test to avoid regressiosn on TP

08ab9ab

molbap requested review from 3outeille and ArthurZucker July 10, 2026 14:45

world size must divide vocab size

3c326e4

molbap commented Jul 10, 2026

View reviewed changes

molbap changed the title ~~Pass tp_plan from lm_heads~~ 🚨 Pass tp_plan from lm_heads Jul 10, 2026

molbap added 2 commits July 10, 2026 18:28

early raise for incompatible vocabs

ae6881e

test our test

d51432a

vasqu approved these changes Jul 10, 2026

View reviewed changes

3outeille approved these changes Jul 13, 2026

View reviewed changes

3outeille mentioned this pull request Jul 13, 2026

FSDP orchestration: apply + loading/saving #46990

Closed

ArthurZucker reviewed Jul 13, 2026

View reviewed changes

ArthurZucker merged commit f5a67a4 into main Jul 15, 2026
196 of 198 checks passed

ArthurZucker deleted the fix_tp_plans_for_lm_heads branch July 15, 2026 18:17

3outeille mentioned this pull request Jul 16, 2026

Wire DistributedConfig through from_pretrained and save_pretrained (FSDP orchestration 2/3) #47353

Closed

arielge mentioned this pull request Jul 22, 2026

TP enablement torch-spyre/hf-adapters#234

Merged

3 tasks

Conversation

molbap commented Jul 10, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Uh oh!

HuggingFaceDocBuilderDev commented Jul 10, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Jul 10, 2026

CI recap

Uh oh!

ArthurZucker left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ArthurZucker commented Jul 15, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

molbap commented Jul 10, 2026 •

edited by github-actions Bot

Loading