Restore BC for the tensor-parallel API by ArthurZucker · Pull Request #48300 · huggingface/transformers

ArthurZucker · 2026-08-25T15:34:44Z

Restore tp_plan from from_pretrained() and add deprecation cycle to it
More guarding regarding qk_rope_head_dim attributes in MLA + expert dim floor-divided by mesh size
Doc updates

`from_pretrained` accepts `tp_plan` and `tp_size` again, a supplied plan is applied to the model instead of being ignored, embeddings are shardable when the weights are untied, `DistributedConfig` is exported publicly, the renamed parallel-style classes keep their old names, and an mps host falls back to cpu/gloo rather than refusing to run.

An unsupported style in a config-derived plan was skipped instead of raising, the expert dimension was floor-divided by the mesh size with no divisibility check, and MLA lost its actionable error for a config without `qk_rope_head_dim`.

HuggingFaceDocBuilderDev · 2026-08-25T15:53:37Z

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.

…into fix-tp-plan-bc

…nto fix-tp-plan-bc

ArthurZucker · 2026-08-26T07:35:53Z

+            # Shard input embeddings for both tied and untied models. Note that architectures use different names for input embeddings.
+            if self._tp_plan:
+                try:
+                    input_embeddings = self.get_input_embeddings()
+                except NotImplementedError:
+                    input_embeddings = None
+                if isinstance(input_embeddings, nn.Embedding):
+                    input_embedding_name = next(
+                        (name for name, module in self.named_modules() if module is input_embeddings), None
+                    )
+                    if input_embedding_name:
+                        self._tp_plan.setdefault(input_embedding_name, "embedding_rowwise")
+


this is ai slope

ArthurZucker · 2026-08-26T07:36:07Z

+
+        def _matches(pattern: str) -> bool:
+            regex_pattern = pattern.replace("*", r"\d+")
+            return any(re.match(regex_pattern, name) for name in model_param_names)
+
+        prefix = getattr(self, "base_model_prefix", "")
+        if prefix:
+            plan = {
+                f"{prefix}.{k}" if not _matches(k) and _matches(f"{prefix}.{k}") else k: v for k, v in plan.items()
+            }
+


this is also ai slope

ArthurZucker · 2026-08-26T07:36:35Z

+        if style not in ALL_PARALLEL_STYLES:
+            raise ValueError(
+                f"Unsupported tensor parallel style '{style}' for layer '{layer_pattern}'. "
+                f"Supported styles are {list(ALL_PARALLEL_STYLES.keys())}"
+            )


no, collect all errors, then value error with all unrecognized

ArthurZucker · 2026-08-26T07:36:46Z

+            logger.warning_once(
+                "PyTorch's built-in DeviceMesh/DTensor stack does not support an MPS mesh. Falling back to CPU."
+            )


ArthurZucker · 2026-08-26T07:37:15Z

-        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",


This was good, if not present and tie, keep it

ArthurZucker · 2026-08-26T07:37:27Z

+            warnings.warn(
+                "Passing `tp_plan` directly to `from_pretrained` is deprecated and will be removed in v5.18. "
+                "Pass it in `distributed_config=DistributedConfig(tp_plan=...)` instead.",
+                FutureWarning,
+                stacklevel=2,


…into fix-tp-plan-bc

github-actions · 2026-08-26T10:03:43Z

CI recap

Dashboard: View test results in Grafana
Latest run: 32953452821:2
Result: failure | Jobs: 1 | Tests: 35 | Failures: 0 | Duration: 7m 13s

ArthurZucker · 2026-08-26T10:43:26Z

 model = AutoModelForImageTextToText.from_pretrained(
    "baidu/ERNIE-4.5-VL-28B-A3B-PT",
-    device_map="auto",  # Use tp_plan="auto" instead to enable Tensor Parallelism!
+    device_map="auto",  # Pass distributed_config=DistributedConfig(tp_size=N) instead to enable TP!


this is not ideal. It's annoying to write.
Can we either:

keep tp plan auto

auto detect torch run no device map -> tp plan auto?

yeah but your suggestions dont take into account other parallelism. I agree that having tp_plan=auto works wonder but if you want to combine with fsdp or pp, then DistributedConfig(tp_size=2, fsdp_size=2, pp_size=2) makes more sense

what I am saying and what you are saying are not against one another! its more about good defaults! DP is never used in inference, PP can be yes -> we would ideally want the torch run command to be caught and auto infer tp size and pp size (inter / intra)

okay let's say we go for DistributedConfig(tp_plan=auto, pp_plan=auto) and WORLD_SIZE=16,

we would intuitively infer tp_size=8 and pp_size=2 but if some reasons, someone would want to do tp_size=4 and pp_size=4, we loose the ability to do so since we auto infer.

On top of that, it's heavily dependant on the infra but sometimes PP can be faster than TP. Here some benchmark Ive ran back in the days (Same DP, batch settings, nodes, and GPU count. PP is 10.8% faster)

PP: dp16_tp1_pp4_mbs1_ga64 — 5211 tok/s/GPU (source: https://huggingface.co/datasets/nanotron/picotron_bench/tree/main/result_blog_posts/saturate_dp_7b_8_node_pp/dp16_tp1_pp4_mbs1_ga64_sl4096)
TP: dp16_tp4_pp1_mbs1_ga64 — 4705 tok/s/GPU (source: https://huggingface.co/datasets/nanotron/picotron_bench/tree/main/result_blog_posts/saturate_dp_7b_8_node_tp/dp16_tp4_pp1_mbs1_ga64_sl4096)

ArthurZucker · 2026-08-26T10:43:58Z

+distributed_config = DistributedConfig(tp_size=int(os.environ["WORLD_SIZE"]))
 model = Llama4ForConditionalGeneration.from_pretrained(
    model_id,
-    tp_plan="auto",
-    device_map="auto",
+    distributed_config=distributed_config,


not a good default at all. why force the use to compute world size, import os etc when a single simple tp plan was enough before

you can use DistributedConfig(tp_plan=auto) if you want but with other parallelism, that would be weird to doDistributedConfig(tp_plan=auto, fsdp_size=2, pp_size=2) for example

absolutely, but we need to remove friction for the more common use case : inference!

…into fix-tp-plan-bc

* Restore backward compatibility for the tensor-parallel API `from_pretrained` accepts `tp_plan` and `tp_size` again, a supplied plan is applied to the model instead of being ignored, embeddings are shardable when the weights are untied, `DistributedConfig` is exported publicly, the renamed parallel-style classes keep their old names, and an mps host falls back to cpu/gloo rather than refusing to run. * Restore the plan validation the refactor turned into silent skips An unsupported style in a config-derived plan was skipped instead of raising, the expert dimension was floor-divided by the mesh size with no divisibility check, and MLA lost its actionable error for a config without `qk_rope_head_dim`. * tp_plan deprecation warning * warning * better * bc modeling utils * bette rnaming * doc * doc * let mixin handle the tied untied of embedding * revert ai slop * revert * collect then raise * update doc * todo * revert * revert doc --------- Co-authored-by: Ferdinand Mom <47445085+3outeille@users.noreply.github.com> Co-authored-by: 3outeille <ferdinand.mom@epita.fr>

ArthurZucker added 2 commits August 26, 2026 00:34

3outeille reviewed Aug 25, 2026

View reviewed changes

Comment thread src/transformers/distributed/configuration_utils.py Outdated

3outeille and others added 13 commits August 26, 2026 01:19

Merge branch 'main' into fix-tp-plan-bc

c56345d

tp_plan deprecation warning

71b5f2a

Merge branch 'fix-tp-plan-bc' of github.com:huggingface/transformers …

f944126

…into fix-tp-plan-bc

warning

c84e49b

better

1f6aa5e

bc modeling utils

770876b

bette rnaming

e653b6a

Merge branch 'main' into fix-tp-plan-bc

6fa033c

doc

f417367

erge branch 'fix-tp-plan-bc' of github.com:huggingface/transformers i…

1f1778e

…nto fix-tp-plan-bc

doc

ee678dc

let mixin handle the tied untied of embedding

95aa691

Merge branch 'main' into fix-tp-plan-bc

63a97d5

ArthurZucker commented Aug 26, 2026

View reviewed changes

3outeille and others added 9 commits August 26, 2026 17:27

Merge branch 'main' into fix-tp-plan-bc

3cabc9c

revert ai slop

46b581a

Merge branch 'fix-tp-plan-bc' of github.com:huggingface/transformers …

861ad35

…into fix-tp-plan-bc

revert

7d92ec8

collect then raise

4efba1e

update doc

c59372f

Merge branch 'main' into fix-tp-plan-bc

7d16d57

todo

9aa263f

Merge branch 'fix-tp-plan-bc' of github.com:huggingface/transformers …

943d457

…into fix-tp-plan-bc

ArthurZucker commented Aug 26, 2026

View reviewed changes

revert

13deb9e

3outeille and others added 3 commits August 26, 2026 20:44

Merge branch 'main' into fix-tp-plan-bc

e535a92

revert doc

aa1eeb0

Merge branch 'fix-tp-plan-bc' of github.com:huggingface/transformers …

c20266e

…into fix-tp-plan-bc

ArthurZucker merged commit f0e42ef into main Aug 26, 2026
108 of 109 checks passed

ArthurZucker deleted the fix-tp-plan-bc branch August 26, 2026 11:54

vasqu added the for patch Tag issues / labels that should be included in the next patch label Aug 26, 2026

zhtmike mentioned this pull request Sep 7, 2026

[RFC][tracking] Upstream patch debt: audited ledger of temporary compatibility patches verl-project/verl-omni#445

Open

37 tasks

Conversation

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

Uh oh!

Uh oh!

Uh oh!

HuggingFaceDocBuilderDev commented Aug 25, 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!

github-actions Bot commented Aug 26, 2026

CI recap

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!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

ArthurZucker commented Aug 25, 2026 •

edited by github-actions Bot

Loading