Restore BC for the tensor-parallel API

docs/source/en/continuous_batching.md

@@ -419,17 +419,18 @@ to the `kernels` package, this is becoming rare).
 
 ## Tensor parallelism
 
-For models too large to fit on a single GPU, shard the weights across devices with tensor parallelism. Load the model with `tp_plan="auto"` and continuous batching reads the tensor parallel size from the model to size the paged KV cache per shard. See [Tensor parallelism](./tensor_parallelism) for the list of supported architectures and how sharding works.
+For models too large to fit on a single GPU, shard the weights across devices with tensor parallelism. Set the number of devices with `DistributedConfig(tp_size=N)`. Continuous batching reads the tensor parallel size from the model to size the paged KV cache per shard. See [Tensor parallelism](./tensor_parallelism) for the list of supported architectures and how sharding works.
 
 ```py
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, DistributedConfig
 from transformers.generation import ContinuousBatchingConfig, GenerationConfig
 
+distributed_config = DistributedConfig(tp_size=4)
 model = AutoModelForCausalLM.from_pretrained(
     "Qwen/Qwen3-32B",
     attn_implementation="paged|flash_attention_2",
-    tp_plan="auto",
+    distributed_config=distributed_config,
 )
 tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-32B")
 
@@ -448,7 +449,7 @@ torchrun --nproc-per-node 4 cb_tp.py
 The tensor parallel size must divide the model's `num_key_value_heads` (check the model config). The paged cache raises an error at startup otherwise, so choose an appropriate `--nproc-per-node`.
 
 > [!WARNING]
-> Don't set `device_map` with `tp_plan`. The two conflict because `device_map` places whole modules on specific GPUs, while `tp_plan` shards those same parameters across all GPUs.
+> Don't set `device_map` with `distributed_config`. The two conflict because `device_map` places whole modules on specific GPUs, while tensor parallelism shards those same parameters across all GPUs.
 
 ## Sliding window attention
 

docs/source/en/expert_parallelism.md

@@ -49,3 +49,5 @@ Launch your inference script with [torchrun](https://pytorch.org/docs/stable/ela
 ```zsh
 torchrun --nproc-per-node 8 your_script.py
 ```
+
+[[autodoc]] DistributedConfig

docs/source/en/experts_interface.md

@@ -156,12 +156,15 @@ This backend requires:
 - A `torch.distributed` process group for the expert-parallel group, which the tensor-parallel wrapping supplies automatically.
 
 ```py
-from transformers import AutoModelForCausalLM
+import os
+
+from transformers import AutoModelForCausalLM, DistributedConfig
 
+distributed_config = DistributedConfig(tp_size=int(os.environ["WORLD_SIZE"]))
 model = AutoModelForCausalLM.from_pretrained(
     "deepseek-ai/DeepSeek-V4",
     experts_implementation="deepgemm_megamoe",
-    tp_plan="auto",
+    distributed_config=distributed_config,
 )
 ```
 

docs/source/en/optimization_overview.md

@@ -123,13 +123,16 @@ outputs = model.generate(**inputs, do_sample=False, max_new_tokens=50, cache_imp
 
 [Parallelism](./perf_infer_gpu_multi) distributes a model across devices so models too big for one device run fast. This approach uses more memory due to sharding overhead and communication to sync results.
 
-[Tensor parallelism](./perf_infer_gpu_multi) splits a model layer across devices. Set `tp_plan="auto"` in [`~PreTrainedModel.from_pretrained`] to enable it.
+[Tensor parallelism](./perf_infer_gpu_multi) splits a model layer across devices. Set the number of devices with `DistributedConfig(tp_size=N)` and pass it to [`~PreTrainedModel.from_pretrained`] to enable it.
 
 ```py
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, DistributedConfig
 
-model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct", tp_plan="auto")
+distributed_config = DistributedConfig(tp_size=4)
+model = AutoModelForCausalLM.from_pretrained(
+    "meta-llama/Meta-Llama-3-8B-Instruct",
+    distributed_config=distributed_config,
+)
 print(model._tp_plan)
 ```
 

docs/source/en/perf_infer_gpu_multi.md

@@ -40,21 +40,27 @@ This guide covers enabling tensor parallelism in Transformers and the available
 
 ## Partitioning a model
 
-Transformers enables tensor parallelism when a model has a `tp_plan`. Choose from two partitioning methods.
+Transformers supports tensor parallelism for models with a predefined plan. Configure the number of tensor parallel devices with `tp_size` in [`DistributedConfig`].
 
-- Set `tp_plan="auto"` for an automatic plan based on the model's predefined configuration.
-- Define and pass a manual `tp_plan`.
+- Set `DistributedConfig(tp_size=N)` to use the model's predefined plan.
+- Define a manual `tp_plan` and pass it to [`DistributedConfig`] with `tp_size`.
+
+You can also set `tp_plan="auto"` in [`DistributedConfig`]. When `tp_size` is omitted, it is inferred from `WORLD_SIZE`. Passing `tp_plan` directly to [`~PreTrainedModel.from_pretrained`] is deprecated and will be removed in v5.18.
 
 <hfoptions id="tp_plan">
 <hfoption id="auto plan">
 
 ```py
-import os
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, DistributedConfig
 
 # model_id = "meta-llama/Llama-4-Scout-17B-16E-Instruct" # better to visualize all the possible strategies
-model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct" , dtype=torch.bfloat16, tp_plan="auto")
+distributed_config = DistributedConfig(tp_size=4)
+model = AutoModelForCausalLM.from_pretrained(
+    "meta-llama/Meta-Llama-3-8B-Instruct",
+    dtype=torch.bfloat16,
+    distributed_config=distributed_config,
+)
 print(model._tp_plan)
 
 tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
@@ -65,7 +71,7 @@ inputs = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
 outputs = model(inputs)
 ```
 
-Launch the inference script with [torchrun](https://pytorch.org/docs/stable/elastic/run.html). Use 4 processes per GPU.
+Launch the inference script with [torchrun](https://pytorch.org/docs/stable/elastic/run.html). Use one process per GPU.
 
 ```bash
 torchrun --nproc-per-node 4 demo.py
@@ -74,12 +80,12 @@ torchrun --nproc-per-node 4 demo.py
 </hfoption>
 <hfoption id="manual plan">
 
-Define a tensor parallel plan for each layer in `tp_plan`. Pass it to [`~PreTrainedModel.from_pretrained`]. The example below uses column and row partitioning. See the [Partitioning strategies](#partitioning-strategies) section for other supported strategies.
+Define a tensor parallel plan for each layer in `tp_plan` and pass it through [`DistributedConfig`]. The example below uses column and row partitioning. See the [Partitioning strategies](#partitioning-strategies) section for other supported strategies.
 
 Manual partitioning requires a deep understanding of model architecture and strategy interactions. Poor partitioning choices create slow models that fail or produce incorrect results. The [Ultra-Scale Playbook](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=tensor_parallelism) explains partitioning strategies in detail.
 
 ```py
-from transformers import AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, DistributedConfig
 
 tp_plan = {
     "model.layers.*.self_attn.q_proj": "colwise",
@@ -89,7 +95,12 @@ tp_plan = {
     ...
 }
 
-model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct", dtype="auto", tp_plan=tp_plan)
+distributed_config = DistributedConfig(tp_size=4, tp_plan=tp_plan)
+model = AutoModelForCausalLM.from_pretrained(
+    "meta-llama/Meta-Llama-3-8B-Instruct",
+    dtype="auto",
+    distributed_config=distributed_config,
+)
 print(model.tp_plan)
 ```
 
@@ -98,7 +109,7 @@ print(model.tp_plan)
 
 ## Partitioning strategies
 
-The [`ParallelInterface`] class defines all partitioning strategies. It maps a string to the strategy implementation. You don't need to interact with this class directly since you set strategies with `tp_plan` in [`~PreTrainedModel.from_pretrained`]. It's useful for checking available strategies.
+The [`ParallelInterface`] class defines all partitioning strategies. It maps a string to the strategy implementation. You don't need to interact with this class directly since you set strategies with `tp_plan` in [`DistributedConfig`]. It's useful for checking available strategies.
 
 ```py
 class ParallelInterface(MutableMapping):
@@ -232,14 +243,22 @@ The example below shows how to implement `ColwiseParallel` with this workflow.
 3. Register the strategy to [`ParallelInterface`] to enable it for use with `tp_plan`.
 
     ```python
+    import torch
+
+    from transformers import AutoModelForCausalLM, DistributedConfig
     from transformers.integrations.tensor_parallel import ParallelInterface
 
     ParallelInterface.register_strategy("colwise_custom", ColwiseParallel)
     tp_plan = {
         "model.layers.*.self_attn.q_proj": "colwise_custom",
         ...
     }
-    model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.bfloat16, tp_plan=tp_plan)
+    distributed_config = DistributedConfig(tp_size=4, tp_plan=tp_plan)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        dtype=torch.bfloat16,
+        distributed_config=distributed_config,
+    )
     ```
 
 ## Benchmarks

docs/source/en/tensor_parallelism.md

@@ -46,24 +46,28 @@ print(config.base_model_tp_plan is not None)
 print(config.base_model_tp_plan)
 ```
 
-If a model supports TP, set `tp_plan="auto"` in [`~PreTrainedModel.from_pretrained`]. Transformers initializes the device mesh and shards the supported layers for you.
+If a model supports TP, create a [`DistributedConfig`] with the number of devices in `tp_size` and pass it to [`~PreTrainedModel.from_pretrained`]. Transformers uses the model's predefined plan, initializes the device mesh, and shards the supported layers for you.
+
+You can also set `tp_plan="auto"` in [`DistributedConfig`]. When `tp_size` is omitted, it is inferred from `WORLD_SIZE`. Passing `tp_plan` directly to [`~PreTrainedModel.from_pretrained`] is deprecated and will be removed in v5.18.
 
 > [!WARNING]
-> Don't use `device_map` with `tp_plan`. The two conflict at the weight-loading level. `device_map` places whole modules on specific GPUs, while `tp_plan` shards those same parameters across all GPUs.
+> Don't use `device_map` with `distributed_config`. The two conflict at the weight-loading level. `device_map` places whole modules on specific GPUs, while tensor parallelism shards those same parameters across all GPUs.
 
 ```py
 import torch
 
-from transformers import AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, DistributedConfig
+
+distributed_config = DistributedConfig(tp_size=4)
 
 model = AutoModelForCausalLM.from_pretrained(
     "Qwen/Qwen3-0.6B",
     dtype=torch.bfloat16,
-    tp_plan="auto",
+    distributed_config=distributed_config,
 )
 ```
 
-[`Trainer`] detects `tp_plan`, reads `tp_size` from the model, and creates a [`~accelerate.parallelism_config.ParallelismConfig`] automatically.
+[`Trainer`] detects the tensor parallel plan, reads `tp_size` from the model, and creates a [`~accelerate.parallelism_config.ParallelismConfig`] automatically.
 
 Launch training on one node with 4 GPUs.
 
@@ -79,12 +83,14 @@ Pass [`~accelerate.parallelism_config.ParallelismConfig`] explicitly when combin
 import torch
 
 from accelerate import ParallelismConfig
-from transformers import AutoModelForCausalLM, TrainingArguments
+from transformers import AutoModelForCausalLM, DistributedConfig, TrainingArguments
+
+distributed_config = DistributedConfig(tp_size=4)
 
 model = AutoModelForCausalLM.from_pretrained(
     "Qwen/Qwen3-0.6B",
     dtype=torch.bfloat16,
-    tp_plan="auto",
+    distributed_config=distributed_config,
 )
 
 parallelism_config = ParallelismConfig(tp_size=4)

docs/source/en/weightconverter.md

@@ -52,7 +52,7 @@ Checkpoint File → from_pretrained() → convert_and_load_state_dict_in_model()
 |------|-------------------|
 | Dynamic loading | Always, for all models |
 | Conversion mapping | Only when the model's class or `model_type` is registered in `_MODEL_TO_CONVERSION_PATTERN` |
-| TP sharding | Only when `tp_plan="auto"` and model has `base_model_tp_plan` |
+| TP sharding | Only when `DistributedConfig(tp_size=N)` is used and the model has `base_model_tp_plan` |
 | Dequantization/deserialization | Only when loading a pre-quantized checkpoint |
 | Quantization | Only when a quantization config is provided and weights are not pre-quantized |
 

src/transformers/__init__.py

@@ -104,7 +104,7 @@
     "debug_utils": [],
     "dependency_versions_check": [],
     "dependency_versions_table": [],
-    "distributed": [],
+    "distributed": ["DistributedConfig"],
     "dynamic_module_utils": [],
     "exporters": [],
     "feature_extraction_sequence_utils": ["SequenceFeatureExtractor"],
@@ -556,6 +556,7 @@
     from .data.datasets import GlueDataTrainingArguments as GlueDataTrainingArguments
     from .data.datasets import SquadDataset as SquadDataset
     from .data.datasets import SquadDataTrainingArguments as SquadDataTrainingArguments
+    from .distributed import DistributedConfig as DistributedConfig
     from .feature_extraction_sequence_utils import SequenceFeatureExtractor as SequenceFeatureExtractor
 
     # Feature Extractor

src/transformers/configuration_utils.py

@@ -348,6 +348,7 @@ def __post_init__(self, **kwargs):
         if per_layer_config is not None:
             self.per_layer_config = per_layer_config
 
+        # TODO: to support models whose input embedding module is not named `embed_tokens` (e.g. GPT-NeoX's `embed_in`).
         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,

src/transformers/distributed/configuration_utils.py

@@ -15,19 +15,21 @@
 import json
 import os
 from dataclasses import asdict, dataclass
+from typing import Literal
 
 
 @dataclass
 class DistributedConfig:
     """
-    Configuration for native distributed training (FSDP2 + TP).
+    Configuration for native distributed inference and training with tensor, pipeline, or FSDP2 parallelism.
 
     Args:
         tp_size (`int`, *optional*):
-            Number of devices for tensor parallelism. If `None` and `fsdp_size` is set, defaults to 1.
-        tp_plan (`dict`, *optional*):
-            Tensor parallel sharding plan. Leave as `None` to use the model's `base_model_tp_plan`.
-            Set explicitly to override.
+            Number of devices for tensor parallelism. If `None` and `tp_plan` is set, defaults to
+            `WORLD_SIZE // (other_parallel_size)`. If `None` and no `tp_plan` is set, defaults to 1.
+        tp_plan (`dict[str, str]` or `"auto"`, *optional*):
+            Tensor parallel sharding plan. Pass `"auto"`, or leave as `None` when `tp_size` is set, to use the
+            model's predefined `base_model_tp_plan`. Pass a dictionary to override the predefined plan.
         enable_sequence_parallel (`bool`, *optional*, defaults to `False`):
             Reserved for sequence parallelism. Not wired up yet.
         enable_expert_parallel (`bool`, *optional*, defaults to `False`):
@@ -38,10 +40,12 @@ class DistributedConfig:
             Whether to enable CPU offloading for FSDP2.
         fsdp_mixed_precision (`bool`, *optional*, defaults to `False`):
             Whether to enable mixed precision for FSDP2.
+        pp_size (`int`, *optional*):
+            Number of devices for pipeline parallelism. If `None` and another parallel mode is set, defaults to 1.
     """
 
     tp_size: int | None = None
-    tp_plan: dict[str, str] | None = None
+    tp_plan: dict[str, str] | Literal["auto"] | None = None
     enable_sequence_parallel: bool = False
     enable_expert_parallel: bool = False
     fsdp_size: int | None = None
@@ -50,15 +54,24 @@ class DistributedConfig:
     pp_size: int | None = None
 
     def __post_init__(self):
-        if self.tp_size is None and self.fsdp_size is None and self.pp_size is None:
+        if self.tp_plan is None and self.tp_size is None and self.fsdp_size is None and self.pp_size is None:
             return
 
-        if self.tp_size is None:
-            self.tp_size = 1
         if self.fsdp_size is None:
             self.fsdp_size = 1
         if self.pp_size is None:
             self.pp_size = 1
+        if self.tp_size is None and self.tp_plan is not None:
+            world_size = int(os.environ.get("WORLD_SIZE", 1))
+            other_parallel_size = self.fsdp_size * self.pp_size
+            if world_size % other_parallel_size != 0:
+                raise ValueError(
+                    f"WORLD_SIZE ({world_size}) must be divisible by fsdp_size * pp_size "
+                    f"({other_parallel_size}) to derive tp_size."
+                )
+            self.tp_size = world_size // other_parallel_size
+        elif self.tp_size is None:
+            self.tp_size = 1
 
         if self.tp_size > 1 and self.fsdp_size > 1 and self.pp_size > 1:
             raise ValueError(

src/transformers/distributed/mixin.py

@@ -24,7 +24,7 @@
 from .fsdp import apply_fully_sharded_data_parallelism, is_fsdp_managed_module
 from .pipeline_parallel import apply_pipeline_parallelism
 from .tensor_parallel import (
-    ALL_PARALLEL_STYLES,
+    _validate_tp_plan_styles,
     apply_tensor_parallelism,
     gather_state_dict_for_save,
 )
@@ -112,12 +112,7 @@ def tp_plan(self, plan: dict[str, str] | None):
         if not isinstance(plan, dict):
             raise ValueError("Can only set a dictionary as `tp_plan`")
 
-        for layer_pattern, parallel_style in plan.items():
-            if parallel_style not in ALL_PARALLEL_STYLES:
-                raise ValueError(
-                    f"Unsupported tensor parallel style '{parallel_style}' for layer '{layer_pattern}'. "
-                    f"Supported styles are {list(ALL_PARALLEL_STYLES.keys())}"
-                )
+        _validate_tp_plan_styles(plan)
 
         model_param_names = [name for name, _ in self.named_parameters()]
         for layer_pattern in plan.keys():
@@ -198,6 +193,8 @@ def maybe_distribute_model(
 
             if distributed_config.tp_size > 1:
                 tp_mesh = device_mesh["tp"] if device_mesh.ndim > 1 else device_mesh
+                if isinstance(distributed_config.tp_plan, dict):
+                    model.tp_plan = distributed_config.tp_plan
                 model = apply_tensor_parallelism(model, tp_mesh)
 
             elif distributed_config.fsdp_size > 1: