🚨 TP: give embeddings their own style, shard input embeddings by default, and shard widths that don't divide by the TP size

docs/source/en/perf_infer_gpu_multi.md

@@ -106,7 +106,8 @@ The `ParallelInterface` class maps each strategy name you can use in a `tp_plan`
 ```py
 class ParallelInterface(GeneralInterface):
     _global_mapping = {
-        "embedding_rowwise": RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()),
+        "embedding_rowwise": EmbeddingParallel(shard_dim=0),
+        "embedding_colwise": EmbeddingParallel(shard_dim=1),
         "colwise_gather_output": ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()),
         "colwise_rep": ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()),
         "colwise": ColwiseParallel(input_layouts=Replicate(), output_layouts=Shard(-1)),
@@ -134,8 +135,9 @@ The table below describes each strategy.
 
 | Strategy | Description |
 |---|---|
-| `ColwiseParallel` | Shards a weight on its output-feature dim (`Shard(0)` for a 2D `nn.Linear` weight, `Shard(1)` for `nn.Embedding`) and shards a 1D bias. Redistributes the input to `Replicate()`, then places the output at `output_layouts`. |
-| `RowwiseParallel` | Shards a weight on its input-feature dim (`Shard(-1)`, or `Shard(0)` for `nn.Embedding`) and replicates the bias, which is added after the reduction. Each rank's forward produces a `Partial()` output holding its share of the sum, and `transform_output_post_forward` redistributes that to `output_layouts`. Reducing to `Replicate()` issues an all-reduce, and reducing to `Shard(1)` issues a reduce-scatter. |
+| `ColwiseParallel` | Shards a linear weight on its output-feature dim (`Shard(0)` for a 2D `nn.Linear` weight) and shards a 1D bias. Redistributes the input to `Replicate()`, then places the output at `output_layouts`. |
+| `RowwiseParallel` | Shards a linear weight on its input-feature dim (`Shard(-1)`) and replicates the bias, which is added after the reduction. Each rank's forward produces a `Partial()` output holding its share of the sum, and `transform_output_post_forward` redistributes that to `output_layouts`. Reducing to `Replicate()` issues an all-reduce, and reducing to `Shard(1)` issues a reduce-scatter. |
+| `EmbeddingParallel` | Shards an `nn.Embedding` with the token ids replicated. `shard_dim=0` (`embedding_rowwise`, the usual choice) shards the vocabulary and all-reduces each rank's partial lookup into the full output. `shard_dim=1` (`embedding_colwise`) shards the hidden dim and all-gathers each rank's slice of the output. |
 | `PackedColwiseParallel` | A variant of `ColwiseParallel` for fused weights, for example `up_proj` and `gate_proj` packed into `gate_up_proj`. Use `split_factor` when a weight packs a number of layers other than two. |
 | `PackedRowwiseParallel` | The row-wise counterpart, for weights packed along the final dim. Replicates 1D parameters. |
 | `SequenceParallel` | Replicates the module's parameters and shards its input on `sequence_dim` (defaults to `1`). Used for norms that operate per-token, such as `LayerNorm` and `RMSNorm`. |

src/transformers/configuration_utils.py

@@ -507,13 +507,6 @@ 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,
-                "embed_tokens": "embedding_rowwise",
-            }
-
         # Remap layer types if needed
         remap_legacy_layer_types(config=self)
 

src/transformers/distributed/mixin.py

@@ -19,6 +19,8 @@
 import warnings
 from typing import TYPE_CHECKING
 
+import torch
+
 from ..dynamic_module_utils import custom_object_save
 from ..utils import is_torch_greater_or_equal, logging
 from ..utils.hub import create_and_tag_model_card
@@ -90,6 +92,9 @@ def init_parallel_plans(self) -> None:
             self._tp_plan.update(self.config.base_model_tp_plan or {})
             self._ep_plan.update(self.config.base_model_ep_plan or {})
             self._fsdp_plan.update(self.config.base_model_fsdp_plan or {})
+            # 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}
 
         for name, module in self.named_children():
             if plan := getattr(module, "_ep_plan", None):

src/transformers/distributed/tensor_parallel.py

@@ -170,6 +170,15 @@ def context_around_forward(self, module, mesh):
     def transform_output_post_forward(self, module, output, mesh):
         return output
 
+    @staticmethod
+    def _from_local_last_dim_shard(local, mesh, size):
+        """Create a `DTensor` for a tensor sharded across its last dimension.
+        The local shard can have a different size from the other ranks when `size` is not divisible by the TP size."""
+        shape = (*local.shape[:-1], size)
+        return DTensor.from_local(
+            local, mesh, [Shard(-1)], run_check=False, shape=shape, stride=torch.empty(shape, device="meta").stride()
+        )
+
     def install_forward(self, module, mesh):
         """Install pre / around / post transforms by replacing module.forward."""
         original_forward = module.forward
@@ -185,7 +194,9 @@ def tp_forward(*args, **kwargs):
 
 
 class ColwiseParallel(TensorParallelLayer):
-    """Column-wise: weight & bias → Shard(0) (Embedding: Shard(1)); input replicated, output Shard(-1)."""
+    """Column-wise for linear layers: weight & bias → Shard(0); input replicated, output Shard(-1).
+
+    The output width (`out_features`) need not divide by the TP size: the last rank then holds a shorter slice."""
 
     def __init__(self, *, input_layouts=None, output_layouts=None, use_local_output: bool = True):
         self.input_layouts = input_layouts or Replicate()
@@ -197,28 +208,16 @@ def should_use_local_tensors(self, module):
         uses_local_inference_kernel = isinstance(module, torch.nn.Linear) and not torch.is_grad_enabled()
         return use_local_quantized_path or uses_local_inference_kernel
 
-    def validate_param(self, module, param, mesh, parameter_name=None):
-        meta = module._parameters.get(param)
-        gathers_output = isinstance(self.output_layouts, Replicate)
-        if meta is None or not gathers_output:
-            return
-
-        shard_dim = 1 if isinstance(module, torch.nn.Embedding) else meta.ndim - 2
-        output_size = meta.shape[shard_dim]
-        tp_size = mesh.size()
-        if output_size % tp_size != 0:
-            parameter_name = parameter_name or param
-            layer_name = parameter_name.rsplit(".", 1)[0]
-            raise ValueError(
-                f"The output size of `{layer_name}` ({output_size}) must be divisible by the tensor parallel size "
-                f"({tp_size}) when gathering a colwise output."
-            )
+    def install_forward(self, module, mesh):
+        if not hasattr(module, "out_features"):
+            raise ValueError(f"{type(module).__name__} has no `out_features`: colwise shards linear layers.")
+        return super().install_forward(module, mesh)
 
     def shard_param(self, module, param, mesh):
         meta = module._parameters.get(param)
         if meta is None:
             return
-        placement = Shard(1) if isinstance(module, torch.nn.Embedding) else Shard(meta.ndim - 2)
+        placement = Shard(meta.ndim - 2)
         module._parameters[param] = torch.nn.Parameter(
             distribute_tensor(meta, mesh, [placement], src_data_rank=None),
             requires_grad=meta.requires_grad,
@@ -261,18 +260,19 @@ def transform_output_post_forward(self, module, output, mesh):
         if self.should_use_local_tensors(module) and self.use_local_output and output_is_local_shard:
             return output
         if not isinstance(output, DTensor):
-            output = DTensor.from_local(output, mesh, [Shard(-1)], run_check=False)
+            # `out_features`, not the weight's shape: a packed or quantized weight doesn't show the full width
+            output = self._from_local_last_dim_shard(output, mesh, module.out_features)
         if output.placements != (self.output_layouts,):
             output = output.redistribute(placements=[self.output_layouts])
         return output.to_local() if self.use_local_output else output
 
 
 class RowwiseParallel(TensorParallelLayer):
-    """Row-wise: weight → Shard(1), bias → Replicate (Embedding: weight → Shard(0)).
+    """Row-wise for linear layers: weight → Shard(1), bias → Replicate; input sharded on the last dim.
 
-    Linear input is sharded on the last dim; Embedding input is replicated. The module
-    forward produces a Partial output which the boundary redistribute reduces to
-    output_layouts (Replicate → allreduce, Shard(1) → reduce-scatter).
+    The module forward produces a Partial output which the boundary redistribute reduces to
+    output_layouts (Replicate → allreduce, Shard(1) → reduce-scatter). The input width (`in_features`) need not
+    divide by the TP size.
     """
 
     def __init__(self, *, input_layouts=None, output_layouts=None, use_local_output: bool = True):
@@ -289,27 +289,32 @@ def shard_param(self, module, param, mesh):
         meta = module._parameters.get(param)
         if meta is None:
             return
-        if isinstance(module, torch.nn.Embedding):
-            placement = Shard(0)
-        else:
-            # bias is replicated (added after the row-reduce); weight shards on input dim (-1)
-            placement = Replicate() if param == "bias" else Shard(-1)
+        # bias is replicated (added after the row-reduce); weight shards on input dim (-1)
+        placement = Replicate() if param == "bias" else Shard(-1)
         module._parameters[param] = torch.nn.Parameter(
             distribute_tensor(meta, mesh, [placement], src_data_rank=None),
             requires_grad=meta.requires_grad,
         )
 
+    def install_forward(self, module, mesh):
+        if not hasattr(module, "in_features"):
+            raise ValueError(f"{type(module).__name__} has no `in_features`: rowwise shards linear layers.")
+        return super().install_forward(module, mesh)
+
     def transform_inputs_pre_forward(self, module, args, kwargs, mesh):
-        # Embedding runtime sharding needs a replicated input; Linear needs Shard(-1).
-        desired = Replicate() if isinstance(module, torch.nn.Embedding) else Shard(-1)
+        desired = Shard(-1)
         x = args[0]
         # A local kernel can use a plain input when the previous layer already split it.
         #  avoid a redundant Tensor -> DTensor -> Tensor round trip.
         input_has_desired_layout = self.input_layouts == desired
         if self.should_use_local_tensors(module) and input_has_desired_layout and not isinstance(x, DTensor):
             return args, kwargs
         if not isinstance(x, DTensor):
-            x = DTensor.from_local(x, mesh, [self.input_layouts], run_check=False)
+            if isinstance(self.input_layouts, Shard):
+                # the input is this rank's slice of the module's declared `in_features`
+                x = self._from_local_last_dim_shard(x, mesh, module.in_features)
+            else:
+                x = DTensor.from_local(x, mesh, [self.input_layouts], run_check=False)
         if x.placements != (desired,):
             x = x.redistribute(placements=[desired])
         if self.should_use_local_tensors(module):
@@ -360,6 +365,50 @@ def transform_output_post_forward(self, module, output, mesh):
         return output
 
 
+class EmbeddingParallel(TensorParallelLayer):
+    """Embedding with the token ids replicated on every rank. `shard_dim=0` is vocab-parallel: each rank holds a slice
+    of the vocabulary, and the ranks' partial lookups are summed. `shard_dim=1` shards the hidden dim: each rank looks
+    up a slice of every vector, and the slices are gathered. The sharded dim need not divide by the TP size."""
+
+    def __init__(self, *, shard_dim: int = 0, use_local_output: bool = True):
+        self.shard_dim = shard_dim
+        self.use_local_output = use_local_output
+
+    def should_use_local_tensors(self, module):
+        return getattr(module, "_hf_quantized_needs_local_tp", False)
+
+    def shard_param(self, module, param, mesh):
+        meta = module._parameters.get(param)
+        if meta is None:
+            return
+        module._parameters[param] = torch.nn.Parameter(
+            distribute_tensor(meta, mesh, [Shard(self.shard_dim)], src_data_rank=None),
+            requires_grad=meta.requires_grad,
+        )
+
+    def transform_inputs_pre_forward(self, module, args, kwargs, mesh):
+        x = args[0]
+        if self.should_use_local_tensors(module) and not isinstance(x, DTensor):
+            return args, kwargs
+        if not isinstance(x, DTensor):
+            x = DTensor.from_local(x, mesh, [Replicate()], run_check=False)
+        elif x.placements != (Replicate(),):
+            x = x.redistribute(placements=[Replicate()])
+        if self.should_use_local_tensors(module):
+            x = x.to_local()
+        return (x,) + args[1:], kwargs
+
+    def transform_output_post_forward(self, module, output, mesh):
+        if not isinstance(output, DTensor):
+            if self.shard_dim == 0:
+                output = DTensor.from_local(output, mesh, [Partial()], run_check=False)
+            else:
+                output = self._from_local_last_dim_shard(output, mesh, module.embedding_dim)
+        if output.placements != (Replicate(),):
+            output = output.redistribute(placements=[Replicate()])
+        return output.to_local() if self.use_local_output else output
+
+
 class ReplicatedWithGradAllReduce(TensorParallelLayer):
     """Replicated parameter whose gradient is partial.
 
@@ -945,7 +994,8 @@ class ParallelInterface(GeneralInterface):
 
     _global_mapping = (
         {
-            "embedding_rowwise": RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()),
+            "embedding_rowwise": EmbeddingParallel(shard_dim=0),
+            "embedding_colwise": EmbeddingParallel(shard_dim=1),
             "colwise_gather_output": ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()),
             "colwise_rep": ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()),
             "colwise": ColwiseParallel(input_layouts=Replicate(), output_layouts=Shard(-1)),

src/transformers/models/doge/configuration_doge.py

@@ -61,8 +61,8 @@ class DogeConfig(PreTrainedConfig):
         "layers.*.mlp.up_proj": "colwise",
         "layers.*.mlp.down_proj": "rowwise",
         "layers.*.mlp.router_gate": "colwise_gather_output",
-        "layers.*.mlp.down_embed": "rowwise_split_input",
-        "layers.*.mlp.up_embed": "rowwise_split_input",
+        "layers.*.mlp.down_embed": "embedding_rowwise",
+        "layers.*.mlp.up_embed": "embedding_rowwise",
     }
     base_model_pp_plan = {
         "embed_tokens": (["input_ids"], ["inputs_embeds"]),

src/transformers/models/doge/modular_doge.py

@@ -91,8 +91,8 @@ class DogeConfig(PreTrainedConfig):
         "layers.*.mlp.up_proj": "colwise",
         "layers.*.mlp.down_proj": "rowwise",
         "layers.*.mlp.router_gate": "colwise_gather_output",
-        "layers.*.mlp.down_embed": "rowwise_split_input",
-        "layers.*.mlp.up_embed": "rowwise_split_input",
+        "layers.*.mlp.down_embed": "embedding_rowwise",
+        "layers.*.mlp.up_embed": "embedding_rowwise",
     }
     base_model_pp_plan = {
         "embed_tokens": (["input_ids"], ["inputs_embeds"]),

src/transformers/models/gemma4/configuration_gemma4.py

@@ -256,15 +256,15 @@ class Gemma4VisionConfig(PreTrainedConfig):
 
     model_type = "gemma4_vision"
     base_model_tp_plan = {
-        "encoder.layers.*.self_attn.q_proj": "colwise",
-        "encoder.layers.*.self_attn.k_proj": "colwise",
-        "encoder.layers.*.self_attn.v_proj": "colwise",
+        "encoder.layers.*.self_attn.q_proj.linear": "colwise",
+        "encoder.layers.*.self_attn.k_proj.linear": "colwise",
+        "encoder.layers.*.self_attn.v_proj.linear": "colwise",
         "encoder.layers.*.self_attn.q_norm": "replicated_with_grad_allreduce",
         "encoder.layers.*.self_attn.k_norm": "replicated_with_grad_allreduce",
-        "encoder.layers.*.self_attn.o_proj": "rowwise",
-        "encoder.layers.*.mlp.gate_proj": "colwise",
-        "encoder.layers.*.mlp.up_proj": "colwise",
-        "encoder.layers.*.mlp.down_proj": "rowwise",
+        "encoder.layers.*.self_attn.o_proj.linear": "rowwise",
+        "encoder.layers.*.mlp.gate_proj.linear": "colwise",
+        "encoder.layers.*.mlp.up_proj.linear": "colwise",
+        "encoder.layers.*.mlp.down_proj.linear": "rowwise",
     }
 
     default_theta = 100.0

src/transformers/models/inkling/configuration_inkling.py

@@ -76,7 +76,6 @@ class InklingTextConfig(PreTrainedConfig):
     model_type = "inkling_text"
     base_config_key = "text_config"
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.mlp.experts.gate_up_proj": "colwise",
         "layers.*.mlp.experts.down_proj": "rowwise",
         "layers.*.mlp.experts": "moe_tp_experts",

src/transformers/models/inkling/modular_inkling.py

@@ -108,7 +108,6 @@ class InklingTextConfig(PreTrainedConfig):
     model_type = "inkling_text"
     base_config_key = "text_config"
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.mlp.experts.gate_up_proj": "colwise",
         "layers.*.mlp.experts.down_proj": "rowwise",
         "layers.*.mlp.experts": "moe_tp_experts",

src/transformers/models/muse_glimmer/configuration_muse_glimmer.py

@@ -98,7 +98,6 @@ class MuseGlimmerTextConfig(PreTrainedConfig):
     model_type = "muse_glimmer_text"
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_proj": "colwise",
         "layers.*.self_attn.k_proj": "colwise",
         "layers.*.self_attn.v_proj": "colwise",

src/transformers/models/muse_glimmer/modular_muse_glimmer.py

@@ -585,7 +585,6 @@ class MuseGlimmerTextConfig(Gemma2Config, PreTrainedConfig):
 
     model_type = "muse_glimmer_text"
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_proj": "colwise",
         "layers.*.self_attn.k_proj": "colwise",
         "layers.*.self_attn.v_proj": "colwise",

src/transformers/models/nemotron_h/configuration_nemotron_h.py

@@ -83,6 +83,7 @@ class NemotronHConfig(PreTrainedConfig):
     model_type = "nemotron_h"
     attribute_map = {"layer_types": "layers_block_type", "num_local_experts": "n_routed_experts"}
     base_model_tp_plan = {
+        "embeddings": "embedding_rowwise",
         "layers.*.mixer.experts.up_proj": "colwise",
         "layers.*.mixer.experts.down_proj": "rowwise",
         "layers.*.mixer.experts": "moe_tp_experts",

src/transformers/models/qwen4_exp/configuration_qwen4_exp.py

@@ -96,7 +96,7 @@ class Qwen4ExpTextConfig(PreTrainedConfig):
         "hyper_connection_mixer.input_mix_weight_down": "rowwise_split_input",
         # It's extremely important to shard this embedding as its size is ~45B == 90 GiB - we shard on dim 1, as the checkpoints
         # for it are sharded on dim0 (this way Concatenate and tp do not work on the same dim and everything is easy)
-        "layers.*.ple.ple_embedding.ngram_embedding": "colwise_gather_output",
+        "layers.*.ple.ple_embedding.ngram_embedding": "embedding_colwise",
     }
     base_model_pp_plan = None
     base_model_ep_plan = {

src/transformers/models/qwen4_exp/modular_qwen4_exp.py

@@ -141,7 +141,7 @@ class Qwen4ExpTextConfig(Qwen3_5MoeTextConfig):
         "hyper_connection_mixer.input_mix_weight_down": "rowwise_split_input",
         # It's extremely important to shard this embedding as its size is ~45B == 90 GiB - we shard on dim 1, as the checkpoints
         # for it are sharded on dim0 (this way Concatenate and tp do not work on the same dim and everything is easy)
-        "layers.*.ple.ple_embedding.ngram_embedding": "colwise_gather_output",
+        "layers.*.ple.ple_embedding.ngram_embedding": "embedding_colwise",
     }
     base_model_fsdp_plan = {
         "embed_tokens": "free_full_weight",