From 10a27d896b10311e0c1fd2c3c1190d0178ed09c8 Mon Sep 17 00:00:00 2001
From: IlyasMoutawwakil <moutawwakil.ilyas.tsi@gmail.com>
Date: Mon, 12 Oct 2026 10:55:21 +0200
Subject: [PATCH 01/13] TP: keep colwise/rowwise for linear layers, give
 embeddings their own styles

`ColwiseParallel` and `RowwiseParallel` no longer special-case `nn.Embedding`; they now require a linear layer.
Embeddings get `EmbeddingRowwiseParallel` (`embedding_rowwise`, vocab-parallel) and `EmbeddingColwiseParallel`
(`embedding_colwise`, hidden-dim parallel, used by qwen4_exp whose ngram table loads as vocab chunks).
Doge's MoE embeddings move to `embedding_rowwise`, and gemma4's vision plan targets the inner `.linear`
of `Gemma4ClippableLinear`, which holds the weights.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---
 docs/source/en/perf_infer_gpu_multi.md        |  9 +-
 .../distributed/tensor_parallel.py            | 94 ++++++++++++++++---
 .../models/doge/configuration_doge.py         |  4 +-
 src/transformers/models/doge/modular_doge.py  |  4 +-
 .../models/gemma4/configuration_gemma4.py     | 14 +--
 .../qwen4_exp/configuration_qwen4_exp.py      |  2 +-
 .../models/qwen4_exp/modular_qwen4_exp.py     |  2 +-
 tests/tensor_parallel/test_tensor_parallel.py | 23 ++---
 8 files changed, 110 insertions(+), 42 deletions(-)

diff --git a/docs/source/en/perf_infer_gpu_multi.md b/docs/source/en/perf_infer_gpu_multi.md
index c7fa4e4e1e64..29401f7fdbf5 100644
--- a/docs/source/en/perf_infer_gpu_multi.md
+++ b/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": EmbeddingRowwiseParallel(),
+        "embedding_colwise": EmbeddingColwiseParallel(),
         "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,10 @@ 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. |
+| `EmbeddingRowwiseParallel` | Shards an `nn.Embedding` on its vocabulary (`Shard(0)`). The token ids stay replicated, and each rank's partial lookup is all-reduced into the full output. |
+| `EmbeddingColwiseParallel` | Shards an `nn.Embedding` on its hidden dim (`Shard(1)`). The token ids stay replicated, and each rank's slice of the output is all-gathered. |
 | `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`. |
diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py
index 00b44df39b63..626790a15ce8 100644
--- a/src/transformers/distributed/tensor_parallel.py
+++ b/src/transformers/distributed/tensor_parallel.py
@@ -170,7 +170,7 @@ 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)."""
 
     def __init__(self, *, input_layouts=None, output_layouts=None, use_local_output: bool = True):
         self.input_layouts = input_layouts or Replicate()
@@ -188,8 +188,7 @@ def validate_param(self, module, param, mesh, parameter_name=None):
         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]
+        output_size = meta.shape[meta.ndim - 2]
         tp_size = mesh.size()
         if output_size % tp_size != 0:
             parameter_name = parameter_name or param
@@ -199,11 +198,16 @@ def validate_param(self, module, param, mesh, parameter_name=None):
                 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,
@@ -253,10 +257,9 @@ def transform_output_post_forward(self, module, output, mesh):
 
 
 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
+    The module forward produces a Partial output which the boundary redistribute reduces to
     output_layouts (Replicate → allreduce, Shard(1) → reduce-scatter).
     """
 
@@ -274,19 +277,20 @@ 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.
@@ -345,6 +349,65 @@ def transform_output_post_forward(self, module, output, mesh):
         return output
 
 
+class EmbeddingRowwiseParallel(TensorParallelLayer):
+    """Vocab-parallel embedding: weight → Shard(0), so each rank holds a slice of the vocabulary; the token ids stay
+    replicated, and the ranks' partial lookups are summed into the full output."""
+
+    shard_dim = 0
+    local_output_placement = Partial()
+
+    def __init__(self, *, use_local_output: bool = True):
+        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):
+            output = DTensor.from_local(output, mesh, [self.local_output_placement], run_check=False)
+        if output.placements != (Replicate(),):
+            output = output.redistribute(placements=[Replicate()])
+        return output.to_local() if self.use_local_output else output
+
+
+class EmbeddingColwiseParallel(EmbeddingRowwiseParallel):
+    """Embedding sharded on its hidden dim: weight → Shard(1), so each rank looks up a slice of every vector; the
+    token ids stay replicated, and the slices are gathered into the full output."""
+
+    shard_dim = 1
+    local_output_placement = Shard(-1)
+
+    def validate_param(self, module, param, mesh, parameter_name=None):
+        meta = module._parameters.get(param)
+        if meta is not None and meta.shape[1] % mesh.size() != 0:
+            layer_name = (parameter_name or param).rsplit(".", 1)[0]
+            raise ValueError(
+                f"The embedding size of `{layer_name}` ({meta.shape[1]}) must be divisible by the tensor parallel "
+                f"size ({mesh.size()})."
+            )
+
+
 class ReplicatedWithGradAllReduce(TensorParallelLayer):
     """Replicated parameter whose gradient is partial.
 
@@ -927,7 +990,8 @@ class ParallelInterface(GeneralInterface):
 
     _global_mapping = (
         {
-            "embedding_rowwise": RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()),
+            "embedding_rowwise": EmbeddingRowwiseParallel(),
+            "embedding_colwise": EmbeddingColwiseParallel(),
             "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)),
diff --git a/src/transformers/models/doge/configuration_doge.py b/src/transformers/models/doge/configuration_doge.py
index 8518d9021458..8ff3d2a4fdb1 100644
--- a/src/transformers/models/doge/configuration_doge.py
+++ b/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"]),
diff --git a/src/transformers/models/doge/modular_doge.py b/src/transformers/models/doge/modular_doge.py
index 897afec72f3f..eb9e83d83ac9 100644
--- a/src/transformers/models/doge/modular_doge.py
+++ b/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"]),
diff --git a/src/transformers/models/gemma4/configuration_gemma4.py b/src/transformers/models/gemma4/configuration_gemma4.py
index 0b78dfc1bf1d..6cc0e2a95fd3 100644
--- a/src/transformers/models/gemma4/configuration_gemma4.py
+++ b/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
diff --git a/src/transformers/models/qwen4_exp/configuration_qwen4_exp.py b/src/transformers/models/qwen4_exp/configuration_qwen4_exp.py
index 62085959be4f..9a7e85300baa 100644
--- a/src/transformers/models/qwen4_exp/configuration_qwen4_exp.py
+++ b/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 = {
diff --git a/src/transformers/models/qwen4_exp/modular_qwen4_exp.py b/src/transformers/models/qwen4_exp/modular_qwen4_exp.py
index 9a94f54073c5..e23ca547f566 100644
--- a/src/transformers/models/qwen4_exp/modular_qwen4_exp.py
+++ b/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",
diff --git a/tests/tensor_parallel/test_tensor_parallel.py b/tests/tensor_parallel/test_tensor_parallel.py
index c83741361fd5..82bbec10dc00 100644
--- a/tests/tensor_parallel/test_tensor_parallel.py
+++ b/tests/tensor_parallel/test_tensor_parallel.py
@@ -251,20 +251,21 @@ def test_rowwise_uneven_local_shapes(self):
             self.assertEqual(weight_shape, (32, expected_size))
             self.assertEqual(bias_shape, (10,))
 
-    def test_embedding_uneven_local_shapes(self):
-        rowwise_embedding = torch.nn.Embedding(10, 10)
-        rowwise_placement = self._get_parameter_placements(rowwise_embedding, RowwiseParallel())["weight"]
+    def test_linear_styles_reject_embedding(self):
+        for style in ("colwise", "rowwise"):
+            with self.subTest(style=style):
+                with self.assertRaises(ValueError) as context:
+                    ALL_PARALLEL_STYLES[style].install_forward(torch.nn.Embedding(8, 16), mesh=None)
+                self.assertIn("linear layers", str(context.exception))
 
-        colwise_embedding = torch.nn.Embedding(10, 10)
-        colwise_placement = self._get_parameter_placements(colwise_embedding, ColwiseParallel())["weight"]
+    def test_embedding_uneven_local_shapes(self):
+        embedding = torch.nn.Embedding(10, 10)
+        placement = self._get_parameter_placements(embedding, ALL_PARALLEL_STYLES["embedding_rowwise"])["weight"]
 
         expected_local_sizes = (4, 4, 2)
         for rank, expected_size in enumerate(expected_local_sizes):
-            rowwise_shape = self._get_local_shape((10, 10), rowwise_placement, world_size=3, rank=rank)
-            colwise_shape = self._get_local_shape((10, 10), colwise_placement, world_size=3, rank=rank)
-
-            self.assertEqual(rowwise_shape, (expected_size, 10))
-            self.assertEqual(colwise_shape, (10, expected_size))
+            local_shape = self._get_local_shape((10, 10), placement, world_size=3, rank=rank)
+            self.assertEqual(local_shape, (expected_size, 10))
 
     def test_shard_tensor_shape_consistency(self):
         world_size = 4
@@ -291,7 +292,7 @@ def test_shard_tensor_shape_consistency(self):
             },
             "embedding_colwise": {
                 "module": torch.nn.Embedding(32, 16),
-                "style": ColwiseParallel(),
+                "style": ALL_PARALLEL_STYLES["embedding_colwise"],
                 "expected_shapes": {"weight": (32, 4)},
             },
         }

From eeeb4794663aed4052389134538ddf615cdb1729 Mon Sep 17 00:00:00 2001
From: IlyasMoutawwakil <moutawwakil.ilyas.tsi@gmail.com>
Date: Mon, 12 Oct 2026 11:10:46 +0200
Subject: [PATCH 02/13] TP: split a linear layer unevenly when its width does
 not divide by the TP size

`DTensor.from_local` assumes even shards, so colwise outputs and rowwise inputs are wrapped with the full
shape taken from the module's `out_features` / `in_features`. The colwise divisibility check goes away, and
the TP tester keeps the tiny vocab of 99 so `lm_head` exercises the uneven split end to end.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---
 .../distributed/tensor_parallel.py            | 44 ++++++++++---------
 tests/tensor_parallel/test_tensor_parallel.py | 12 -----
 tests/test_tensor_parallel_mixin.py           | 13 +++---
 3 files changed, 29 insertions(+), 40 deletions(-)

diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py
index 626790a15ce8..b8b7cba59a8e 100644
--- a/src/transformers/distributed/tensor_parallel.py
+++ b/src/transformers/distributed/tensor_parallel.py
@@ -155,6 +155,18 @@ 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, placement, size):
+        """Wrap this rank's slice of a tensor split along its last dimension, whose full width is `size`.
+
+        When `size` doesn't divide by the number of ranks, the last rank holds a shorter slice. `DTensor.from_local`
+        can't tell that from one rank's slice and would assume every rank holds the same width, so we pass the full
+        shape explicitly."""
+        shape = (*local.shape[:-1], size)
+        return DTensor.from_local(
+            local, mesh, [placement], 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
@@ -170,7 +182,9 @@ def tp_forward(*args, **kwargs):
 
 
 class ColwiseParallel(TensorParallelLayer):
-    """Column-wise for linear layers: weight & bias → Shard(0); 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()
@@ -182,22 +196,6 @@ 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
-
-        output_size = meta.shape[meta.ndim - 2]
-        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.")
@@ -250,7 +248,8 @@ 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, Shard(-1), 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
@@ -260,7 +259,8 @@ class RowwiseParallel(TensorParallelLayer):
     """Row-wise for linear layers: weight → Shard(1), bias → Replicate; input sharded on the last dim.
 
     The module forward produces a Partial output which the boundary redistribute reduces to
-    output_layouts (Replicate → allreduce, Shard(1) → reduce-scatter).
+    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):
@@ -298,7 +298,11 @@ def transform_inputs_pre_forward(self, module, args, kwargs, mesh):
         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, self.input_layouts, 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):
diff --git a/tests/tensor_parallel/test_tensor_parallel.py b/tests/tensor_parallel/test_tensor_parallel.py
index 82bbec10dc00..9cee944a5587 100644
--- a/tests/tensor_parallel/test_tensor_parallel.py
+++ b/tests/tensor_parallel/test_tensor_parallel.py
@@ -211,18 +211,6 @@ def _make_dtensor_shard_op(self, mesh, placement, param_shape, local_shape):
         op._axis0_local_size = local_shape[0]
         return op
 
-    def test_colwise_gather_output_rejects_indivisible_out_features(self):
-        model = torch.nn.Module()
-        model.lm_head = torch.nn.Linear(8, 99)
-        model.tp_plan = {"lm_head": "colwise_gather_output"}
-        device_mesh = self.MockDeviceMesh(world_size=2, rank=0)
-
-        with self.assertRaises(ValueError) as context:
-            tensor_parallel.apply_tensor_parallelism(model, device_mesh)
-
-        self.assertIn("lm_head", str(context.exception))
-        self.assertIn("divisible", str(context.exception))
-
     def test_colwise_uneven_local_shapes(self):
         module = torch.nn.Module()
         module.register_parameter("weight", torch.nn.Parameter(torch.empty(10, 32)))
diff --git a/tests/test_tensor_parallel_mixin.py b/tests/test_tensor_parallel_mixin.py
index 2789f8798494..34975ff130ad 100644
--- a/tests/test_tensor_parallel_mixin.py
+++ b/tests/test_tensor_parallel_mixin.py
@@ -281,10 +281,9 @@ def _test_tp_backward_impl(rank, model_path, model_class, atol, rtol):
                             # interleaved slicing
                             grad = get_packed_grad_shard(grad, world_size, rank, dim)
                         else:
-                            # regular slicing
-                            shard_size = grad_tp.size(dim)
-                            start = rank * shard_size
-                            grad = grad.narrow(dim, start, shard_size)
+                            # regular slicing: every shard but the last is ceil(size / world_size) long
+                            start = rank * -(-grad.size(dim) // world_size)
+                            grad = grad.narrow(dim, start, grad_tp.size(dim))
                         break
 
             if not torch.allclose(grad.cpu(), grad_tp.cpu(), atol=atol, rtol=rtol):
@@ -524,12 +523,10 @@ def _get_tp_model_class(self):
         return self.all_model_classes[0]
 
     def _get_tp_config(self, tie_word_embeddings: bool | None = None):
-        """Tiny config with `vocab_size` rounded up to a multiple of the world size, as sharded dims (typically `lm_head`) have to be split across ranks."""
+        """The tester's tiny config. Its vocab stays as it is: 99 does not divide by the world size, which is
+        the uneven sharding of `lm_head` that real vocabs need."""
         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
         if tie_word_embeddings is not None:
             if hasattr(text_config, "tie_word_embeddings"):
                 text_config.tie_word_embeddings = tie_word_embeddings

From 3bbb894d836b5990d05a415a2689fed3a9e5a365 Mon Sep 17 00:00:00 2001
From: IlyasMoutawwakil <moutawwakil.ilyas.tsi@gmail.com>
Date: Tue, 13 Oct 2026 16:13:13 +0200
Subject: [PATCH 03/13] TP: one EmbeddingParallel style for both embedding
 shardings

`embedding_rowwise` and `embedding_colwise` become `EmbeddingParallel(shard_dim=0)` and
`EmbeddingParallel(shard_dim=1)` instead of two classes, and a hidden-sharded embedding no longer
needs a hidden size divisible by the TP size: its local output is wrapped with
`_from_local_last_dim_shard`, like the linear layers'. The helper always shards the last dim, so it
drops its placement argument.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---
 docs/source/en/perf_infer_gpu_multi.md        |  7 ++-
 .../distributed/tensor_parallel.py            | 54 +++++++------------
 tests/tensor_parallel/test_tensor_parallel.py | 10 ++--
 3 files changed, 27 insertions(+), 44 deletions(-)

diff --git a/docs/source/en/perf_infer_gpu_multi.md b/docs/source/en/perf_infer_gpu_multi.md
index 29401f7fdbf5..31972298904f 100644
--- a/docs/source/en/perf_infer_gpu_multi.md
+++ b/docs/source/en/perf_infer_gpu_multi.md
@@ -106,8 +106,8 @@ The `ParallelInterface` class maps each strategy name you can use in a `tp_plan`
 ```py
 class ParallelInterface(GeneralInterface):
     _global_mapping = {
-        "embedding_rowwise": EmbeddingRowwiseParallel(),
-        "embedding_colwise": EmbeddingColwiseParallel(),
+        "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)),
@@ -137,8 +137,7 @@ The table below describes each strategy.
 |---|---|
 | `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. |
-| `EmbeddingRowwiseParallel` | Shards an `nn.Embedding` on its vocabulary (`Shard(0)`). The token ids stay replicated, and each rank's partial lookup is all-reduced into the full output. |
-| `EmbeddingColwiseParallel` | Shards an `nn.Embedding` on its hidden dim (`Shard(1)`). The token ids stay replicated, and each rank's slice of the output is all-gathered. |
+| `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`. |
diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py
index b8b7cba59a8e..d07dd912a112 100644
--- a/src/transformers/distributed/tensor_parallel.py
+++ b/src/transformers/distributed/tensor_parallel.py
@@ -156,15 +156,12 @@ def transform_output_post_forward(self, module, output, mesh):
         return output
 
     @staticmethod
-    def _from_local_last_dim_shard(local, mesh, placement, size):
-        """Wrap this rank's slice of a tensor split along its last dimension, whose full width is `size`.
-
-        When `size` doesn't divide by the number of ranks, the last rank holds a shorter slice. `DTensor.from_local`
-        can't tell that from one rank's slice and would assume every rank holds the same width, so we pass the full
-        shape explicitly."""
+    def _from_local_last_dim_shard(local, mesh, size):
+        """Wrap this rank's slice of a tensor sharded on its last dim, whose full width is `size`. `DTensor.from_local`
+        would assume every rank holds an equal slice, which isn't so when `size` doesn't divide by the TP size."""
         shape = (*local.shape[:-1], size)
         return DTensor.from_local(
-            local, mesh, [placement], run_check=False, shape=shape, stride=torch.empty(shape, device="meta").stride()
+            local, mesh, [Shard(-1)], run_check=False, shape=shape, stride=torch.empty(shape, device="meta").stride()
         )
 
     def install_forward(self, module, mesh):
@@ -249,7 +246,7 @@ def transform_output_post_forward(self, module, output, mesh):
             return output
         if not isinstance(output, DTensor):
             # `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, Shard(-1), module.out_features)
+            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
@@ -300,7 +297,7 @@ def transform_inputs_pre_forward(self, module, args, kwargs, mesh):
         if not isinstance(x, DTensor):
             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, self.input_layouts, module.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,):
@@ -353,14 +350,13 @@ def transform_output_post_forward(self, module, output, mesh):
         return output
 
 
-class EmbeddingRowwiseParallel(TensorParallelLayer):
-    """Vocab-parallel embedding: weight → Shard(0), so each rank holds a slice of the vocabulary; the token ids stay
-    replicated, and the ranks' partial lookups are summed into the full output."""
-
-    shard_dim = 0
-    local_output_placement = Partial()
+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, *, use_local_output: bool = True):
+    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):
@@ -389,29 +385,15 @@ def transform_inputs_pre_forward(self, module, args, kwargs, mesh):
 
     def transform_output_post_forward(self, module, output, mesh):
         if not isinstance(output, DTensor):
-            output = DTensor.from_local(output, mesh, [self.local_output_placement], run_check=False)
+            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 EmbeddingColwiseParallel(EmbeddingRowwiseParallel):
-    """Embedding sharded on its hidden dim: weight → Shard(1), so each rank looks up a slice of every vector; the
-    token ids stay replicated, and the slices are gathered into the full output."""
-
-    shard_dim = 1
-    local_output_placement = Shard(-1)
-
-    def validate_param(self, module, param, mesh, parameter_name=None):
-        meta = module._parameters.get(param)
-        if meta is not None and meta.shape[1] % mesh.size() != 0:
-            layer_name = (parameter_name or param).rsplit(".", 1)[0]
-            raise ValueError(
-                f"The embedding size of `{layer_name}` ({meta.shape[1]}) must be divisible by the tensor parallel "
-                f"size ({mesh.size()})."
-            )
-
-
 class ReplicatedWithGradAllReduce(TensorParallelLayer):
     """Replicated parameter whose gradient is partial.
 
@@ -994,8 +976,8 @@ class ParallelInterface(GeneralInterface):
 
     _global_mapping = (
         {
-            "embedding_rowwise": EmbeddingRowwiseParallel(),
-            "embedding_colwise": EmbeddingColwiseParallel(),
+            "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)),
diff --git a/tests/tensor_parallel/test_tensor_parallel.py b/tests/tensor_parallel/test_tensor_parallel.py
index 9cee944a5587..0142b67191a7 100644
--- a/tests/tensor_parallel/test_tensor_parallel.py
+++ b/tests/tensor_parallel/test_tensor_parallel.py
@@ -247,13 +247,15 @@ def test_linear_styles_reject_embedding(self):
                 self.assertIn("linear layers", str(context.exception))
 
     def test_embedding_uneven_local_shapes(self):
-        embedding = torch.nn.Embedding(10, 10)
-        placement = self._get_parameter_placements(embedding, ALL_PARALLEL_STYLES["embedding_rowwise"])["weight"]
+        rowwise = self._get_parameter_placements(torch.nn.Embedding(10, 10), ALL_PARALLEL_STYLES["embedding_rowwise"])
+        colwise = self._get_parameter_placements(torch.nn.Embedding(10, 10), ALL_PARALLEL_STYLES["embedding_colwise"])
 
         expected_local_sizes = (4, 4, 2)
         for rank, expected_size in enumerate(expected_local_sizes):
-            local_shape = self._get_local_shape((10, 10), placement, world_size=3, rank=rank)
-            self.assertEqual(local_shape, (expected_size, 10))
+            rowwise_shape = self._get_local_shape((10, 10), rowwise["weight"], world_size=3, rank=rank)
+            colwise_shape = self._get_local_shape((10, 10), colwise["weight"], world_size=3, rank=rank)
+            self.assertEqual(rowwise_shape, (expected_size, 10))
+            self.assertEqual(colwise_shape, (10, expected_size))
 
     def test_shard_tensor_shape_consistency(self):
         world_size = 4

From 94d8f784c168c21537cffcc624a8834bf82e6648 Mon Sep 17 00:00:00 2001
From: IlyasMoutawwakil <moutawwakil.ilyas.tsi@gmail.com>
Date: Tue, 13 Oct 2026 16:32:14 +0200
Subject: [PATCH 04/13] TP: reject packed weights whose chunks don't divide by
 the number of ranks

The strided split of a packed weight (an MoE experts' fused gate_up_proj) can't hold uneven
slices, so fail up front with a clear error instead of building a misshapen DTensor.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---
 .../distributed/tensor_parallel.py            | 23 +++++++++++++++++++
 tests/tensor_parallel/test_tensor_parallel.py | 16 +++++++++++++
 2 files changed, 39 insertions(+)

diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py
index d07dd912a112..575fc6899f5e 100644
--- a/src/transformers/distributed/tensor_parallel.py
+++ b/src/transformers/distributed/tensor_parallel.py
@@ -498,6 +498,17 @@ def transform_output_post_forward(self, module, output, mesh):
         return output
 
 
+def _validate_packed_split(meta, dim, split_factor, mesh, parameter_name):
+    """A packed weight shards each of its `split_factor` chunks across the mesh's ranks, so each chunk must divide by
+    their number: unlike plain colwise/rowwise, the strided split can't hold uneven slices."""
+    size = meta.shape[dim]
+    if size % (split_factor * mesh.size()) != 0:
+        raise ValueError(
+            f"`{parameter_name}` packs {split_factor} chunks into {size} features: each chunk must divide by the "
+            f"number of ranks ({mesh.size()})."
+        )
+
+
 class PackedColwiseParallel(TensorParallelLayer):
     """Column-wise parallel style for fused linear weights packed along the output dimension."""
 
@@ -523,6 +534,12 @@ def _packed_output_shard_dim(self, param_ndim: int) -> int:
             return -1
         return param_ndim - 2
 
+    def validate_param(self, module, param, mesh, parameter_name=None):
+        meta = module._parameters.get(param)
+        if meta is not None:
+            dim = self._packed_output_shard_dim(meta.ndim)
+            _validate_packed_split(meta, dim, self.split_factor, mesh, parameter_name or param)
+
     def shard_param(self, module, param, mesh):
         meta = module._parameters.get(param)
         if meta is None:
@@ -566,6 +583,12 @@ class PackedRowwiseParallel(TensorParallelLayer):
     def __init__(self, *, split_factor: int = 2):
         self.split_factor = split_factor
 
+    def validate_param(self, module, param, mesh, parameter_name=None):
+        meta = module._parameters.get(param)
+        # a 1D parameter is replicated
+        if meta is not None and meta.ndim > 1:
+            _validate_packed_split(meta, -1, self.split_factor, mesh, parameter_name or param)
+
     def shard_param(self, module, param, mesh):
         meta = module._parameters.get(param)
         if meta is None:
diff --git a/tests/tensor_parallel/test_tensor_parallel.py b/tests/tensor_parallel/test_tensor_parallel.py
index 0142b67191a7..533d586171d9 100644
--- a/tests/tensor_parallel/test_tensor_parallel.py
+++ b/tests/tensor_parallel/test_tensor_parallel.py
@@ -246,6 +246,22 @@ def test_linear_styles_reject_embedding(self):
                     ALL_PARALLEL_STYLES[style].install_forward(torch.nn.Embedding(8, 16), mesh=None)
                 self.assertIn("linear layers", str(context.exception))
 
+    def test_packed_styles_reject_indivisible_chunks(self):
+        # 2 packed chunks of 3 features: each chunk can't split across 2 ranks
+        module = torch.nn.Module()
+        module.register_parameter("weight", torch.nn.Parameter(torch.empty(6, 6)))
+        mesh = self.MockDeviceMesh(world_size=2, rank=0)
+        for style in ("packed_colwise", "packed_rowwise"):
+            with self.subTest(style=style):
+                with self.assertRaises(ValueError) as context:
+                    ALL_PARALLEL_STYLES[style].validate_param(
+                        module, "weight", mesh, parameter_name="gate_up_proj.weight"
+                    )
+                self.assertIn("gate_up_proj.weight", str(context.exception))
+        module.weight = torch.nn.Parameter(torch.empty(8, 8))
+        for style in ("packed_colwise", "packed_rowwise"):
+            ALL_PARALLEL_STYLES[style].validate_param(module, "weight", mesh, parameter_name="gate_up_proj.weight")
+
     def test_embedding_uneven_local_shapes(self):
         rowwise = self._get_parameter_placements(torch.nn.Embedding(10, 10), ALL_PARALLEL_STYLES["embedding_rowwise"])
         colwise = self._get_parameter_placements(torch.nn.Embedding(10, 10), ALL_PARALLEL_STYLES["embedding_colwise"])

From 71929e47a0cc7d3e0dfc654e7716766382f684f8 Mon Sep 17 00:00:00 2001
From: IlyasMoutawwakil <moutawwakil.ilyas.tsi@gmail.com>
Date: Tue, 13 Oct 2026 16:33:14 +0200
Subject: [PATCH 05/13] TP: inline the packed divisibility check into both
 packed styles

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---
 .../distributed/tensor_parallel.py            | 31 +++++++++----------
 1 file changed, 14 insertions(+), 17 deletions(-)

diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py
index 575fc6899f5e..aa1bd5c1f9df 100644
--- a/src/transformers/distributed/tensor_parallel.py
+++ b/src/transformers/distributed/tensor_parallel.py
@@ -498,17 +498,6 @@ def transform_output_post_forward(self, module, output, mesh):
         return output
 
 
-def _validate_packed_split(meta, dim, split_factor, mesh, parameter_name):
-    """A packed weight shards each of its `split_factor` chunks across the mesh's ranks, so each chunk must divide by
-    their number: unlike plain colwise/rowwise, the strided split can't hold uneven slices."""
-    size = meta.shape[dim]
-    if size % (split_factor * mesh.size()) != 0:
-        raise ValueError(
-            f"`{parameter_name}` packs {split_factor} chunks into {size} features: each chunk must divide by the "
-            f"number of ranks ({mesh.size()})."
-        )
-
-
 class PackedColwiseParallel(TensorParallelLayer):
     """Column-wise parallel style for fused linear weights packed along the output dimension."""
 
@@ -536,9 +525,14 @@ def _packed_output_shard_dim(self, param_ndim: int) -> int:
 
     def validate_param(self, module, param, mesh, parameter_name=None):
         meta = module._parameters.get(param)
-        if meta is not None:
-            dim = self._packed_output_shard_dim(meta.ndim)
-            _validate_packed_split(meta, dim, self.split_factor, mesh, parameter_name or param)
+        if meta is None:
+            return
+        # each packed chunk is split across the ranks: the strided split can't hold uneven slices
+        if meta.shape[self._packed_output_shard_dim(meta.ndim)] % (self.split_factor * mesh.size()):
+            raise ValueError(
+                f"`{parameter_name or param}` packs {self.split_factor} chunks: each must divide by the number of "
+                f"ranks ({mesh.size()})."
+            )
 
     def shard_param(self, module, param, mesh):
         meta = module._parameters.get(param)
@@ -585,9 +579,12 @@ def __init__(self, *, split_factor: int = 2):
 
     def validate_param(self, module, param, mesh, parameter_name=None):
         meta = module._parameters.get(param)
-        # a 1D parameter is replicated
-        if meta is not None and meta.ndim > 1:
-            _validate_packed_split(meta, -1, self.split_factor, mesh, parameter_name or param)
+        # each packed chunk is split across the ranks (a 1D parameter is replicated)
+        if meta is not None and meta.ndim > 1 and meta.shape[-1] % (self.split_factor * mesh.size()):
+            raise ValueError(
+                f"`{parameter_name or param}` packs {self.split_factor} chunks: each must divide by the number of "
+                f"ranks ({mesh.size()})."
+            )
 
     def shard_param(self, module, param, mesh):
         meta = module._parameters.get(param)

From 087c77f4505a5dc884688b993e6c21127d866e25 Mon Sep 17 00:00:00 2001
From: IlyasMoutawwakil <moutawwakil.ilyas.tsi@gmail.com>
Date: Tue, 13 Oct 2026 16:45:14 +0200
Subject: [PATCH 06/13] TP: let packed weights shard unevenly too

The packed-weight divisibility check was unneeded: DTensor's strided split (with
`src_data_rank=None`, as transformers distributes) and the checkpoint loader both give each rank
matching ceil-sized slices of every packed chunk, so an uneven MoE `gate_up_proj` works. Verified
with a TP forward and backward of a tiny Qwen3-MoE whose `moe_intermediate_size=36` splits 8/8/8/8/4
over 5 ranks. The TP mixin's packed-gradient slicing now uses the same ceil-sized shards.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---
 .../distributed/tensor_parallel.py            | 20 -------------------
 tests/tensor_parallel/test_tensor_parallel.py | 16 ---------------
 tests/test_tensor_parallel_mixin.py           |  7 ++++---
 3 files changed, 4 insertions(+), 39 deletions(-)

diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py
index aa1bd5c1f9df..d07dd912a112 100644
--- a/src/transformers/distributed/tensor_parallel.py
+++ b/src/transformers/distributed/tensor_parallel.py
@@ -523,17 +523,6 @@ def _packed_output_shard_dim(self, param_ndim: int) -> int:
             return -1
         return param_ndim - 2
 
-    def validate_param(self, module, param, mesh, parameter_name=None):
-        meta = module._parameters.get(param)
-        if meta is None:
-            return
-        # each packed chunk is split across the ranks: the strided split can't hold uneven slices
-        if meta.shape[self._packed_output_shard_dim(meta.ndim)] % (self.split_factor * mesh.size()):
-            raise ValueError(
-                f"`{parameter_name or param}` packs {self.split_factor} chunks: each must divide by the number of "
-                f"ranks ({mesh.size()})."
-            )
-
     def shard_param(self, module, param, mesh):
         meta = module._parameters.get(param)
         if meta is None:
@@ -577,15 +566,6 @@ class PackedRowwiseParallel(TensorParallelLayer):
     def __init__(self, *, split_factor: int = 2):
         self.split_factor = split_factor
 
-    def validate_param(self, module, param, mesh, parameter_name=None):
-        meta = module._parameters.get(param)
-        # each packed chunk is split across the ranks (a 1D parameter is replicated)
-        if meta is not None and meta.ndim > 1 and meta.shape[-1] % (self.split_factor * mesh.size()):
-            raise ValueError(
-                f"`{parameter_name or param}` packs {self.split_factor} chunks: each must divide by the number of "
-                f"ranks ({mesh.size()})."
-            )
-
     def shard_param(self, module, param, mesh):
         meta = module._parameters.get(param)
         if meta is None:
diff --git a/tests/tensor_parallel/test_tensor_parallel.py b/tests/tensor_parallel/test_tensor_parallel.py
index 533d586171d9..0142b67191a7 100644
--- a/tests/tensor_parallel/test_tensor_parallel.py
+++ b/tests/tensor_parallel/test_tensor_parallel.py
@@ -246,22 +246,6 @@ def test_linear_styles_reject_embedding(self):
                     ALL_PARALLEL_STYLES[style].install_forward(torch.nn.Embedding(8, 16), mesh=None)
                 self.assertIn("linear layers", str(context.exception))
 
-    def test_packed_styles_reject_indivisible_chunks(self):
-        # 2 packed chunks of 3 features: each chunk can't split across 2 ranks
-        module = torch.nn.Module()
-        module.register_parameter("weight", torch.nn.Parameter(torch.empty(6, 6)))
-        mesh = self.MockDeviceMesh(world_size=2, rank=0)
-        for style in ("packed_colwise", "packed_rowwise"):
-            with self.subTest(style=style):
-                with self.assertRaises(ValueError) as context:
-                    ALL_PARALLEL_STYLES[style].validate_param(
-                        module, "weight", mesh, parameter_name="gate_up_proj.weight"
-                    )
-                self.assertIn("gate_up_proj.weight", str(context.exception))
-        module.weight = torch.nn.Parameter(torch.empty(8, 8))
-        for style in ("packed_colwise", "packed_rowwise"):
-            ALL_PARALLEL_STYLES[style].validate_param(module, "weight", mesh, parameter_name="gate_up_proj.weight")
-
     def test_embedding_uneven_local_shapes(self):
         rowwise = self._get_parameter_placements(torch.nn.Embedding(10, 10), ALL_PARALLEL_STYLES["embedding_rowwise"])
         colwise = self._get_parameter_placements(torch.nn.Embedding(10, 10), ALL_PARALLEL_STYLES["embedding_colwise"])
diff --git a/tests/test_tensor_parallel_mixin.py b/tests/test_tensor_parallel_mixin.py
index 128c15bcfc77..d385b931b589 100644
--- a/tests/test_tensor_parallel_mixin.py
+++ b/tests/test_tensor_parallel_mixin.py
@@ -92,14 +92,15 @@ def get_packed_grad_shard(grad, world_size, rank, dim):
     total_size = grad.shape[dim]
     # Packed weights have 2 blocks (gate and up)
     block_size = total_size // 2
-    shard_block_size = block_size // world_size
+    # every shard but the last is ceil(block_size / world_size) long, like DTensor's split
+    shard_block_size = -(-block_size // world_size)
 
     # Build interleaved indices
     indices = []
     for block_idx in range(2):  # gate block, then up block
         block_offset = block_idx * block_size
-        start = block_offset + rank * shard_block_size
-        stop = block_offset + (rank + 1) * shard_block_size
+        start = block_offset + min(rank * shard_block_size, block_size)
+        stop = block_offset + min((rank + 1) * shard_block_size, block_size)
         indices.extend(range(start, stop))
 
     # Select along the sharded dimension

From 0f4c42d0d1fc59865e3ef902971ca4ac57e34758 Mon Sep 17 00:00:00 2001
From: IlyasMoutawwakil <moutawwakil.ilyas.tsi@gmail.com>
Date: Wed, 14 Oct 2026 10:08:10 +0200
Subject: [PATCH 07/13] TP: shard every model's input embeddings with
 `embedding_rowwise`

Every TP plan now vocab-shards the model's input embeddings: 89 plans with a plain `nn.Embedding`,
the scaled-embedding family (Gemma*, VaultGemma, MiniCPM3, NeoMME, T5Gemma*), Qwen4-Exp's
`embed_tokens`, and the models whose only plan was a colwise `lm_head` (AFMoE, DBRX, GPT-OSS,
HunYuan, LFM2, PhiMoE). Replicated, the embedding cost `vocab x hidden` on every rank, and a tied
`lm_head` sharded on the vocabulary did not match the embedding weight it shares.

`test_tp_plan_shards_input_embeddings` keeps every model with a TP plan sharding its embeddings.
The models that now get TP tests for the first time skip what can't pass: PhiMoE's backward (it
routes by Gumbel sampling in training) and DBRX's tied forward (it doesn't support tying).

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---
 .../models/afmoe/configuration_afmoe.py        |  1 +
 .../models/apertus/configuration_apertus.py    |  1 +
 .../models/apertus/modular_apertus.py          |  1 +
 .../models/arcee/configuration_arcee.py        |  1 +
 src/transformers/models/arcee/modular_arcee.py |  1 +
 .../models/aria/configuration_aria.py          |  1 +
 src/transformers/models/aria/modular_aria.py   |  1 +
 .../models/axk1/configuration_axk1.py          |  1 +
 .../models/axk2/configuration_axk2.py          |  1 +
 src/transformers/models/axk2/modular_axk2.py   |  1 +
 .../models/cohere/configuration_cohere.py      |  1 +
 .../models/cohere2/configuration_cohere2.py    |  1 +
 .../models/cohere2/modular_cohere2.py          |  1 +
 .../cohere2_moe/configuration_cohere2_moe.py   |  1 +
 .../configuration_cohere_compass.py            |  1 +
 .../cosmos3_edge/configuration_cosmos3_edge.py |  1 +
 .../cosmos3_edge/modular_cosmos3_edge.py       |  1 +
 .../models/cwm/configuration_cwm.py            |  1 +
 .../models/dbrx/configuration_dbrx.py          |  1 +
 .../configuration_deepseek_ocr2.py             |  2 ++
 .../deepseek_ocr2/modular_deepseek_ocr2.py     |  1 +
 .../deepseek_v2/configuration_deepseek_v2.py   |  1 +
 .../models/deepseek_v2/modular_deepseek_v2.py  |  1 +
 .../deepseek_v3/configuration_deepseek_v3.py   |  1 +
 .../deepseek_v32/configuration_deepseek_v32.py |  1 +
 .../deepseek_v32/modular_deepseek_v32.py       |  1 +
 .../configuration_diffusion_gemma.py           |  1 +
 .../models/dots1/configuration_dots1.py        |  1 +
 src/transformers/models/dots1/modular_dots1.py |  1 +
 .../configuration_embedding_gemma2.py          |  1 +
 .../modular_embedding_gemma2.py                |  1 +
 .../models/ernie4_5/configuration_ernie4_5.py  |  1 +
 .../ernie4_5_moe/configuration_ernie4_5_moe.py |  1 +
 .../configuration_ernie4_5_vl_moe.py           |  1 +
 .../ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py |  1 +
 .../models/esmc/configuration_esmc.py          |  1 +
 .../models/eurobert/configuration_eurobert.py  |  1 +
 .../models/exaone4/configuration_exaone4.py    |  1 +
 .../models/exaone4/modular_exaone4.py          |  1 +
 .../exaone_moe/configuration_exaone_moe.py     |  1 +
 .../flex_olmo/configuration_flex_olmo.py       |  1 +
 .../models/flex_olmo/modular_flex_olmo.py      |  1 +
 .../models/gemma/configuration_gemma.py        |  1 +
 src/transformers/models/gemma/modular_gemma.py |  1 +
 .../models/gemma2/configuration_gemma2.py      |  1 +
 .../models/gemma2/modular_gemma2.py            |  1 +
 .../models/gemma3/configuration_gemma3.py      |  1 +
 .../models/gemma3/modular_gemma3.py            |  1 +
 .../models/gemma3n/configuration_gemma3n.py    |  1 +
 .../models/gemma3n/modular_gemma3n.py          |  1 +
 .../models/gemma4/configuration_gemma4.py      |  1 +
 .../configuration_gemma4_unified.py            |  1 +
 .../gemma4_unified/modular_gemma4_unified.py   |  1 +
 .../models/glm/configuration_glm.py            |  1 +
 .../models/glm4/configuration_glm4.py          |  1 +
 .../models/glm4_moe/configuration_glm4_moe.py  |  1 +
 .../models/glm4_moe/modular_glm4_moe.py        |  1 +
 .../configuration_glm4_moe_lite.py             |  1 +
 .../glm4_moe_lite/modular_glm4_moe_lite.py     |  1 +
 .../models/glm4v/configuration_glm4v.py        |  1 +
 src/transformers/models/glm4v/modular_glm4v.py |  1 +
 .../glm4v_moe/configuration_glm4v_moe.py       |  1 +
 .../models/glm4v_moe/modular_glm4v_moe.py      |  1 +
 .../glm5_next/configuration_glm5_next.py       |  1 +
 .../glm_image/configuration_glm_image.py       |  1 +
 .../glm_moe_dsa/configuration_glm_moe_dsa.py   |  1 +
 .../models/glm_ocr/configuration_glm_ocr.py    |  1 +
 .../models/gpt_oss/configuration_gpt_oss.py    |  1 +
 .../models/granite/configuration_granite.py    |  1 +
 .../configuration_granite4_vision.py           |  1 +
 .../granite_swa/configuration_granite_swa.py   |  1 +
 .../models/granite_swa/modular_granite_swa.py  |  1 +
 .../configuration_granitemoe_swa.py            |  1 +
 .../granitemoe_swa/modular_granitemoe_swa.py   |  1 +
 .../models/helium/configuration_helium.py      |  1 +
 .../configuration_higgs_audio_v2.py            |  1 +
 .../models/hrm_text/configuration_hrm_text.py  |  1 +
 .../models/hrm_text/modular_hrm_text.py        |  1 +
 .../configuration_hunyuan_v1_dense.py          |  1 +
 .../configuration_hunyuan_v1_moe.py            |  1 +
 .../hunyuan_vl/configuration_hunyuan_vl.py     |  1 +
 .../models/hy_v3/configuration_hy_v3.py        |  1 +
 src/transformers/models/hy_v3/modular_hy_v3.py |  1 +
 .../models/hy_v4/configuration_hy_v4.py        |  1 +
 src/transformers/models/hy_v4/modular_hy_v4.py |  1 +
 .../hyperclovax/configuration_hyperclovax.py   |  1 +
 .../models/jais2/configuration_jais2.py        |  1 +
 src/transformers/models/jais2/modular_jais2.py |  1 +
 .../kimi_linear/configuration_kimi_linear.py   |  1 +
 .../models/laguna/configuration_laguna.py      |  1 +
 .../models/laguna/modular_laguna.py            |  1 +
 .../models/lfm2/configuration_lfm2.py          |  1 +
 .../models/lfm2_moe/configuration_lfm2_moe.py  |  1 +
 .../models/llama/configuration_llama.py        |  1 +
 .../configuration_longcat_flash.py             |  1 +
 .../models/mellum/configuration_mellum.py      |  1 +
 .../configuration_mimo_v2_flash.py             |  1 +
 .../mimo_v2_flash/modular_mimo_v2_flash.py     |  1 +
 .../models/minicpm3/configuration_minicpm3.py  |  1 +
 .../models/minicpm3/modular_minicpm3.py        |  1 +
 .../models/minimax/configuration_minimax.py    |  1 +
 .../models/minimax/modular_minimax.py          |  1 +
 .../minimax_m2/configuration_minimax_m2.py     |  1 +
 .../models/minimax_m2/modular_minimax_m2.py    |  1 +
 .../configuration_minimax_m3_vl.py             |  1 +
 .../ministral/configuration_ministral.py       |  1 +
 .../ministral3/configuration_ministral3.py     |  1 +
 .../models/mistral/configuration_mistral.py    |  1 +
 .../models/mistral4/configuration_mistral4.py  |  1 +
 .../models/mixtral/configuration_mixtral.py    |  1 +
 .../models/nanochat/configuration_nanochat.py  |  1 +
 .../models/neomme/configuration_neomme.py      |  1 +
 .../models/olmo/configuration_olmo.py          |  1 +
 .../models/olmo2/configuration_olmo2.py        |  1 +
 src/transformers/models/olmo2/modular_olmo2.py |  1 +
 .../models/olmo3/configuration_olmo3.py        |  1 +
 src/transformers/models/olmo3/modular_olmo3.py |  1 +
 .../olmo_hybrid/configuration_olmo_hybrid.py   |  1 +
 .../models/olmo_hybrid/modular_olmo_hybrid.py  |  1 +
 .../models/olmoe/configuration_olmoe.py        |  1 +
 .../configuration_openai_privacy_filter.py     |  1 +
 .../paddleocr_vl/configuration_paddleocr_vl.py |  1 +
 .../models/phi/configuration_phi.py            |  1 +
 .../models/phi3/configuration_phi3.py          |  1 +
 .../configuration_phi4_multimodal.py           |  1 +
 .../models/phimoe/configuration_phimoe.py      |  1 +
 .../models/qwen2/configuration_qwen2.py        |  1 +
 .../qwen2_5_omni/configuration_qwen2_5_omni.py |  1 +
 .../qwen2_5_omni/modular_qwen2_5_omni.py       |  1 +
 .../qwen2_5_vl/configuration_qwen2_5_vl.py     |  1 +
 .../qwen2_moe/configuration_qwen2_moe.py       |  1 +
 .../models/qwen2_vl/configuration_qwen2_vl.py  |  1 +
 .../models/qwen3/configuration_qwen3.py        |  1 +
 .../models/qwen3_5/configuration_qwen3_5.py    |  1 +
 .../models/qwen3_5/modular_qwen3_5.py          |  1 +
 .../qwen3_5_moe/configuration_qwen3_5_moe.py   |  1 +
 .../models/qwen3_5_moe/modular_qwen3_5_moe.py  |  1 +
 .../qwen3_moe/configuration_qwen3_moe.py       |  1 +
 .../qwen3_next/configuration_qwen3_next.py     |  1 +
 .../configuration_qwen3_omni_moe.py            |  3 +++
 .../qwen3_omni_moe/modular_qwen3_omni_moe.py   |  1 +
 .../qwen3_vl_moe/configuration_qwen3_vl_moe.py |  1 +
 .../qwen3_vl_moe/modular_qwen3_vl_moe.py       |  1 +
 .../qwen4_exp/configuration_qwen4_exp.py       |  1 +
 .../models/qwen4_exp/modular_qwen4_exp.py      |  1 +
 .../models/seed_oss/configuration_seed_oss.py  |  1 +
 .../models/smollm3/configuration_smollm3.py    |  1 +
 .../models/smollm3/modular_smollm3.py          |  1 +
 .../solar_open/configuration_solar_open.py     |  1 +
 .../models/solar_open/modular_solar_open.py    |  1 +
 .../starcoder2/configuration_starcoder2.py     |  1 +
 .../models/step3p7/configuration_step3p7.py    |  1 +
 .../models/step3p7/modular_step3p7.py          |  1 +
 .../models/t5gemma/configuration_t5gemma.py    |  1 +
 .../models/t5gemma2/configuration_t5gemma2.py  |  2 ++
 .../vaultgemma/configuration_vaultgemma.py     |  1 +
 .../configuration_voxtral_realtime.py          |  1 +
 .../models/youtu/configuration_youtu.py        |  1 +
 src/transformers/models/youtu/modular_youtu.py |  1 +
 tests/models/dbrx/test_modeling_dbrx.py        |  4 ++++
 tests/models/phimoe/test_modeling_phimoe.py    |  4 ++++
 tests/test_tensor_parallel_mixin.py            | 18 ++++++++++++++++++
 162 files changed, 189 insertions(+)

diff --git a/src/transformers/models/afmoe/configuration_afmoe.py b/src/transformers/models/afmoe/configuration_afmoe.py
index 40083d09e68a..7f6d0293894c 100644
--- a/src/transformers/models/afmoe/configuration_afmoe.py
+++ b/src/transformers/models/afmoe/configuration_afmoe.py
@@ -53,6 +53,7 @@ class AfmoeConfig(PreTrainedConfig):
     """
 
     model_type = "afmoe"
+    base_model_tp_plan = {"embed_tokens": "embedding_rowwise"}
     keys_to_ignore_at_inference = ["past_key_values"]
 
     # Default pipeline parallel plan for base model
diff --git a/src/transformers/models/apertus/configuration_apertus.py b/src/transformers/models/apertus/configuration_apertus.py
index 1e0122160b30..0894ca731ea5 100644
--- a/src/transformers/models/apertus/configuration_apertus.py
+++ b/src/transformers/models/apertus/configuration_apertus.py
@@ -46,6 +46,7 @@ class ApertusConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     default_theta = 12000000.0
     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",
diff --git a/src/transformers/models/apertus/modular_apertus.py b/src/transformers/models/apertus/modular_apertus.py
index 03ae896f8b07..f771fa6b5510 100644
--- a/src/transformers/models/apertus/modular_apertus.py
+++ b/src/transformers/models/apertus/modular_apertus.py
@@ -64,6 +64,7 @@ class ApertusConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     default_theta = 12000000.0
     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",
diff --git a/src/transformers/models/arcee/configuration_arcee.py b/src/transformers/models/arcee/configuration_arcee.py
index ee711d608204..587612e0293b 100644
--- a/src/transformers/models/arcee/configuration_arcee.py
+++ b/src/transformers/models/arcee/configuration_arcee.py
@@ -46,6 +46,7 @@ class ArceeConfig(PreTrainedConfig):
     model_type = "arcee"
     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",
diff --git a/src/transformers/models/arcee/modular_arcee.py b/src/transformers/models/arcee/modular_arcee.py
index 91cd8e13f1ed..f3e153908722 100644
--- a/src/transformers/models/arcee/modular_arcee.py
+++ b/src/transformers/models/arcee/modular_arcee.py
@@ -50,6 +50,7 @@ class ArceeConfig(LlamaConfig):
 
     model_type = "arcee"
     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",
diff --git a/src/transformers/models/aria/configuration_aria.py b/src/transformers/models/aria/configuration_aria.py
index 6da1884bf8cb..63549636af58 100644
--- a/src/transformers/models/aria/configuration_aria.py
+++ b/src/transformers/models/aria/configuration_aria.py
@@ -41,6 +41,7 @@ class AriaTextConfig(PreTrainedConfig):
     model_type = "aria_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",
diff --git a/src/transformers/models/aria/modular_aria.py b/src/transformers/models/aria/modular_aria.py
index 6e1484f6af03..2ae9aa68136c 100644
--- a/src/transformers/models/aria/modular_aria.py
+++ b/src/transformers/models/aria/modular_aria.py
@@ -81,6 +81,7 @@ class AriaTextConfig(LlamaConfig):
     model_type = "aria_text"
     base_config_key = "text_config"
     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",
diff --git a/src/transformers/models/axk1/configuration_axk1.py b/src/transformers/models/axk1/configuration_axk1.py
index 134ee88b01a6..db38fe7bd0ba 100644
--- a/src/transformers/models/axk1/configuration_axk1.py
+++ b/src/transformers/models/axk1/configuration_axk1.py
@@ -53,6 +53,7 @@ class AXK1Config(PreTrainedConfig):
     model_type = "axk1"
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.mlp.experts.gate_up_proj": "packed_colwise",
         "layers.*.mlp.experts.down_proj": "rowwise",
         "layers.*.mlp.experts": "moe_tp_experts",
diff --git a/src/transformers/models/axk2/configuration_axk2.py b/src/transformers/models/axk2/configuration_axk2.py
index a5f57ff65048..4343ab85ba24 100644
--- a/src/transformers/models/axk2/configuration_axk2.py
+++ b/src/transformers/models/axk2/configuration_axk2.py
@@ -62,6 +62,7 @@ class AXK2Config(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
 
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_gate_proj": "colwise",
         "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
         "layers.*.self_attn.kv_b_proj": "colwise",
diff --git a/src/transformers/models/axk2/modular_axk2.py b/src/transformers/models/axk2/modular_axk2.py
index d18dd84f4d7d..125220f402b6 100644
--- a/src/transformers/models/axk2/modular_axk2.py
+++ b/src/transformers/models/axk2/modular_axk2.py
@@ -123,6 +123,7 @@ class AXK2Config(DeepseekV32Config):
     mlp_bias = AttributeError()
 
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_gate_proj": "colwise",
         "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
         "layers.*.self_attn.kv_b_proj": "colwise",
diff --git a/src/transformers/models/cohere/configuration_cohere.py b/src/transformers/models/cohere/configuration_cohere.py
index 44dbb4c01c20..2909d12972ec 100644
--- a/src/transformers/models/cohere/configuration_cohere.py
+++ b/src/transformers/models/cohere/configuration_cohere.py
@@ -50,6 +50,7 @@ class CohereConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     default_theta = 500000.0
     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",
diff --git a/src/transformers/models/cohere2/configuration_cohere2.py b/src/transformers/models/cohere2/configuration_cohere2.py
index 48c2df360354..4a6b1c750580 100644
--- a/src/transformers/models/cohere2/configuration_cohere2.py
+++ b/src/transformers/models/cohere2/configuration_cohere2.py
@@ -49,6 +49,7 @@ class Cohere2Config(PreTrainedConfig):
     model_type = "cohere2"
     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",
diff --git a/src/transformers/models/cohere2/modular_cohere2.py b/src/transformers/models/cohere2/modular_cohere2.py
index 9fd81c926971..0c5f9dd6ac0d 100644
--- a/src/transformers/models/cohere2/modular_cohere2.py
+++ b/src/transformers/models/cohere2/modular_cohere2.py
@@ -69,6 +69,7 @@ class Cohere2Config(PreTrainedConfig):
     model_type = "cohere2"
     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",
diff --git a/src/transformers/models/cohere2_moe/configuration_cohere2_moe.py b/src/transformers/models/cohere2_moe/configuration_cohere2_moe.py
index e94bed7ac32c..0ed5318195a4 100644
--- a/src/transformers/models/cohere2_moe/configuration_cohere2_moe.py
+++ b/src/transformers/models/cohere2_moe/configuration_cohere2_moe.py
@@ -71,6 +71,7 @@ class Cohere2MoeConfig(PreTrainedConfig):
     model_type = "cohere2_moe"
     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",
diff --git a/src/transformers/models/cohere_compass/configuration_cohere_compass.py b/src/transformers/models/cohere_compass/configuration_cohere_compass.py
index 74459f7a8ea2..68bd2d7ed0b9 100644
--- a/src/transformers/models/cohere_compass/configuration_cohere_compass.py
+++ b/src/transformers/models/cohere_compass/configuration_cohere_compass.py
@@ -70,6 +70,7 @@ class CohereCompassTextConfig(PreTrainedConfig):
     model_type = "cohere_compass_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",
diff --git a/src/transformers/models/cosmos3_edge/configuration_cosmos3_edge.py b/src/transformers/models/cosmos3_edge/configuration_cosmos3_edge.py
index 76790b3d135b..838775bc868c 100644
--- a/src/transformers/models/cosmos3_edge/configuration_cosmos3_edge.py
+++ b/src/transformers/models/cosmos3_edge/configuration_cosmos3_edge.py
@@ -43,6 +43,7 @@ class Cosmos3EdgeTextConfig(PreTrainedConfig):
     model_type = "cosmos3_edge_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",
diff --git a/src/transformers/models/cosmos3_edge/modular_cosmos3_edge.py b/src/transformers/models/cosmos3_edge/modular_cosmos3_edge.py
index def47bc67e31..5172b29876aa 100644
--- a/src/transformers/models/cosmos3_edge/modular_cosmos3_edge.py
+++ b/src/transformers/models/cosmos3_edge/modular_cosmos3_edge.py
@@ -89,6 +89,7 @@ class Cosmos3EdgeTextConfig(LlamaConfig):
     default_theta = 100_000_000.0
     ignore_keys_at_rope_validation = {"mrope_section"}
     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",
diff --git a/src/transformers/models/cwm/configuration_cwm.py b/src/transformers/models/cwm/configuration_cwm.py
index ecc3743da19d..6a8c6c563ada 100644
--- a/src/transformers/models/cwm/configuration_cwm.py
+++ b/src/transformers/models/cwm/configuration_cwm.py
@@ -46,6 +46,7 @@ class CwmConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `CwmModel`
     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",
diff --git a/src/transformers/models/dbrx/configuration_dbrx.py b/src/transformers/models/dbrx/configuration_dbrx.py
index 17c57149d545..41ea4fdb234b 100644
--- a/src/transformers/models/dbrx/configuration_dbrx.py
+++ b/src/transformers/models/dbrx/configuration_dbrx.py
@@ -111,6 +111,7 @@ class DbrxConfig(PreTrainedConfig):
     """
 
     model_type = "dbrx"
+    base_model_tp_plan = {"wte": "embedding_rowwise"}
     sub_configs = {"attn_config": DbrxAttentionConfig, "ffn_config": DbrxFFNConfig}
     attribute_map = {
         "num_attention_heads": "n_heads",
diff --git a/src/transformers/models/deepseek_ocr2/configuration_deepseek_ocr2.py b/src/transformers/models/deepseek_ocr2/configuration_deepseek_ocr2.py
index 7adb7745d895..02a9c5a6509d 100644
--- a/src/transformers/models/deepseek_ocr2/configuration_deepseek_ocr2.py
+++ b/src/transformers/models/deepseek_ocr2/configuration_deepseek_ocr2.py
@@ -91,6 +91,7 @@ class DeepseekOcr2VisionEncoderConfig(PreTrainedConfig):
 
     # Default tensor parallel plan for base model `DeepseekOcr2VisionEncoder`
     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",
@@ -182,6 +183,7 @@ class DeepseekOcr2TextConfig(PreTrainedConfig):
 
     # Override DeepseekV2's MLA TP plan with standard MHA projections
     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",
diff --git a/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py b/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py
index f06bb355700e..03ac37bdda64 100644
--- a/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py
+++ b/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py
@@ -571,6 +571,7 @@ class DeepseekOcr2TextConfig(DeepseekV2Config):
 
     # Override DeepseekV2's MLA TP plan with standard MHA projections
     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",
diff --git a/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py b/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py
index 0b9093626489..f12ea083878b 100644
--- a/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py
+++ b/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py
@@ -52,6 +52,7 @@ class DeepseekV2Config(PreTrainedConfig):
     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.q_b_proj": "colwise",
         "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
diff --git a/src/transformers/models/deepseek_v2/modular_deepseek_v2.py b/src/transformers/models/deepseek_v2/modular_deepseek_v2.py
index f31d8d4528ca..b1381bf921ab 100644
--- a/src/transformers/models/deepseek_v2/modular_deepseek_v2.py
+++ b/src/transformers/models/deepseek_v2/modular_deepseek_v2.py
@@ -70,6 +70,7 @@ class DeepseekV2Config(LlamaConfig):
     """
 
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_proj": "colwise",
         "layers.*.self_attn.q_b_proj": "colwise",
         "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
diff --git a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py
index a030c1db5d03..25c71936f5bb 100644
--- a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py
+++ b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py
@@ -53,6 +53,7 @@ class DeepseekV3Config(PreTrainedConfig):
     model_type = "deepseek_v3"
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.mlp.experts.gate_up_proj": "packed_colwise",
         "layers.*.mlp.experts.down_proj": "rowwise",
         "layers.*.mlp.experts": "moe_tp_experts",
diff --git a/src/transformers/models/deepseek_v32/configuration_deepseek_v32.py b/src/transformers/models/deepseek_v32/configuration_deepseek_v32.py
index 48eb2047edcb..a01ca4607fcf 100644
--- a/src/transformers/models/deepseek_v32/configuration_deepseek_v32.py
+++ b/src/transformers/models/deepseek_v32/configuration_deepseek_v32.py
@@ -57,6 +57,7 @@ class DeepseekV32Config(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
 
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_b_proj": "colwise",
         "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
         "layers.*.self_attn.kv_b_proj": "colwise",
diff --git a/src/transformers/models/deepseek_v32/modular_deepseek_v32.py b/src/transformers/models/deepseek_v32/modular_deepseek_v32.py
index 45e979579c8b..bb829b5342d1 100644
--- a/src/transformers/models/deepseek_v32/modular_deepseek_v32.py
+++ b/src/transformers/models/deepseek_v32/modular_deepseek_v32.py
@@ -85,6 +85,7 @@ class DeepseekV32Config(Glm4MoeLiteConfig):
     ```"""
 
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_b_proj": "colwise",
         "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
         "layers.*.self_attn.kv_b_proj": "colwise",
diff --git a/src/transformers/models/diffusion_gemma/configuration_diffusion_gemma.py b/src/transformers/models/diffusion_gemma/configuration_diffusion_gemma.py
index 64460b97be3b..78a79c32b83d 100644
--- a/src/transformers/models/diffusion_gemma/configuration_diffusion_gemma.py
+++ b/src/transformers/models/diffusion_gemma/configuration_diffusion_gemma.py
@@ -47,6 +47,7 @@ class DiffusionGemmaTextConfig(PreTrainedConfig):
     model_type = "diffusion_gemma_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",
diff --git a/src/transformers/models/dots1/configuration_dots1.py b/src/transformers/models/dots1/configuration_dots1.py
index f04022e58df4..944f06dd88c5 100644
--- a/src/transformers/models/dots1/configuration_dots1.py
+++ b/src/transformers/models/dots1/configuration_dots1.py
@@ -48,6 +48,7 @@ class Dots1Config(PreTrainedConfig):
     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",
diff --git a/src/transformers/models/dots1/modular_dots1.py b/src/transformers/models/dots1/modular_dots1.py
index ea3a83ec51ed..47beb2f3e264 100644
--- a/src/transformers/models/dots1/modular_dots1.py
+++ b/src/transformers/models/dots1/modular_dots1.py
@@ -65,6 +65,7 @@ class Dots1Config(PreTrainedConfig):
     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",
diff --git a/src/transformers/models/embedding_gemma2/configuration_embedding_gemma2.py b/src/transformers/models/embedding_gemma2/configuration_embedding_gemma2.py
index b3c173cced61..e0be111fce05 100644
--- a/src/transformers/models/embedding_gemma2/configuration_embedding_gemma2.py
+++ b/src/transformers/models/embedding_gemma2/configuration_embedding_gemma2.py
@@ -51,6 +51,7 @@ class EmbeddingGemma2TextConfig(PreTrainedConfig):
 
     model_type = "embedding_gemma2_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",
diff --git a/src/transformers/models/embedding_gemma2/modular_embedding_gemma2.py b/src/transformers/models/embedding_gemma2/modular_embedding_gemma2.py
index fe8fefbaf443..c3f2f2e4620b 100644
--- a/src/transformers/models/embedding_gemma2/modular_embedding_gemma2.py
+++ b/src/transformers/models/embedding_gemma2/modular_embedding_gemma2.py
@@ -83,6 +83,7 @@ class EmbeddingGemma2TextConfig(PreTrainedConfig):
 
     model_type = "embedding_gemma2_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",
diff --git a/src/transformers/models/ernie4_5/configuration_ernie4_5.py b/src/transformers/models/ernie4_5/configuration_ernie4_5.py
index 896fb0b99402..ec422f957085 100644
--- a/src/transformers/models/ernie4_5/configuration_ernie4_5.py
+++ b/src/transformers/models/ernie4_5/configuration_ernie4_5.py
@@ -47,6 +47,7 @@ class Ernie4_5Config(PreTrainedConfig):
     default_theta = 500000.0
     # Default tensor parallel plan for base model `Ernie4_5Model`
     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",
diff --git a/src/transformers/models/ernie4_5_moe/configuration_ernie4_5_moe.py b/src/transformers/models/ernie4_5_moe/configuration_ernie4_5_moe.py
index 142ff7a2b25c..f536a6b66390 100644
--- a/src/transformers/models/ernie4_5_moe/configuration_ernie4_5_moe.py
+++ b/src/transformers/models/ernie4_5_moe/configuration_ernie4_5_moe.py
@@ -63,6 +63,7 @@ class Ernie4_5_MoeConfig(PreTrainedConfig):
 
     # Default tensor parallel plan for base model `Ernie4_5_MoE`
     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",
diff --git a/src/transformers/models/ernie4_5_vl_moe/configuration_ernie4_5_vl_moe.py b/src/transformers/models/ernie4_5_vl_moe/configuration_ernie4_5_vl_moe.py
index 3cb2029656c8..8cfb60726553 100644
--- a/src/transformers/models/ernie4_5_vl_moe/configuration_ernie4_5_vl_moe.py
+++ b/src/transformers/models/ernie4_5_vl_moe/configuration_ernie4_5_vl_moe.py
@@ -86,6 +86,7 @@ class Ernie4_5_VLMoeTextConfig(PreTrainedConfig):
     default_theta = 500000.0
 
     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",
diff --git a/src/transformers/models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py b/src/transformers/models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py
index 883b09855edf..e823b4ef7ae7 100644
--- a/src/transformers/models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py
+++ b/src/transformers/models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py
@@ -120,6 +120,7 @@ class Ernie4_5_VLMoeTextConfig(Ernie4_5_MoeConfig):
     base_config_key = "text_config"
 
     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",
diff --git a/src/transformers/models/esmc/configuration_esmc.py b/src/transformers/models/esmc/configuration_esmc.py
index 219091d5001f..038561995f99 100644
--- a/src/transformers/models/esmc/configuration_esmc.py
+++ b/src/transformers/models/esmc/configuration_esmc.py
@@ -53,6 +53,7 @@ class EsmcConfig(PreTrainedConfig):
     model_type = "esmc"
     # Default tensor parallel plan for base model `EsmcModel`
     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",
diff --git a/src/transformers/models/eurobert/configuration_eurobert.py b/src/transformers/models/eurobert/configuration_eurobert.py
index f64c4f7e5a11..189a2354852c 100644
--- a/src/transformers/models/eurobert/configuration_eurobert.py
+++ b/src/transformers/models/eurobert/configuration_eurobert.py
@@ -51,6 +51,7 @@ class EuroBertConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `EuroBertModel`
     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",
diff --git a/src/transformers/models/exaone4/configuration_exaone4.py b/src/transformers/models/exaone4/configuration_exaone4.py
index f29cab8dd8ea..0550473cf7a9 100644
--- a/src/transformers/models/exaone4/configuration_exaone4.py
+++ b/src/transformers/models/exaone4/configuration_exaone4.py
@@ -60,6 +60,7 @@ class Exaone4Config(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `LlamaModel`
     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",
diff --git a/src/transformers/models/exaone4/modular_exaone4.py b/src/transformers/models/exaone4/modular_exaone4.py
index b20b18717d2d..1bd405e9dc92 100644
--- a/src/transformers/models/exaone4/modular_exaone4.py
+++ b/src/transformers/models/exaone4/modular_exaone4.py
@@ -89,6 +89,7 @@ class Exaone4Config(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `LlamaModel`
     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",
diff --git a/src/transformers/models/exaone_moe/configuration_exaone_moe.py b/src/transformers/models/exaone_moe/configuration_exaone_moe.py
index 9cdab1097f8e..13561a4b3325 100644
--- a/src/transformers/models/exaone_moe/configuration_exaone_moe.py
+++ b/src/transformers/models/exaone_moe/configuration_exaone_moe.py
@@ -66,6 +66,7 @@ class ExaoneMoeConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `LlamaModel`
     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",
diff --git a/src/transformers/models/flex_olmo/configuration_flex_olmo.py b/src/transformers/models/flex_olmo/configuration_flex_olmo.py
index 921ef2b2597c..82cd1b44717a 100644
--- a/src/transformers/models/flex_olmo/configuration_flex_olmo.py
+++ b/src/transformers/models/flex_olmo/configuration_flex_olmo.py
@@ -50,6 +50,7 @@ class FlexOlmoConfig(PreTrainedConfig):
     attribute_map = {"num_local_experts": "num_experts"}
     default_theta = 500000.0
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
         "layers.*.self_attn.k_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
         "layers.*.self_attn.v_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
diff --git a/src/transformers/models/flex_olmo/modular_flex_olmo.py b/src/transformers/models/flex_olmo/modular_flex_olmo.py
index 956bea08670c..b623367445b9 100644
--- a/src/transformers/models/flex_olmo/modular_flex_olmo.py
+++ b/src/transformers/models/flex_olmo/modular_flex_olmo.py
@@ -60,6 +60,7 @@ class FlexOlmoConfig(PreTrainedConfig):
     attribute_map = {"num_local_experts": "num_experts"}
     default_theta = 500000.0
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
         "layers.*.self_attn.k_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
         "layers.*.self_attn.v_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
diff --git a/src/transformers/models/gemma/configuration_gemma.py b/src/transformers/models/gemma/configuration_gemma.py
index 67f8ae5f143d..ce0f63e14c75 100644
--- a/src/transformers/models/gemma/configuration_gemma.py
+++ b/src/transformers/models/gemma/configuration_gemma.py
@@ -50,6 +50,7 @@ class GemmaConfig(PreTrainedConfig):
     model_type = "gemma"
     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",
diff --git a/src/transformers/models/gemma/modular_gemma.py b/src/transformers/models/gemma/modular_gemma.py
index 272d32a864f7..b9e314c1bf60 100644
--- a/src/transformers/models/gemma/modular_gemma.py
+++ b/src/transformers/models/gemma/modular_gemma.py
@@ -66,6 +66,7 @@ class GemmaConfig(PreTrainedConfig):
     model_type = "gemma"
     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",
diff --git a/src/transformers/models/gemma2/configuration_gemma2.py b/src/transformers/models/gemma2/configuration_gemma2.py
index 11d7b012099f..3cf0a00ad208 100644
--- a/src/transformers/models/gemma2/configuration_gemma2.py
+++ b/src/transformers/models/gemma2/configuration_gemma2.py
@@ -51,6 +51,7 @@ class Gemma2Config(PreTrainedConfig):
     model_type = "gemma2"
     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",
diff --git a/src/transformers/models/gemma2/modular_gemma2.py b/src/transformers/models/gemma2/modular_gemma2.py
index ede2b994b568..d49df566b5e0 100644
--- a/src/transformers/models/gemma2/modular_gemma2.py
+++ b/src/transformers/models/gemma2/modular_gemma2.py
@@ -77,6 +77,7 @@ class Gemma2Config(PreTrainedConfig):
     model_type = "gemma2"
     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",
diff --git a/src/transformers/models/gemma3/configuration_gemma3.py b/src/transformers/models/gemma3/configuration_gemma3.py
index e160c8e9b0eb..e1f933f1e76a 100644
--- a/src/transformers/models/gemma3/configuration_gemma3.py
+++ b/src/transformers/models/gemma3/configuration_gemma3.py
@@ -55,6 +55,7 @@ class Gemma3TextConfig(PreTrainedConfig):
     model_type = "gemma3_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",
diff --git a/src/transformers/models/gemma3/modular_gemma3.py b/src/transformers/models/gemma3/modular_gemma3.py
index 180a0bbea9cb..2710829160d8 100644
--- a/src/transformers/models/gemma3/modular_gemma3.py
+++ b/src/transformers/models/gemma3/modular_gemma3.py
@@ -91,6 +91,7 @@ class Gemma3TextConfig(Gemma2Config, PreTrainedConfig):
 
     model_type = "gemma3_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",
diff --git a/src/transformers/models/gemma3n/configuration_gemma3n.py b/src/transformers/models/gemma3n/configuration_gemma3n.py
index adb25b426292..21d189fed03b 100644
--- a/src/transformers/models/gemma3n/configuration_gemma3n.py
+++ b/src/transformers/models/gemma3n/configuration_gemma3n.py
@@ -76,6 +76,7 @@ class Gemma3nTextConfig(PreTrainedConfig):
     model_type = "gemma3n_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",
diff --git a/src/transformers/models/gemma3n/modular_gemma3n.py b/src/transformers/models/gemma3n/modular_gemma3n.py
index 565afe3b07cb..5eb022580332 100644
--- a/src/transformers/models/gemma3n/modular_gemma3n.py
+++ b/src/transformers/models/gemma3n/modular_gemma3n.py
@@ -116,6 +116,7 @@ class Gemma3nTextConfig(Gemma3TextConfig):
 
     model_type = "gemma3n_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",
diff --git a/src/transformers/models/gemma4/configuration_gemma4.py b/src/transformers/models/gemma4/configuration_gemma4.py
index 6cc0e2a95fd3..b60a664ba5d3 100644
--- a/src/transformers/models/gemma4/configuration_gemma4.py
+++ b/src/transformers/models/gemma4/configuration_gemma4.py
@@ -120,6 +120,7 @@ class Gemma4TextConfig(PreTrainedConfig):
     model_type = "gemma4_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",
diff --git a/src/transformers/models/gemma4_unified/configuration_gemma4_unified.py b/src/transformers/models/gemma4_unified/configuration_gemma4_unified.py
index fa8538e7fe7a..d5a7e5afdcbe 100644
--- a/src/transformers/models/gemma4_unified/configuration_gemma4_unified.py
+++ b/src/transformers/models/gemma4_unified/configuration_gemma4_unified.py
@@ -104,6 +104,7 @@ class Gemma4UnifiedTextConfig(PreTrainedConfig):
     model_type = "gemma4_unified_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",
diff --git a/src/transformers/models/gemma4_unified/modular_gemma4_unified.py b/src/transformers/models/gemma4_unified/modular_gemma4_unified.py
index 8adcae22652f..6ae414a44c69 100644
--- a/src/transformers/models/gemma4_unified/modular_gemma4_unified.py
+++ b/src/transformers/models/gemma4_unified/modular_gemma4_unified.py
@@ -444,6 +444,7 @@ class Gemma4UnifiedTextConfig(Gemma4TextConfig):
 
     model_type = "gemma4_unified_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",
diff --git a/src/transformers/models/glm/configuration_glm.py b/src/transformers/models/glm/configuration_glm.py
index 98525012a23b..16e54eea0722 100644
--- a/src/transformers/models/glm/configuration_glm.py
+++ b/src/transformers/models/glm/configuration_glm.py
@@ -40,6 +40,7 @@ class GlmConfig(PreTrainedConfig):
     model_type = "glm"
     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",
diff --git a/src/transformers/models/glm4/configuration_glm4.py b/src/transformers/models/glm4/configuration_glm4.py
index f33129607fec..c17bb3f84063 100644
--- a/src/transformers/models/glm4/configuration_glm4.py
+++ b/src/transformers/models/glm4/configuration_glm4.py
@@ -40,6 +40,7 @@ class Glm4Config(PreTrainedConfig):
     model_type = "glm4"
     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",
diff --git a/src/transformers/models/glm4_moe/configuration_glm4_moe.py b/src/transformers/models/glm4_moe/configuration_glm4_moe.py
index 884066913209..be4f01978aab 100644
--- a/src/transformers/models/glm4_moe/configuration_glm4_moe.py
+++ b/src/transformers/models/glm4_moe/configuration_glm4_moe.py
@@ -58,6 +58,7 @@ class Glm4MoeConfig(PreTrainedConfig):
 
     # Default tensor parallel plan for base model `Glm4Moe`
     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",
diff --git a/src/transformers/models/glm4_moe/modular_glm4_moe.py b/src/transformers/models/glm4_moe/modular_glm4_moe.py
index 7bef7d9ca2da..a7c757863297 100644
--- a/src/transformers/models/glm4_moe/modular_glm4_moe.py
+++ b/src/transformers/models/glm4_moe/modular_glm4_moe.py
@@ -71,6 +71,7 @@ class Glm4MoeConfig(PreTrainedConfig):
 
     # Default tensor parallel plan for base model `Glm4Moe`
     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",
diff --git a/src/transformers/models/glm4_moe_lite/configuration_glm4_moe_lite.py b/src/transformers/models/glm4_moe_lite/configuration_glm4_moe_lite.py
index b02df0b43b53..4bd74a037457 100644
--- a/src/transformers/models/glm4_moe_lite/configuration_glm4_moe_lite.py
+++ b/src/transformers/models/glm4_moe_lite/configuration_glm4_moe_lite.py
@@ -52,6 +52,7 @@ class Glm4MoeLiteConfig(PreTrainedConfig):
     model_type = "glm4_moe_lite"
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_b_proj": "colwise",
         "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
         "layers.*.self_attn.kv_b_proj": "colwise",
diff --git a/src/transformers/models/glm4_moe_lite/modular_glm4_moe_lite.py b/src/transformers/models/glm4_moe_lite/modular_glm4_moe_lite.py
index db05ee5e7b6f..09acb163f1d5 100644
--- a/src/transformers/models/glm4_moe_lite/modular_glm4_moe_lite.py
+++ b/src/transformers/models/glm4_moe_lite/modular_glm4_moe_lite.py
@@ -60,6 +60,7 @@ class Glm4MoeLiteConfig(PreTrainedConfig):
     model_type = "glm4_moe_lite"
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_b_proj": "colwise",
         "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
         "layers.*.self_attn.kv_b_proj": "colwise",
diff --git a/src/transformers/models/glm4v/configuration_glm4v.py b/src/transformers/models/glm4v/configuration_glm4v.py
index 594503efc61a..dc6bca78ac05 100644
--- a/src/transformers/models/glm4v/configuration_glm4v.py
+++ b/src/transformers/models/glm4v/configuration_glm4v.py
@@ -93,6 +93,7 @@ class Glm4vTextConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `Glm4v`
     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",
diff --git a/src/transformers/models/glm4v/modular_glm4v.py b/src/transformers/models/glm4v/modular_glm4v.py
index dcfd68bdb28e..ebcb93b23834 100644
--- a/src/transformers/models/glm4v/modular_glm4v.py
+++ b/src/transformers/models/glm4v/modular_glm4v.py
@@ -133,6 +133,7 @@ class Glm4vTextConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `Glm4v`
     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",
diff --git a/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py b/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py
index f1b8ce70ea63..4fad96436b1e 100644
--- a/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py
+++ b/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py
@@ -53,6 +53,7 @@ class Glm4vMoeTextConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `Glm4vMoe`
     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",
diff --git a/src/transformers/models/glm4v_moe/modular_glm4v_moe.py b/src/transformers/models/glm4v_moe/modular_glm4v_moe.py
index 1143f430c4fb..aa75bd70b05f 100644
--- a/src/transformers/models/glm4v_moe/modular_glm4v_moe.py
+++ b/src/transformers/models/glm4v_moe/modular_glm4v_moe.py
@@ -86,6 +86,7 @@ class Glm4vMoeTextConfig(Glm4MoeConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `Glm4vMoe`
     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",
diff --git a/src/transformers/models/glm5_next/configuration_glm5_next.py b/src/transformers/models/glm5_next/configuration_glm5_next.py
index c239123a5b6d..07decbe6730f 100644
--- a/src/transformers/models/glm5_next/configuration_glm5_next.py
+++ b/src/transformers/models/glm5_next/configuration_glm5_next.py
@@ -70,6 +70,7 @@ class Glm5NextTextConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
 
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_b_proj": "colwise",
         "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
         "layers.*.self_attn.kv_b_proj": "colwise",
diff --git a/src/transformers/models/glm_image/configuration_glm_image.py b/src/transformers/models/glm_image/configuration_glm_image.py
index 7305ad3d0b3c..42339bcf6d4f 100644
--- a/src/transformers/models/glm_image/configuration_glm_image.py
+++ b/src/transformers/models/glm_image/configuration_glm_image.py
@@ -105,6 +105,7 @@ class GlmImageTextConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `GlmImage`
     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",
diff --git a/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py
index 7f6f457469ed..d7cc943eb1c8 100644
--- a/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py
+++ b/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py
@@ -62,6 +62,7 @@ class GlmMoeDsaConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
 
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_b_proj": "colwise",
         "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
         "layers.*.self_attn.kv_b_proj": "colwise",
diff --git a/src/transformers/models/glm_ocr/configuration_glm_ocr.py b/src/transformers/models/glm_ocr/configuration_glm_ocr.py
index 3aec6d0a41cb..624c1edeb90f 100644
--- a/src/transformers/models/glm_ocr/configuration_glm_ocr.py
+++ b/src/transformers/models/glm_ocr/configuration_glm_ocr.py
@@ -94,6 +94,7 @@ class GlmOcrTextConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `GlmOcr`
     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",
diff --git a/src/transformers/models/gpt_oss/configuration_gpt_oss.py b/src/transformers/models/gpt_oss/configuration_gpt_oss.py
index ebe7caedbdf6..53edd0253c37 100644
--- a/src/transformers/models/gpt_oss/configuration_gpt_oss.py
+++ b/src/transformers/models/gpt_oss/configuration_gpt_oss.py
@@ -30,6 +30,7 @@ class GptOssConfig(PreTrainedConfig):
     """
 
     model_type = "gpt_oss"
+    base_model_tp_plan = {"embed_tokens": "embedding_rowwise"}
     attribute_map = {
         "num_experts": "num_local_experts",
     }
diff --git a/src/transformers/models/granite/configuration_granite.py b/src/transformers/models/granite/configuration_granite.py
index e026cbbe5ff3..f1e10d0e7062 100644
--- a/src/transformers/models/granite/configuration_granite.py
+++ b/src/transformers/models/granite/configuration_granite.py
@@ -47,6 +47,7 @@ class GraniteConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `GraniteModel`
     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",
diff --git a/src/transformers/models/granite4_vision/configuration_granite4_vision.py b/src/transformers/models/granite4_vision/configuration_granite4_vision.py
index 6f8a26af2119..24ae0410a941 100644
--- a/src/transformers/models/granite4_vision/configuration_granite4_vision.py
+++ b/src/transformers/models/granite4_vision/configuration_granite4_vision.py
@@ -50,6 +50,7 @@ class Granite4VisionTextConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `Granite4VisionTextModel`
     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",
diff --git a/src/transformers/models/granite_swa/configuration_granite_swa.py b/src/transformers/models/granite_swa/configuration_granite_swa.py
index 678edf63afac..7cce18b6c25e 100644
--- a/src/transformers/models/granite_swa/configuration_granite_swa.py
+++ b/src/transformers/models/granite_swa/configuration_granite_swa.py
@@ -58,6 +58,7 @@ class GraniteSWAConfig(PreTrainedConfig):
     # Extends Granite's plan with the learnable per-head `sinks`, sharded across heads (colwise)
     # to match the q/k/v head-sharding so TP keeps each rank's sink slice aligned with its heads.
     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",
diff --git a/src/transformers/models/granite_swa/modular_granite_swa.py b/src/transformers/models/granite_swa/modular_granite_swa.py
index 042d33c47f25..e0e34f1c9287 100644
--- a/src/transformers/models/granite_swa/modular_granite_swa.py
+++ b/src/transformers/models/granite_swa/modular_granite_swa.py
@@ -88,6 +88,7 @@ class GraniteSWAConfig(GraniteConfig):
     # Extends Granite's plan with the learnable per-head `sinks`, sharded across heads (colwise)
     # to match the q/k/v head-sharding so TP keeps each rank's sink slice aligned with its heads.
     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",
diff --git a/src/transformers/models/granitemoe_swa/configuration_granitemoe_swa.py b/src/transformers/models/granitemoe_swa/configuration_granitemoe_swa.py
index 578ab17db32f..df16a882cd7c 100644
--- a/src/transformers/models/granitemoe_swa/configuration_granitemoe_swa.py
+++ b/src/transformers/models/granitemoe_swa/configuration_granitemoe_swa.py
@@ -90,6 +90,7 @@ class GraniteMoeSWAConfig(PreTrainedConfig):
     # with the router replicated. The optional shared expert (`shared_mlp`, off by default) is left
     # replicated -- it is small and its full output sums consistently with the all-reduced MoE output.
     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",
diff --git a/src/transformers/models/granitemoe_swa/modular_granitemoe_swa.py b/src/transformers/models/granitemoe_swa/modular_granitemoe_swa.py
index 63cc0679cbb8..1857feae135e 100644
--- a/src/transformers/models/granitemoe_swa/modular_granitemoe_swa.py
+++ b/src/transformers/models/granitemoe_swa/modular_granitemoe_swa.py
@@ -87,6 +87,7 @@ class GraniteMoeSWAConfig(GraniteMoeSharedConfig):
     # with the router replicated. The optional shared expert (`shared_mlp`, off by default) is left
     # replicated -- it is small and its full output sums consistently with the all-reduced MoE output.
     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",
diff --git a/src/transformers/models/helium/configuration_helium.py b/src/transformers/models/helium/configuration_helium.py
index caa966b58a9d..fe6271f0eee8 100644
--- a/src/transformers/models/helium/configuration_helium.py
+++ b/src/transformers/models/helium/configuration_helium.py
@@ -41,6 +41,7 @@ class HeliumConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     default_theta = 100000.0
     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",
diff --git a/src/transformers/models/higgs_audio_v2/configuration_higgs_audio_v2.py b/src/transformers/models/higgs_audio_v2/configuration_higgs_audio_v2.py
index 97823eb79576..495806a9d154 100644
--- a/src/transformers/models/higgs_audio_v2/configuration_higgs_audio_v2.py
+++ b/src/transformers/models/higgs_audio_v2/configuration_higgs_audio_v2.py
@@ -59,6 +59,7 @@ class HiggsAudioV2Config(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `HiggsAudioV2Model`
     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",
diff --git a/src/transformers/models/hrm_text/configuration_hrm_text.py b/src/transformers/models/hrm_text/configuration_hrm_text.py
index 61d29a8311b0..d86a155ac4b0 100644
--- a/src/transformers/models/hrm_text/configuration_hrm_text.py
+++ b/src/transformers/models/hrm_text/configuration_hrm_text.py
@@ -56,6 +56,7 @@ class HrmTextConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
 
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         **{f"{stack}.layers.*.self_attn.q_proj": "colwise" for stack in ("L_module", "H_module")},
         **{f"{stack}.layers.*.self_attn.k_proj": "colwise" for stack in ("L_module", "H_module")},
         **{f"{stack}.layers.*.self_attn.v_proj": "colwise" for stack in ("L_module", "H_module")},
diff --git a/src/transformers/models/hrm_text/modular_hrm_text.py b/src/transformers/models/hrm_text/modular_hrm_text.py
index 274f9d3291c0..35975aa2b5b9 100644
--- a/src/transformers/models/hrm_text/modular_hrm_text.py
+++ b/src/transformers/models/hrm_text/modular_hrm_text.py
@@ -73,6 +73,7 @@ class HrmTextConfig(LlamaConfig):
     model_type = "hrm_text"
 
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         **{f"{stack}.layers.*.self_attn.q_proj": "colwise" for stack in ("L_module", "H_module")},
         **{f"{stack}.layers.*.self_attn.k_proj": "colwise" for stack in ("L_module", "H_module")},
         **{f"{stack}.layers.*.self_attn.v_proj": "colwise" for stack in ("L_module", "H_module")},
diff --git a/src/transformers/models/hunyuan_v1_dense/configuration_hunyuan_v1_dense.py b/src/transformers/models/hunyuan_v1_dense/configuration_hunyuan_v1_dense.py
index 7ac120cec14f..52fdc8f05a89 100644
--- a/src/transformers/models/hunyuan_v1_dense/configuration_hunyuan_v1_dense.py
+++ b/src/transformers/models/hunyuan_v1_dense/configuration_hunyuan_v1_dense.py
@@ -31,6 +31,7 @@ class HunYuanDenseV1Config(PreTrainedConfig):
     """
 
     model_type = "hunyuan_v1_dense"
+    base_model_tp_plan = {"embed_tokens": "embedding_rowwise"}
     keys_to_ignore_at_inference = ["past_key_values"]
 
     vocab_size: int = 290943
diff --git a/src/transformers/models/hunyuan_v1_moe/configuration_hunyuan_v1_moe.py b/src/transformers/models/hunyuan_v1_moe/configuration_hunyuan_v1_moe.py
index f20b240507b8..5ad0d0bb1244 100644
--- a/src/transformers/models/hunyuan_v1_moe/configuration_hunyuan_v1_moe.py
+++ b/src/transformers/models/hunyuan_v1_moe/configuration_hunyuan_v1_moe.py
@@ -32,6 +32,7 @@ class HunYuanMoEV1Config(PreTrainedConfig):
     """
 
     model_type = "hunyuan_v1_moe"
+    base_model_tp_plan = {"embed_tokens": "embedding_rowwise"}
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_ep_plan = {
         "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
diff --git a/src/transformers/models/hunyuan_vl/configuration_hunyuan_vl.py b/src/transformers/models/hunyuan_vl/configuration_hunyuan_vl.py
index 6a6c2d339a16..b79e2b2cf427 100644
--- a/src/transformers/models/hunyuan_vl/configuration_hunyuan_vl.py
+++ b/src/transformers/models/hunyuan_vl/configuration_hunyuan_vl.py
@@ -119,6 +119,7 @@ class HunYuanVLTextConfig(PreTrainedConfig):
     """
 
     model_type = "hunyuan_vl_text"
+    base_model_tp_plan = {"embed_tokens": "embedding_rowwise"}
     keys_to_ignore_at_inference = ["past_key_values"]
 
     vocab_size: int = 290943
diff --git a/src/transformers/models/hy_v3/configuration_hy_v3.py b/src/transformers/models/hy_v3/configuration_hy_v3.py
index 6266f301095c..807346e23e61 100644
--- a/src/transformers/models/hy_v3/configuration_hy_v3.py
+++ b/src/transformers/models/hy_v3/configuration_hy_v3.py
@@ -51,6 +51,7 @@ class HYV3Config(PreTrainedConfig):
         "num_local_experts": "num_experts",
     }
     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",
diff --git a/src/transformers/models/hy_v3/modular_hy_v3.py b/src/transformers/models/hy_v3/modular_hy_v3.py
index 7aa409123b40..62058204fe6b 100644
--- a/src/transformers/models/hy_v3/modular_hy_v3.py
+++ b/src/transformers/models/hy_v3/modular_hy_v3.py
@@ -76,6 +76,7 @@ class HYV3Config(PreTrainedConfig):
         "num_local_experts": "num_experts",
     }
     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",
diff --git a/src/transformers/models/hy_v4/configuration_hy_v4.py b/src/transformers/models/hy_v4/configuration_hy_v4.py
index ac7750fd9cb3..2131005a3cb5 100644
--- a/src/transformers/models/hy_v4/configuration_hy_v4.py
+++ b/src/transformers/models/hy_v4/configuration_hy_v4.py
@@ -62,6 +62,7 @@ class HYV4Config(PreTrainedConfig):
         "num_local_experts": "n_routed_experts",
     }
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_b_proj": "colwise",
         "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
         "layers.*.self_attn.kv_b_proj": "colwise",
diff --git a/src/transformers/models/hy_v4/modular_hy_v4.py b/src/transformers/models/hy_v4/modular_hy_v4.py
index 85b1aef0d3bb..f8a6fa313ca3 100644
--- a/src/transformers/models/hy_v4/modular_hy_v4.py
+++ b/src/transformers/models/hy_v4/modular_hy_v4.py
@@ -91,6 +91,7 @@ class HYV4Config(PreTrainedConfig):
         "num_local_experts": "n_routed_experts",
     }
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_b_proj": "colwise",
         "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
         "layers.*.self_attn.kv_b_proj": "colwise",
diff --git a/src/transformers/models/hyperclovax/configuration_hyperclovax.py b/src/transformers/models/hyperclovax/configuration_hyperclovax.py
index 430a56bf0249..299e3ccab5e8 100644
--- a/src/transformers/models/hyperclovax/configuration_hyperclovax.py
+++ b/src/transformers/models/hyperclovax/configuration_hyperclovax.py
@@ -62,6 +62,7 @@ class HyperCLOVAXConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `HyperCLOVAXModel`
     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",
diff --git a/src/transformers/models/jais2/configuration_jais2.py b/src/transformers/models/jais2/configuration_jais2.py
index 7139661d7575..d4bc5bba7aaa 100644
--- a/src/transformers/models/jais2/configuration_jais2.py
+++ b/src/transformers/models/jais2/configuration_jais2.py
@@ -47,6 +47,7 @@ class Jais2Config(PreTrainedConfig):
     model_type = "jais2"
     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",
diff --git a/src/transformers/models/jais2/modular_jais2.py b/src/transformers/models/jais2/modular_jais2.py
index 3fbdf2c8cd46..0fd270cdce53 100644
--- a/src/transformers/models/jais2/modular_jais2.py
+++ b/src/transformers/models/jais2/modular_jais2.py
@@ -31,6 +31,7 @@
 @strict
 class Jais2Config(LlamaConfig):
     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",
diff --git a/src/transformers/models/kimi_linear/configuration_kimi_linear.py b/src/transformers/models/kimi_linear/configuration_kimi_linear.py
index 90fbf51356ec..0ce8cb988c86 100644
--- a/src/transformers/models/kimi_linear/configuration_kimi_linear.py
+++ b/src/transformers/models/kimi_linear/configuration_kimi_linear.py
@@ -44,6 +44,7 @@ class KimiLinearConfig(PreTrainedConfig):
     model_type = "kimi_linear"
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.mlp.experts.gate_up_proj": "packed_colwise",
         "layers.*.mlp.experts.down_proj": "rowwise",
         "layers.*.mlp.experts": "moe_tp_experts",
diff --git a/src/transformers/models/laguna/configuration_laguna.py b/src/transformers/models/laguna/configuration_laguna.py
index 024937999e4b..3e09f4107bc2 100644
--- a/src/transformers/models/laguna/configuration_laguna.py
+++ b/src/transformers/models/laguna/configuration_laguna.py
@@ -61,6 +61,7 @@ class LagunaConfig(PreTrainedConfig):
     model_type = "laguna"
     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",
diff --git a/src/transformers/models/laguna/modular_laguna.py b/src/transformers/models/laguna/modular_laguna.py
index 79b7796c5ee3..f787576473a7 100644
--- a/src/transformers/models/laguna/modular_laguna.py
+++ b/src/transformers/models/laguna/modular_laguna.py
@@ -79,6 +79,7 @@ class LagunaConfig(Qwen2MoeConfig):
 
     model_type = "laguna"
     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",
diff --git a/src/transformers/models/lfm2/configuration_lfm2.py b/src/transformers/models/lfm2/configuration_lfm2.py
index f58ce6270d31..6f291a0f1c7d 100644
--- a/src/transformers/models/lfm2/configuration_lfm2.py
+++ b/src/transformers/models/lfm2/configuration_lfm2.py
@@ -52,6 +52,7 @@ class Lfm2Config(PreTrainedConfig):
     """
 
     model_type = "lfm2"
+    base_model_tp_plan = {"embed_tokens": "embedding_rowwise"}
     keys_to_ignore_at_inference = ["past_key_values"]
     default_theta = 1000000.0
 
diff --git a/src/transformers/models/lfm2_moe/configuration_lfm2_moe.py b/src/transformers/models/lfm2_moe/configuration_lfm2_moe.py
index 0e09fdd607cc..f1a7c4759f47 100644
--- a/src/transformers/models/lfm2_moe/configuration_lfm2_moe.py
+++ b/src/transformers/models/lfm2_moe/configuration_lfm2_moe.py
@@ -47,6 +47,7 @@ class Lfm2MoeConfig(PreTrainedConfig):
     """
 
     model_type = "lfm2_moe"
+    base_model_tp_plan = {"embed_tokens": "embedding_rowwise"}
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_ep_plan = {
         "layers.*.feed_forward.experts.gate_up_proj": "grouped_gemm",
diff --git a/src/transformers/models/llama/configuration_llama.py b/src/transformers/models/llama/configuration_llama.py
index 6960a6970592..f6a08f49e37a 100644
--- a/src/transformers/models/llama/configuration_llama.py
+++ b/src/transformers/models/llama/configuration_llama.py
@@ -47,6 +47,7 @@ class LlamaConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `LlamaModel`
     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",
diff --git a/src/transformers/models/longcat_flash/configuration_longcat_flash.py b/src/transformers/models/longcat_flash/configuration_longcat_flash.py
index 5a7101897c3e..2befa78aaefe 100644
--- a/src/transformers/models/longcat_flash/configuration_longcat_flash.py
+++ b/src/transformers/models/longcat_flash/configuration_longcat_flash.py
@@ -63,6 +63,7 @@ class LongcatFlashConfig(PreTrainedConfig):
     }
     default_theta = 10000000.0
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.*.q_b_proj": "colwise",
         "layers.*.self_attn.*.kv_a_proj_with_mqa": "mla_kv_a_proj",
         "layers.*.self_attn.*.kv_b_proj": "colwise",
diff --git a/src/transformers/models/mellum/configuration_mellum.py b/src/transformers/models/mellum/configuration_mellum.py
index 9b3374607297..8c892961485c 100644
--- a/src/transformers/models/mellum/configuration_mellum.py
+++ b/src/transformers/models/mellum/configuration_mellum.py
@@ -51,6 +51,7 @@ class MellumConfig(PreTrainedConfig):
 
     # Default tensor parallel plan for base model `Mellum`
     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",
diff --git a/src/transformers/models/mimo_v2_flash/configuration_mimo_v2_flash.py b/src/transformers/models/mimo_v2_flash/configuration_mimo_v2_flash.py
index e40cc863153a..86a525f45bf3 100644
--- a/src/transformers/models/mimo_v2_flash/configuration_mimo_v2_flash.py
+++ b/src/transformers/models/mimo_v2_flash/configuration_mimo_v2_flash.py
@@ -49,6 +49,7 @@ class MiMoV2FlashConfig(PreTrainedConfig):
     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",
diff --git a/src/transformers/models/mimo_v2_flash/modular_mimo_v2_flash.py b/src/transformers/models/mimo_v2_flash/modular_mimo_v2_flash.py
index 4c5fd1a9483f..674a1d3a2a76 100644
--- a/src/transformers/models/mimo_v2_flash/modular_mimo_v2_flash.py
+++ b/src/transformers/models/mimo_v2_flash/modular_mimo_v2_flash.py
@@ -65,6 +65,7 @@ class MiMoV2FlashConfig(Glm4MoeConfig):
     attribute_map = {"num_local_experts": "n_routed_experts"}
 
     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",
diff --git a/src/transformers/models/minicpm3/configuration_minicpm3.py b/src/transformers/models/minicpm3/configuration_minicpm3.py
index 5367c5497a96..7283613254da 100644
--- a/src/transformers/models/minicpm3/configuration_minicpm3.py
+++ b/src/transformers/models/minicpm3/configuration_minicpm3.py
@@ -68,6 +68,7 @@ class MiniCPM3Config(PreTrainedConfig):
     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.q_b_proj": "colwise",
         "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
diff --git a/src/transformers/models/minicpm3/modular_minicpm3.py b/src/transformers/models/minicpm3/modular_minicpm3.py
index 0e229a39c613..2b5fe15e68d5 100644
--- a/src/transformers/models/minicpm3/modular_minicpm3.py
+++ b/src/transformers/models/minicpm3/modular_minicpm3.py
@@ -82,6 +82,7 @@ class MiniCPM3Config(LlamaConfig):
     """
 
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_proj": "colwise",
         "layers.*.self_attn.q_b_proj": "colwise",
         "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
diff --git a/src/transformers/models/minimax/configuration_minimax.py b/src/transformers/models/minimax/configuration_minimax.py
index a7fd614a1d62..3ac1f05b8074 100644
--- a/src/transformers/models/minimax/configuration_minimax.py
+++ b/src/transformers/models/minimax/configuration_minimax.py
@@ -62,6 +62,7 @@ class MiniMaxConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     default_theta = 1000000.0
     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",
diff --git a/src/transformers/models/minimax/modular_minimax.py b/src/transformers/models/minimax/modular_minimax.py
index 37c41a63c126..5c72d2dae97c 100644
--- a/src/transformers/models/minimax/modular_minimax.py
+++ b/src/transformers/models/minimax/modular_minimax.py
@@ -89,6 +89,7 @@ class MiniMaxConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     default_theta = 1000000.0
     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",
diff --git a/src/transformers/models/minimax_m2/configuration_minimax_m2.py b/src/transformers/models/minimax_m2/configuration_minimax_m2.py
index 3e52eb9f736e..000f8844dae2 100644
--- a/src/transformers/models/minimax_m2/configuration_minimax_m2.py
+++ b/src/transformers/models/minimax_m2/configuration_minimax_m2.py
@@ -48,6 +48,7 @@ class MiniMaxM2Config(PreTrainedConfig):
     model_type = "minimax_m2"
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_proj": "colwise_gather_output",
         "layers.*.self_attn.k_proj": "colwise_gather_output",
         "layers.*.self_attn.v_proj": "colwise_gather_output",
diff --git a/src/transformers/models/minimax_m2/modular_minimax_m2.py b/src/transformers/models/minimax_m2/modular_minimax_m2.py
index 57090e31a927..e29fbc5b82e8 100644
--- a/src/transformers/models/minimax_m2/modular_minimax_m2.py
+++ b/src/transformers/models/minimax_m2/modular_minimax_m2.py
@@ -67,6 +67,7 @@ class MiniMaxM2Config(PreTrainedConfig):
     model_type = "minimax_m2"
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_proj": "colwise_gather_output",
         "layers.*.self_attn.k_proj": "colwise_gather_output",
         "layers.*.self_attn.v_proj": "colwise_gather_output",
diff --git a/src/transformers/models/minimax_m3_vl/configuration_minimax_m3_vl.py b/src/transformers/models/minimax_m3_vl/configuration_minimax_m3_vl.py
index 4069a570f96d..e42b7a975da1 100644
--- a/src/transformers/models/minimax_m3_vl/configuration_minimax_m3_vl.py
+++ b/src/transformers/models/minimax_m3_vl/configuration_minimax_m3_vl.py
@@ -55,6 +55,7 @@ class MiniMaxM3VLTextConfig(PreTrainedConfig):
     model_type = "minimax_m3_vl_text"
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_proj": "colwise_gather_output",
         "layers.*.self_attn.k_proj": "colwise_gather_output",
         "layers.*.self_attn.v_proj": "colwise_gather_output",
diff --git a/src/transformers/models/ministral/configuration_ministral.py b/src/transformers/models/ministral/configuration_ministral.py
index 4ff445e6808a..706968c2872c 100644
--- a/src/transformers/models/ministral/configuration_ministral.py
+++ b/src/transformers/models/ministral/configuration_ministral.py
@@ -49,6 +49,7 @@ class MinistralConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `MinistralModel`
     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",
diff --git a/src/transformers/models/ministral3/configuration_ministral3.py b/src/transformers/models/ministral3/configuration_ministral3.py
index ff0ca8610280..0597792a9ac5 100644
--- a/src/transformers/models/ministral3/configuration_ministral3.py
+++ b/src/transformers/models/ministral3/configuration_ministral3.py
@@ -55,6 +55,7 @@ class Ministral3Config(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `MistralModel`
     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",
diff --git a/src/transformers/models/mistral/configuration_mistral.py b/src/transformers/models/mistral/configuration_mistral.py
index c57193f58d7b..a4d8e020f50d 100644
--- a/src/transformers/models/mistral/configuration_mistral.py
+++ b/src/transformers/models/mistral/configuration_mistral.py
@@ -46,6 +46,7 @@ class MistralConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `MistralModel`
     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",
diff --git a/src/transformers/models/mistral4/configuration_mistral4.py b/src/transformers/models/mistral4/configuration_mistral4.py
index eb3c2f3f787d..0e04e7fac8df 100644
--- a/src/transformers/models/mistral4/configuration_mistral4.py
+++ b/src/transformers/models/mistral4/configuration_mistral4.py
@@ -47,6 +47,7 @@ class Mistral4Config(PreTrainedConfig):
     model_type = "mistral4"
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.mlp.experts.gate_up_proj": "packed_colwise",
         "layers.*.mlp.experts.down_proj": "rowwise",
         "layers.*.mlp.experts": "moe_tp_experts",
diff --git a/src/transformers/models/mixtral/configuration_mixtral.py b/src/transformers/models/mixtral/configuration_mixtral.py
index c1158852cf55..4e57ab005c80 100644
--- a/src/transformers/models/mixtral/configuration_mixtral.py
+++ b/src/transformers/models/mixtral/configuration_mixtral.py
@@ -43,6 +43,7 @@ class MixtralConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     default_theta = 1000000.0
     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",
diff --git a/src/transformers/models/nanochat/configuration_nanochat.py b/src/transformers/models/nanochat/configuration_nanochat.py
index 24a0ab7b6d09..af18cf077772 100644
--- a/src/transformers/models/nanochat/configuration_nanochat.py
+++ b/src/transformers/models/nanochat/configuration_nanochat.py
@@ -43,6 +43,7 @@ class NanoChatConfig(PreTrainedConfig):
     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",
diff --git a/src/transformers/models/neomme/configuration_neomme.py b/src/transformers/models/neomme/configuration_neomme.py
index 6a1506dd1e5f..5418d551fd64 100644
--- a/src/transformers/models/neomme/configuration_neomme.py
+++ b/src/transformers/models/neomme/configuration_neomme.py
@@ -59,6 +59,7 @@ class NeoMMEConfig(PreTrainedConfig):
 
     model_type = "neomme"
     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",
diff --git a/src/transformers/models/olmo/configuration_olmo.py b/src/transformers/models/olmo/configuration_olmo.py
index 186cc3a704fb..7d13bc90b3aa 100644
--- a/src/transformers/models/olmo/configuration_olmo.py
+++ b/src/transformers/models/olmo/configuration_olmo.py
@@ -50,6 +50,7 @@ class OlmoConfig(PreTrainedConfig):
     model_type = "olmo"
     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",
diff --git a/src/transformers/models/olmo2/configuration_olmo2.py b/src/transformers/models/olmo2/configuration_olmo2.py
index a99a49194305..3b4754bf8a24 100644
--- a/src/transformers/models/olmo2/configuration_olmo2.py
+++ b/src/transformers/models/olmo2/configuration_olmo2.py
@@ -53,6 +53,7 @@ class Olmo2Config(PreTrainedConfig):
     model_type = "olmo2"
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
         "layers.*.self_attn.k_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
         "layers.*.self_attn.v_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
diff --git a/src/transformers/models/olmo2/modular_olmo2.py b/src/transformers/models/olmo2/modular_olmo2.py
index 4f8dcffeaae0..d0fbe9213788 100644
--- a/src/transformers/models/olmo2/modular_olmo2.py
+++ b/src/transformers/models/olmo2/modular_olmo2.py
@@ -67,6 +67,7 @@ class Olmo2Config(OlmoConfig):
 
     model_type = "olmo2"
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
         "layers.*.self_attn.k_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
         "layers.*.self_attn.v_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
diff --git a/src/transformers/models/olmo3/configuration_olmo3.py b/src/transformers/models/olmo3/configuration_olmo3.py
index f43d5c02b9c5..d527e08eccf2 100644
--- a/src/transformers/models/olmo3/configuration_olmo3.py
+++ b/src/transformers/models/olmo3/configuration_olmo3.py
@@ -48,6 +48,7 @@ class Olmo3Config(PreTrainedConfig):
     model_type = "olmo3"
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
         "layers.*.self_attn.k_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
         "layers.*.self_attn.v_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
diff --git a/src/transformers/models/olmo3/modular_olmo3.py b/src/transformers/models/olmo3/modular_olmo3.py
index 015a3461a073..b5292f9873e8 100644
--- a/src/transformers/models/olmo3/modular_olmo3.py
+++ b/src/transformers/models/olmo3/modular_olmo3.py
@@ -66,6 +66,7 @@ class Olmo3Config(Olmo2Config):
     default_theta = 500000.0
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
         "layers.*.self_attn.k_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
         "layers.*.self_attn.v_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
diff --git a/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py
index 0f7c32e8799e..5d113e458764 100644
--- a/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py
+++ b/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py
@@ -73,6 +73,7 @@ class OlmoHybridConfig(PreTrainedConfig):
     model_type = "olmo_hybrid"
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
         "layers.*.self_attn.k_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
         "layers.*.self_attn.v_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
diff --git a/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py
index d0981294b01a..b93f0cec3865 100644
--- a/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py
+++ b/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py
@@ -107,6 +107,7 @@ class OlmoHybridConfig(LlamaConfig):
 
     model_type = "olmo_hybrid"
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
         "layers.*.self_attn.k_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
         "layers.*.self_attn.v_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
diff --git a/src/transformers/models/olmoe/configuration_olmoe.py b/src/transformers/models/olmoe/configuration_olmoe.py
index 699703b0d94d..442df02573c4 100644
--- a/src/transformers/models/olmoe/configuration_olmoe.py
+++ b/src/transformers/models/olmoe/configuration_olmoe.py
@@ -46,6 +46,7 @@ class OlmoeConfig(PreTrainedConfig):
 
     # Default tensor parallel plan for base model `Olmoe`
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_proj": "colwise_gather_output",  # due to the norm, we have to gather
         "layers.*.self_attn.k_proj": "colwise_gather_output",  # due to the norm, we have to gather
         "layers.*.self_attn.v_proj": "colwise_gather_output",  # due to the norm, we have to gather
diff --git a/src/transformers/models/openai_privacy_filter/configuration_openai_privacy_filter.py b/src/transformers/models/openai_privacy_filter/configuration_openai_privacy_filter.py
index b256f8aadf22..efc0c81291ed 100644
--- a/src/transformers/models/openai_privacy_filter/configuration_openai_privacy_filter.py
+++ b/src/transformers/models/openai_privacy_filter/configuration_openai_privacy_filter.py
@@ -54,6 +54,7 @@ class OpenAIPrivacyFilterConfig(PreTrainedConfig):
     """
 
     model_type = "openai_privacy_filter"
+    base_model_tp_plan = {"embed_tokens": "embedding_rowwise"}
     attribute_map = {
         "num_experts": "num_local_experts",
     }
diff --git a/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py b/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py
index efcbf9fe7cd2..ecc13765d4d3 100644
--- a/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py
+++ b/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py
@@ -97,6 +97,7 @@ class PaddleOCRTextConfig(PreTrainedConfig):
     default_theta = 500000.0
     # Default tensor parallel plan for base model `PaddleOCRTextModel`
     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",
diff --git a/src/transformers/models/phi/configuration_phi.py b/src/transformers/models/phi/configuration_phi.py
index 2a65f0f16aec..3720214a1247 100644
--- a/src/transformers/models/phi/configuration_phi.py
+++ b/src/transformers/models/phi/configuration_phi.py
@@ -46,6 +46,7 @@ class PhiConfig(PreTrainedConfig):
     model_type = "phi"
     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",
diff --git a/src/transformers/models/phi3/configuration_phi3.py b/src/transformers/models/phi3/configuration_phi3.py
index f85502f8205d..725c2fb232a2 100644
--- a/src/transformers/models/phi3/configuration_phi3.py
+++ b/src/transformers/models/phi3/configuration_phi3.py
@@ -47,6 +47,7 @@ class Phi3Config(PreTrainedConfig):
     model_type = "phi3"
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.qkv_proj": "colwise_gather_output",  # we need to replicate here due to the slicing of qkv
         "layers.*.self_attn.o_proj": "rowwise_split_input",  # input is replicated due to the slicing of qkv
         "layers.*.mlp.gate_up_proj": "colwise_gather_output",  # we need to replicate here due to the `chunk` operation
diff --git a/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py b/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py
index a6f47121057d..f4c3dbe4dd50 100644
--- a/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py
+++ b/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py
@@ -181,6 +181,7 @@ class Phi4MultimodalConfig(PreTrainedConfig):
     model_type = "phi4_multimodal"
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.qkv_proj": "colwise_gather_output",  # we need to replicate here due to the slicing of qkv
         "layers.*.self_attn.o_proj": "rowwise_split_input",  # input is replicated due to the slicing of qkv
         "layers.*.mlp.gate_up_proj": "colwise_gather_output",  # we need to replicate here due to the `chunk` operation
diff --git a/src/transformers/models/phimoe/configuration_phimoe.py b/src/transformers/models/phimoe/configuration_phimoe.py
index 89afec64ce3d..7d83152ce250 100644
--- a/src/transformers/models/phimoe/configuration_phimoe.py
+++ b/src/transformers/models/phimoe/configuration_phimoe.py
@@ -45,6 +45,7 @@ class PhimoeConfig(PreTrainedConfig):
     ```"""
 
     model_type = "phimoe"
+    base_model_tp_plan = {"embed_tokens": "embedding_rowwise"}
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_ep_plan = {
         "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
diff --git a/src/transformers/models/qwen2/configuration_qwen2.py b/src/transformers/models/qwen2/configuration_qwen2.py
index ae41af7b211f..0815d2ca0c64 100644
--- a/src/transformers/models/qwen2/configuration_qwen2.py
+++ b/src/transformers/models/qwen2/configuration_qwen2.py
@@ -44,6 +44,7 @@ class Qwen2Config(PreTrainedConfig):
 
     # Default tensor parallel plan for base model `Qwen2`
     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",
diff --git a/src/transformers/models/qwen2_5_omni/configuration_qwen2_5_omni.py b/src/transformers/models/qwen2_5_omni/configuration_qwen2_5_omni.py
index c4d597e7a44b..3481bbd99b2f 100644
--- a/src/transformers/models/qwen2_5_omni/configuration_qwen2_5_omni.py
+++ b/src/transformers/models/qwen2_5_omni/configuration_qwen2_5_omni.py
@@ -154,6 +154,7 @@ class Qwen2_5OmniTextConfig(PreTrainedConfig):
 
     # Default tensor parallel plan for base model `Qwen25OmniText`
     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",
diff --git a/src/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py b/src/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py
index 3dbdad07c981..6592c85fd5d1 100644
--- a/src/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py
+++ b/src/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py
@@ -265,6 +265,7 @@ class Qwen2_5OmniTextConfig(PreTrainedConfig):
 
     # Default tensor parallel plan for base model `Qwen25OmniText`
     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",
diff --git a/src/transformers/models/qwen2_5_vl/configuration_qwen2_5_vl.py b/src/transformers/models/qwen2_5_vl/configuration_qwen2_5_vl.py
index d73087ba5cce..edef654058eb 100644
--- a/src/transformers/models/qwen2_5_vl/configuration_qwen2_5_vl.py
+++ b/src/transformers/models/qwen2_5_vl/configuration_qwen2_5_vl.py
@@ -91,6 +91,7 @@ class Qwen2_5_VLTextConfig(PreTrainedConfig):
     default_theta = 1000000.0
     # Default tensor parallel plan for base model `Qwen2_5_VL`
     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",
diff --git a/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py b/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py
index 6f034ef7dde6..ef1c003266b7 100644
--- a/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py
+++ b/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py
@@ -54,6 +54,7 @@ class Qwen2MoeConfig(PreTrainedConfig):
 
     # Default tensor parallel plan for base model `Qwen2Moe`
     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",
diff --git a/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py b/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py
index 6a6fea8893e6..ac73eae70fbf 100644
--- a/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py
+++ b/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py
@@ -68,6 +68,7 @@ class Qwen2VLTextConfig(PreTrainedConfig):
     default_theta = 1000000.0
     # Default tensor parallel plan for base model `Qwen2VL`
     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",
diff --git a/src/transformers/models/qwen3/configuration_qwen3.py b/src/transformers/models/qwen3/configuration_qwen3.py
index 07ad0bb24b33..9f79fef44d18 100644
--- a/src/transformers/models/qwen3/configuration_qwen3.py
+++ b/src/transformers/models/qwen3/configuration_qwen3.py
@@ -43,6 +43,7 @@ class Qwen3Config(PreTrainedConfig):
 
     # Default tensor parallel plan for base model `Qwen3`
     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",
diff --git a/src/transformers/models/qwen3_5/configuration_qwen3_5.py b/src/transformers/models/qwen3_5/configuration_qwen3_5.py
index ee41a3d3b961..cbda226d8194 100644
--- a/src/transformers/models/qwen3_5/configuration_qwen3_5.py
+++ b/src/transformers/models/qwen3_5/configuration_qwen3_5.py
@@ -57,6 +57,7 @@ class Qwen3_5TextConfig(PreTrainedConfig):
     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",
diff --git a/src/transformers/models/qwen3_5/modular_qwen3_5.py b/src/transformers/models/qwen3_5/modular_qwen3_5.py
index 5202822a69fd..92ad3d17757f 100644
--- a/src/transformers/models/qwen3_5/modular_qwen3_5.py
+++ b/src/transformers/models/qwen3_5/modular_qwen3_5.py
@@ -101,6 +101,7 @@ class Qwen3_5TextConfig(Qwen3NextConfig):
     base_config_key = "text_config"
 
     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",
diff --git a/src/transformers/models/qwen3_5_moe/configuration_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/configuration_qwen3_5_moe.py
index 84739697ed85..0dfc76911558 100644
--- a/src/transformers/models/qwen3_5_moe/configuration_qwen3_5_moe.py
+++ b/src/transformers/models/qwen3_5_moe/configuration_qwen3_5_moe.py
@@ -57,6 +57,7 @@ class Qwen3_5MoeTextConfig(PreTrainedConfig):
     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",
diff --git a/src/transformers/models/qwen3_5_moe/modular_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/modular_qwen3_5_moe.py
index 3e9cab7916c1..7030ae93cb4f 100644
--- a/src/transformers/models/qwen3_5_moe/modular_qwen3_5_moe.py
+++ b/src/transformers/models/qwen3_5_moe/modular_qwen3_5_moe.py
@@ -93,6 +93,7 @@ class Qwen3_5MoeTextConfig(Qwen3NextConfig):
     base_config_key = "text_config"
 
     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",
diff --git a/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py b/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py
index 62cf801edfcf..f3136f021c84 100644
--- a/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py
+++ b/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py
@@ -54,6 +54,7 @@ class Qwen3MoeConfig(PreTrainedConfig):
 
     # Default tensor parallel plan for base model `Qwen3Moe`
     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",
diff --git a/src/transformers/models/qwen3_next/configuration_qwen3_next.py b/src/transformers/models/qwen3_next/configuration_qwen3_next.py
index bcba1a251bea..f418d8ff733e 100644
--- a/src/transformers/models/qwen3_next/configuration_qwen3_next.py
+++ b/src/transformers/models/qwen3_next/configuration_qwen3_next.py
@@ -59,6 +59,7 @@ class Qwen3NextConfig(PreTrainedConfig):
     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",
diff --git a/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py b/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py
index c358d4d913bb..5c020fe4483b 100644
--- a/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py
+++ b/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py
@@ -134,6 +134,7 @@ class Qwen3OmniMoeTextConfig(PreTrainedConfig):
 
     # Default tensor parallel plan for base model `Qwen3OmniMoeText`
     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",
@@ -246,6 +247,7 @@ class Qwen3OmniMoeTalkerCodePredictorConfig(PreTrainedConfig):
 
     # Default tensor parallel plan for base model `Qwen3OmniMoeTalkerCodePredictor`
     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",
@@ -335,6 +337,7 @@ class Qwen3OmniMoeTalkerTextConfig(PreTrainedConfig):
 
     # Default tensor parallel plan for base model `Qwen3OmniMoeTalkerText`
     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",
diff --git a/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py b/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py
index 3d29469afb0c..5dfc0fde4b45 100644
--- a/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py
+++ b/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py
@@ -285,6 +285,7 @@ class Qwen3OmniMoeTextConfig(PreTrainedConfig):
 
     # Default tensor parallel plan for base model `Qwen3OmniMoeText`
     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",
diff --git a/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py b/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py
index def956bb83a1..15e03e4c2878 100644
--- a/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py
+++ b/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py
@@ -56,6 +56,7 @@ class Qwen3VLMoeTextConfig(PreTrainedConfig):
     }
     # Default tensor parallel plan for base model `Qwen3VLMoe`
     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",
diff --git a/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py b/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py
index 0902fd903091..1f61a19dcb35 100644
--- a/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py
+++ b/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py
@@ -83,6 +83,7 @@ class Qwen3VLMoeTextConfig(Qwen3MoeConfig):
     default_theta = 500000.0
     # Default tensor parallel plan for base model `Qwen3VLMoe`
     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",
diff --git a/src/transformers/models/qwen4_exp/configuration_qwen4_exp.py b/src/transformers/models/qwen4_exp/configuration_qwen4_exp.py
index 9a7e85300baa..3cb9d97978a4 100644
--- a/src/transformers/models/qwen4_exp/configuration_qwen4_exp.py
+++ b/src/transformers/models/qwen4_exp/configuration_qwen4_exp.py
@@ -97,6 +97,7 @@ class Qwen4ExpTextConfig(PreTrainedConfig):
         # 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": "embedding_colwise",
+        "embed_tokens": "embedding_rowwise",
     }
     base_model_pp_plan = None
     base_model_ep_plan = {
diff --git a/src/transformers/models/qwen4_exp/modular_qwen4_exp.py b/src/transformers/models/qwen4_exp/modular_qwen4_exp.py
index e23ca547f566..7a6efa3558ae 100644
--- a/src/transformers/models/qwen4_exp/modular_qwen4_exp.py
+++ b/src/transformers/models/qwen4_exp/modular_qwen4_exp.py
@@ -142,6 +142,7 @@ class Qwen4ExpTextConfig(Qwen3_5MoeTextConfig):
         # 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": "embedding_colwise",
+        "embed_tokens": "embedding_rowwise",
     }
     base_model_fsdp_plan = {
         "embed_tokens": "free_full_weight",
diff --git a/src/transformers/models/seed_oss/configuration_seed_oss.py b/src/transformers/models/seed_oss/configuration_seed_oss.py
index b1221fcf53ce..e4d42a0ea5b8 100644
--- a/src/transformers/models/seed_oss/configuration_seed_oss.py
+++ b/src/transformers/models/seed_oss/configuration_seed_oss.py
@@ -45,6 +45,7 @@ class SeedOssConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `SeedOssModel`
     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",
diff --git a/src/transformers/models/smollm3/configuration_smollm3.py b/src/transformers/models/smollm3/configuration_smollm3.py
index f48c979a1dd3..9053ada91fce 100644
--- a/src/transformers/models/smollm3/configuration_smollm3.py
+++ b/src/transformers/models/smollm3/configuration_smollm3.py
@@ -55,6 +55,7 @@ class SmolLM3Config(PreTrainedConfig):
     default_theta = 2000000.0
 
     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",
diff --git a/src/transformers/models/smollm3/modular_smollm3.py b/src/transformers/models/smollm3/modular_smollm3.py
index f75017ad2645..fe7a04fddc9d 100644
--- a/src/transformers/models/smollm3/modular_smollm3.py
+++ b/src/transformers/models/smollm3/modular_smollm3.py
@@ -71,6 +71,7 @@ class SmolLM3Config(PreTrainedConfig):
     default_theta = 2000000.0
 
     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",
diff --git a/src/transformers/models/solar_open/configuration_solar_open.py b/src/transformers/models/solar_open/configuration_solar_open.py
index 5d7b88efc59b..ef49b7bbd3a8 100644
--- a/src/transformers/models/solar_open/configuration_solar_open.py
+++ b/src/transformers/models/solar_open/configuration_solar_open.py
@@ -38,6 +38,7 @@ class SolarOpenConfig(PreTrainedConfig):
 
     # Default tensor parallel plan for base model `SolarOpenModel`
     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",
diff --git a/src/transformers/models/solar_open/modular_solar_open.py b/src/transformers/models/solar_open/modular_solar_open.py
index 74236a466411..6a50edbac46d 100644
--- a/src/transformers/models/solar_open/modular_solar_open.py
+++ b/src/transformers/models/solar_open/modular_solar_open.py
@@ -44,6 +44,7 @@ class SolarOpenConfig(Glm4MoeConfig):
 
     # Default tensor parallel plan for base model `SolarOpenModel`
     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",
diff --git a/src/transformers/models/starcoder2/configuration_starcoder2.py b/src/transformers/models/starcoder2/configuration_starcoder2.py
index 59efa94fc5f4..19fef828c593 100644
--- a/src/transformers/models/starcoder2/configuration_starcoder2.py
+++ b/src/transformers/models/starcoder2/configuration_starcoder2.py
@@ -45,6 +45,7 @@ class Starcoder2Config(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `Starcoder2`
     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",
diff --git a/src/transformers/models/step3p7/configuration_step3p7.py b/src/transformers/models/step3p7/configuration_step3p7.py
index 56c747be1dfa..89728fc16178 100644
--- a/src/transformers/models/step3p7/configuration_step3p7.py
+++ b/src/transformers/models/step3p7/configuration_step3p7.py
@@ -113,6 +113,7 @@ class Step3p7TextConfig(PreTrainedConfig):
     # gates their gathered output). Spelled out in full, not `{**MiniMaxM3VLTextConfig.base_model_tp_plan, ...}`:
     # generated files have no cross-model imports, so that name wouldn't resolve there.
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_proj": "colwise_gather_output",
         "layers.*.self_attn.k_proj": "colwise_gather_output",
         "layers.*.self_attn.v_proj": "colwise_gather_output",
diff --git a/src/transformers/models/step3p7/modular_step3p7.py b/src/transformers/models/step3p7/modular_step3p7.py
index 384f2cbfd5cf..fb659c197f01 100644
--- a/src/transformers/models/step3p7/modular_step3p7.py
+++ b/src/transformers/models/step3p7/modular_step3p7.py
@@ -166,6 +166,7 @@ class Step3p7TextConfig(MiniMaxM3VLTextConfig):
     # gates their gathered output). Spelled out in full, not `{**MiniMaxM3VLTextConfig.base_model_tp_plan, ...}`:
     # generated files have no cross-model imports, so that name wouldn't resolve there.
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_proj": "colwise_gather_output",
         "layers.*.self_attn.k_proj": "colwise_gather_output",
         "layers.*.self_attn.v_proj": "colwise_gather_output",
diff --git a/src/transformers/models/t5gemma/configuration_t5gemma.py b/src/transformers/models/t5gemma/configuration_t5gemma.py
index 9509b7df9d09..a2051edaf06d 100644
--- a/src/transformers/models/t5gemma/configuration_t5gemma.py
+++ b/src/transformers/models/t5gemma/configuration_t5gemma.py
@@ -51,6 +51,7 @@ class T5GemmaModuleConfig(PreTrainedConfig):
     model_type = "t5_gemma_module"
     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",
diff --git a/src/transformers/models/t5gemma2/configuration_t5gemma2.py b/src/transformers/models/t5gemma2/configuration_t5gemma2.py
index 65bcc1a43ac4..692f87e36ac6 100644
--- a/src/transformers/models/t5gemma2/configuration_t5gemma2.py
+++ b/src/transformers/models/t5gemma2/configuration_t5gemma2.py
@@ -42,6 +42,7 @@ class T5Gemma2TextConfig(PreTrainedConfig):
     model_type = "t5gemma2_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",
@@ -210,6 +211,7 @@ class T5Gemma2DecoderConfig(PreTrainedConfig):
     model_type = "t5gemma2_decoder"
     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",
diff --git a/src/transformers/models/vaultgemma/configuration_vaultgemma.py b/src/transformers/models/vaultgemma/configuration_vaultgemma.py
index a60b7e8edc0c..0cc086235448 100644
--- a/src/transformers/models/vaultgemma/configuration_vaultgemma.py
+++ b/src/transformers/models/vaultgemma/configuration_vaultgemma.py
@@ -50,6 +50,7 @@ class VaultGemmaConfig(PreTrainedConfig):
     model_type = "vaultgemma"
     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",
diff --git a/src/transformers/models/voxtral_realtime/configuration_voxtral_realtime.py b/src/transformers/models/voxtral_realtime/configuration_voxtral_realtime.py
index 3ba420f94ade..82ff27fb55e6 100644
--- a/src/transformers/models/voxtral_realtime/configuration_voxtral_realtime.py
+++ b/src/transformers/models/voxtral_realtime/configuration_voxtral_realtime.py
@@ -27,6 +27,7 @@ class VoxtralRealtimeTextConfig(PreTrainedConfig):
     model_type = "voxtral_realtime_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",
diff --git a/src/transformers/models/youtu/configuration_youtu.py b/src/transformers/models/youtu/configuration_youtu.py
index 6d9f2cef1f96..75c9989333fe 100644
--- a/src/transformers/models/youtu/configuration_youtu.py
+++ b/src/transformers/models/youtu/configuration_youtu.py
@@ -51,6 +51,7 @@ class YoutuConfig(PreTrainedConfig):
     model_type = "youtu"
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.mlp.gate_proj": "colwise",
         "layers.*.mlp.up_proj": "colwise",
         "layers.*.mlp.down_proj": "rowwise",
diff --git a/src/transformers/models/youtu/modular_youtu.py b/src/transformers/models/youtu/modular_youtu.py
index c5ba1f03da53..0f33d53ee937 100644
--- a/src/transformers/models/youtu/modular_youtu.py
+++ b/src/transformers/models/youtu/modular_youtu.py
@@ -60,6 +60,7 @@ class YoutuConfig(DeepseekV3Config):
 
     model_type = "youtu"
     base_model_tp_plan = {
+        "embed_tokens": "embedding_rowwise",
         "layers.*.mlp.gate_proj": "colwise",
         "layers.*.mlp.up_proj": "colwise",
         "layers.*.mlp.down_proj": "rowwise",
diff --git a/tests/models/dbrx/test_modeling_dbrx.py b/tests/models/dbrx/test_modeling_dbrx.py
index 63e10bb41388..d2458c67af53 100644
--- a/tests/models/dbrx/test_modeling_dbrx.py
+++ b/tests/models/dbrx/test_modeling_dbrx.py
@@ -95,6 +95,10 @@ def test_model_from_pretrained(self):
 
     # Offload does not work with Dbrx models because of the forward of DbrxExperts where we chunk the experts.
     # The issue is that the offloaded weights of the mlp layer are still on meta device (w1_chunked, v1_chunked, w2_chunked)
+    @unittest.skip(reason="DBRX doesn't support tied word embeddings")
+    def test_tp_forward_1(self):
+        pass
+
     @unittest.skip(reason="Dbrx models do not work with offload")
     def test_cpu_offload(self):
         pass
diff --git a/tests/models/phimoe/test_modeling_phimoe.py b/tests/models/phimoe/test_modeling_phimoe.py
index 7598d74c4dc7..0be32266462d 100644
--- a/tests/models/phimoe/test_modeling_phimoe.py
+++ b/tests/models/phimoe/test_modeling_phimoe.py
@@ -101,6 +101,10 @@ def is_pipeline_test_to_skip(
     ):
         return True
 
+    @unittest.skip("PhiMoE routes by Gumbel sampling in training, so TP and reference backward pick different experts")
+    def test_tp_backward(self):
+        pass
+
     @unittest.skip("PhiMoE's RoPE has custom parameterization")
     def test_model_rope_scaling_frequencies(self):
         pass
diff --git a/tests/test_tensor_parallel_mixin.py b/tests/test_tensor_parallel_mixin.py
index d385b931b589..af559bfcea5c 100644
--- a/tests/test_tensor_parallel_mixin.py
+++ b/tests/test_tensor_parallel_mixin.py
@@ -10,6 +10,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import copy
 import os
 import socket
 import tempfile
@@ -581,6 +582,23 @@ def _skip_if_not_supported(self, expert_parallel: bool = False):
         # if hasattr(config, "vision_config") and config.vision_config is not None:
         #     self.skipTest("VLM models are not yet supported in TP tests")
 
+    def test_tp_plan_shards_input_embeddings(self):
+        """A model with a TP plan must shard its input embeddings: replicated, they cost `vocab x hidden` on every
+        rank, and a tied `lm_head` sharded on the vocabulary would not match the embedding weight it shares."""
+        config, _ = self.model_tester.prepare_config_and_inputs_for_common()
+        for model_class in self.all_model_classes:
+            model = model_class(copy.deepcopy(config))
+            if not model.tp_plan:
+                continue
+            embeddings = model.get_input_embeddings()
+            if not isinstance(embeddings, torch.nn.Embedding):
+                continue
+            name = next(name for name, module in model.named_modules() if module is embeddings)
+            self.assertIsNotNone(
+                _get_parameter_plan(f"{name}.weight", model.tp_plan, is_weight=True),
+                f"{model_class.__name__}: the TP plan leaves `{name}` replicated; shard it with `embedding_rowwise`",
+            )
+
     @parameterized.expand([(False,), (True,)])
     @is_tensor_parallel_test
     def test_tp_forward(self, tie_word_embeddings):

From 11051a6d91c016735f9f9edd269e12c4dbdeb498 Mon Sep 17 00:00:00 2001
From: IlyasMoutawwakil <moutawwakil.ilyas.tsi@gmail.com>
Date: Wed, 14 Oct 2026 11:15:01 +0200
Subject: [PATCH 08/13] TP: shard `embed_tokens` by default instead of in every
 plan

Drop the `tie_word_embeddings` gate on the config-level `embedding_rowwise` default (explicit plan entries win), remove the per-config `embed_tokens` entries, and leave the lm_head-only models out of the embedding rollout; `test_tp_plan_shards_input_embeddings` now only checks decoder plans.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---
 src/transformers/configuration_utils.py                    | 7 ++-----
 src/transformers/models/afmoe/configuration_afmoe.py       | 1 -
 src/transformers/models/apertus/configuration_apertus.py   | 1 -
 src/transformers/models/apertus/modular_apertus.py         | 1 -
 src/transformers/models/arcee/configuration_arcee.py       | 1 -
 src/transformers/models/arcee/modular_arcee.py             | 1 -
 src/transformers/models/aria/configuration_aria.py         | 1 -
 src/transformers/models/aria/modular_aria.py               | 1 -
 src/transformers/models/axk1/configuration_axk1.py         | 1 -
 src/transformers/models/axk2/configuration_axk2.py         | 1 -
 src/transformers/models/axk2/modular_axk2.py               | 1 -
 src/transformers/models/cohere/configuration_cohere.py     | 1 -
 src/transformers/models/cohere2/configuration_cohere2.py   | 1 -
 src/transformers/models/cohere2/modular_cohere2.py         | 1 -
 .../models/cohere2_moe/configuration_cohere2_moe.py        | 1 -
 .../models/cohere_compass/configuration_cohere_compass.py  | 1 -
 .../models/cosmos3_edge/configuration_cosmos3_edge.py      | 1 -
 .../models/cosmos3_edge/modular_cosmos3_edge.py            | 1 -
 src/transformers/models/cwm/configuration_cwm.py           | 1 -
 src/transformers/models/dbrx/configuration_dbrx.py         | 1 -
 .../models/deepseek_ocr2/configuration_deepseek_ocr2.py    | 2 --
 .../models/deepseek_ocr2/modular_deepseek_ocr2.py          | 1 -
 .../models/deepseek_v2/configuration_deepseek_v2.py        | 1 -
 src/transformers/models/deepseek_v2/modular_deepseek_v2.py | 1 -
 .../models/deepseek_v3/configuration_deepseek_v3.py        | 1 -
 .../models/deepseek_v32/configuration_deepseek_v32.py      | 1 -
 .../models/deepseek_v32/modular_deepseek_v32.py            | 1 -
 .../diffusion_gemma/configuration_diffusion_gemma.py       | 1 -
 src/transformers/models/dots1/configuration_dots1.py       | 1 -
 src/transformers/models/dots1/modular_dots1.py             | 1 -
 .../embedding_gemma2/configuration_embedding_gemma2.py     | 1 -
 .../models/embedding_gemma2/modular_embedding_gemma2.py    | 1 -
 src/transformers/models/ernie4_5/configuration_ernie4_5.py | 1 -
 .../models/ernie4_5_moe/configuration_ernie4_5_moe.py      | 1 -
 .../ernie4_5_vl_moe/configuration_ernie4_5_vl_moe.py       | 1 -
 .../models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py      | 1 -
 src/transformers/models/esmc/configuration_esmc.py         | 1 -
 src/transformers/models/eurobert/configuration_eurobert.py | 1 -
 src/transformers/models/exaone4/configuration_exaone4.py   | 1 -
 src/transformers/models/exaone4/modular_exaone4.py         | 1 -
 .../models/exaone_moe/configuration_exaone_moe.py          | 1 -
 .../models/flex_olmo/configuration_flex_olmo.py            | 1 -
 src/transformers/models/flex_olmo/modular_flex_olmo.py     | 1 -
 src/transformers/models/gemma/configuration_gemma.py       | 1 -
 src/transformers/models/gemma/modular_gemma.py             | 1 -
 src/transformers/models/gemma2/configuration_gemma2.py     | 1 -
 src/transformers/models/gemma2/modular_gemma2.py           | 1 -
 src/transformers/models/gemma3/configuration_gemma3.py     | 1 -
 src/transformers/models/gemma3/modular_gemma3.py           | 1 -
 src/transformers/models/gemma3n/configuration_gemma3n.py   | 1 -
 src/transformers/models/gemma3n/modular_gemma3n.py         | 1 -
 src/transformers/models/gemma4/configuration_gemma4.py     | 1 -
 .../models/gemma4_unified/configuration_gemma4_unified.py  | 1 -
 .../models/gemma4_unified/modular_gemma4_unified.py        | 1 -
 src/transformers/models/glm/configuration_glm.py           | 1 -
 src/transformers/models/glm4/configuration_glm4.py         | 1 -
 src/transformers/models/glm4_moe/configuration_glm4_moe.py | 1 -
 src/transformers/models/glm4_moe/modular_glm4_moe.py       | 1 -
 .../models/glm4_moe_lite/configuration_glm4_moe_lite.py    | 1 -
 .../models/glm4_moe_lite/modular_glm4_moe_lite.py          | 1 -
 src/transformers/models/glm4v/configuration_glm4v.py       | 1 -
 src/transformers/models/glm4v/modular_glm4v.py             | 1 -
 .../models/glm4v_moe/configuration_glm4v_moe.py            | 1 -
 src/transformers/models/glm4v_moe/modular_glm4v_moe.py     | 1 -
 .../models/glm5_next/configuration_glm5_next.py            | 1 -
 .../models/glm_image/configuration_glm_image.py            | 1 -
 .../models/glm_moe_dsa/configuration_glm_moe_dsa.py        | 1 -
 src/transformers/models/glm_ocr/configuration_glm_ocr.py   | 1 -
 src/transformers/models/gpt_oss/configuration_gpt_oss.py   | 1 -
 src/transformers/models/granite/configuration_granite.py   | 1 -
 .../granite4_vision/configuration_granite4_vision.py       | 1 -
 .../models/granite_swa/configuration_granite_swa.py        | 1 -
 src/transformers/models/granite_swa/modular_granite_swa.py | 1 -
 .../models/granitemoe_swa/configuration_granitemoe_swa.py  | 1 -
 .../models/granitemoe_swa/modular_granitemoe_swa.py        | 1 -
 src/transformers/models/helium/configuration_helium.py     | 1 -
 .../models/higgs_audio_v2/configuration_higgs_audio_v2.py  | 1 -
 src/transformers/models/hrm_text/configuration_hrm_text.py | 1 -
 src/transformers/models/hrm_text/modular_hrm_text.py       | 1 -
 .../hunyuan_v1_dense/configuration_hunyuan_v1_dense.py     | 1 -
 .../models/hunyuan_v1_moe/configuration_hunyuan_v1_moe.py  | 1 -
 .../models/hunyuan_vl/configuration_hunyuan_vl.py          | 1 -
 src/transformers/models/hy_v3/configuration_hy_v3.py       | 1 -
 src/transformers/models/hy_v3/modular_hy_v3.py             | 1 -
 src/transformers/models/hy_v4/configuration_hy_v4.py       | 1 -
 src/transformers/models/hy_v4/modular_hy_v4.py             | 1 -
 .../models/hyperclovax/configuration_hyperclovax.py        | 1 -
 src/transformers/models/inkling/configuration_inkling.py   | 1 -
 src/transformers/models/inkling/modular_inkling.py         | 1 -
 src/transformers/models/jais2/configuration_jais2.py       | 1 -
 src/transformers/models/jais2/modular_jais2.py             | 1 -
 .../models/kimi_linear/configuration_kimi_linear.py        | 1 -
 src/transformers/models/laguna/configuration_laguna.py     | 1 -
 src/transformers/models/laguna/modular_laguna.py           | 1 -
 src/transformers/models/lfm2/configuration_lfm2.py         | 1 -
 src/transformers/models/lfm2_moe/configuration_lfm2_moe.py | 1 -
 src/transformers/models/llama/configuration_llama.py       | 1 -
 .../models/longcat_flash/configuration_longcat_flash.py    | 1 -
 src/transformers/models/mellum/configuration_mellum.py     | 1 -
 .../models/mimo_v2_flash/configuration_mimo_v2_flash.py    | 1 -
 .../models/mimo_v2_flash/modular_mimo_v2_flash.py          | 1 -
 src/transformers/models/minicpm3/configuration_minicpm3.py | 1 -
 src/transformers/models/minicpm3/modular_minicpm3.py       | 1 -
 src/transformers/models/minimax/configuration_minimax.py   | 1 -
 src/transformers/models/minimax/modular_minimax.py         | 1 -
 .../models/minimax_m2/configuration_minimax_m2.py          | 1 -
 src/transformers/models/minimax_m2/modular_minimax_m2.py   | 1 -
 .../models/minimax_m3_vl/configuration_minimax_m3_vl.py    | 1 -
 .../models/ministral/configuration_ministral.py            | 1 -
 .../models/ministral3/configuration_ministral3.py          | 1 -
 src/transformers/models/mistral/configuration_mistral.py   | 1 -
 src/transformers/models/mistral4/configuration_mistral4.py | 1 -
 src/transformers/models/mixtral/configuration_mixtral.py   | 1 -
 .../models/muse_glimmer/configuration_muse_glimmer.py      | 1 -
 .../models/muse_glimmer/modular_muse_glimmer.py            | 1 -
 src/transformers/models/nanochat/configuration_nanochat.py | 1 -
 src/transformers/models/neomme/configuration_neomme.py     | 1 -
 src/transformers/models/olmo/configuration_olmo.py         | 1 -
 src/transformers/models/olmo2/configuration_olmo2.py       | 1 -
 src/transformers/models/olmo2/modular_olmo2.py             | 1 -
 src/transformers/models/olmo3/configuration_olmo3.py       | 1 -
 src/transformers/models/olmo3/modular_olmo3.py             | 1 -
 .../models/olmo_hybrid/configuration_olmo_hybrid.py        | 1 -
 src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py | 1 -
 src/transformers/models/olmoe/configuration_olmoe.py       | 1 -
 .../configuration_openai_privacy_filter.py                 | 1 -
 .../models/paddleocr_vl/configuration_paddleocr_vl.py      | 1 -
 src/transformers/models/phi/configuration_phi.py           | 1 -
 src/transformers/models/phi3/configuration_phi3.py         | 1 -
 .../phi4_multimodal/configuration_phi4_multimodal.py       | 1 -
 src/transformers/models/phimoe/configuration_phimoe.py     | 1 -
 src/transformers/models/qwen2/configuration_qwen2.py       | 1 -
 .../models/qwen2_5_omni/configuration_qwen2_5_omni.py      | 1 -
 .../models/qwen2_5_omni/modular_qwen2_5_omni.py            | 1 -
 .../models/qwen2_5_vl/configuration_qwen2_5_vl.py          | 1 -
 .../models/qwen2_moe/configuration_qwen2_moe.py            | 1 -
 src/transformers/models/qwen2_vl/configuration_qwen2_vl.py | 1 -
 src/transformers/models/qwen3/configuration_qwen3.py       | 1 -
 src/transformers/models/qwen3_5/configuration_qwen3_5.py   | 1 -
 src/transformers/models/qwen3_5/modular_qwen3_5.py         | 1 -
 .../models/qwen3_5_moe/configuration_qwen3_5_moe.py        | 1 -
 src/transformers/models/qwen3_5_moe/modular_qwen3_5_moe.py | 1 -
 .../models/qwen3_moe/configuration_qwen3_moe.py            | 1 -
 .../models/qwen3_next/configuration_qwen3_next.py          | 1 -
 .../models/qwen3_omni_moe/configuration_qwen3_omni_moe.py  | 3 ---
 .../models/qwen3_omni_moe/modular_qwen3_omni_moe.py        | 1 -
 .../models/qwen3_vl_moe/configuration_qwen3_vl_moe.py      | 1 -
 .../models/qwen3_vl_moe/modular_qwen3_vl_moe.py            | 1 -
 .../models/qwen4_exp/configuration_qwen4_exp.py            | 1 -
 src/transformers/models/qwen4_exp/modular_qwen4_exp.py     | 1 -
 src/transformers/models/seed_oss/configuration_seed_oss.py | 1 -
 src/transformers/models/smollm3/configuration_smollm3.py   | 1 -
 src/transformers/models/smollm3/modular_smollm3.py         | 1 -
 .../models/solar_open/configuration_solar_open.py          | 1 -
 src/transformers/models/solar_open/modular_solar_open.py   | 1 -
 .../models/starcoder2/configuration_starcoder2.py          | 1 -
 src/transformers/models/step3p7/configuration_step3p7.py   | 1 -
 src/transformers/models/step3p7/modular_step3p7.py         | 1 -
 src/transformers/models/t5gemma/configuration_t5gemma.py   | 1 -
 src/transformers/models/t5gemma2/configuration_t5gemma2.py | 2 --
 .../models/vaultgemma/configuration_vaultgemma.py          | 1 -
 .../voxtral_realtime/configuration_voxtral_realtime.py     | 1 -
 src/transformers/models/youtu/configuration_youtu.py       | 1 -
 src/transformers/models/youtu/modular_youtu.py             | 1 -
 tests/models/dbrx/test_modeling_dbrx.py                    | 4 ----
 tests/models/phimoe/test_modeling_phimoe.py                | 4 ----
 tests/test_tensor_parallel_mixin.py                        | 7 ++++---
 167 files changed, 6 insertions(+), 183 deletions(-)

diff --git a/src/transformers/configuration_utils.py b/src/transformers/configuration_utils.py
index 3f32558c4ad6..786bbbf4a2b9 100755
--- a/src/transformers/configuration_utils.py
+++ b/src/transformers/configuration_utils.py
@@ -507,11 +507,8 @@ def __post_init__(self, **kwargs):
             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",
-            }
+        if self.base_model_tp_plan is not None:
+            self.base_model_tp_plan = {"embed_tokens": "embedding_rowwise", **self.base_model_tp_plan}
 
         # Remap layer types if needed
         remap_legacy_layer_types(config=self)
diff --git a/src/transformers/models/afmoe/configuration_afmoe.py b/src/transformers/models/afmoe/configuration_afmoe.py
index 7f6d0293894c..40083d09e68a 100644
--- a/src/transformers/models/afmoe/configuration_afmoe.py
+++ b/src/transformers/models/afmoe/configuration_afmoe.py
@@ -53,7 +53,6 @@ class AfmoeConfig(PreTrainedConfig):
     """
 
     model_type = "afmoe"
-    base_model_tp_plan = {"embed_tokens": "embedding_rowwise"}
     keys_to_ignore_at_inference = ["past_key_values"]
 
     # Default pipeline parallel plan for base model
diff --git a/src/transformers/models/apertus/configuration_apertus.py b/src/transformers/models/apertus/configuration_apertus.py
index 0894ca731ea5..1e0122160b30 100644
--- a/src/transformers/models/apertus/configuration_apertus.py
+++ b/src/transformers/models/apertus/configuration_apertus.py
@@ -46,7 +46,6 @@ class ApertusConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     default_theta = 12000000.0
     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",
diff --git a/src/transformers/models/apertus/modular_apertus.py b/src/transformers/models/apertus/modular_apertus.py
index f771fa6b5510..03ae896f8b07 100644
--- a/src/transformers/models/apertus/modular_apertus.py
+++ b/src/transformers/models/apertus/modular_apertus.py
@@ -64,7 +64,6 @@ class ApertusConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     default_theta = 12000000.0
     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",
diff --git a/src/transformers/models/arcee/configuration_arcee.py b/src/transformers/models/arcee/configuration_arcee.py
index 587612e0293b..ee711d608204 100644
--- a/src/transformers/models/arcee/configuration_arcee.py
+++ b/src/transformers/models/arcee/configuration_arcee.py
@@ -46,7 +46,6 @@ class ArceeConfig(PreTrainedConfig):
     model_type = "arcee"
     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",
diff --git a/src/transformers/models/arcee/modular_arcee.py b/src/transformers/models/arcee/modular_arcee.py
index f3e153908722..91cd8e13f1ed 100644
--- a/src/transformers/models/arcee/modular_arcee.py
+++ b/src/transformers/models/arcee/modular_arcee.py
@@ -50,7 +50,6 @@ class ArceeConfig(LlamaConfig):
 
     model_type = "arcee"
     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",
diff --git a/src/transformers/models/aria/configuration_aria.py b/src/transformers/models/aria/configuration_aria.py
index 63549636af58..6da1884bf8cb 100644
--- a/src/transformers/models/aria/configuration_aria.py
+++ b/src/transformers/models/aria/configuration_aria.py
@@ -41,7 +41,6 @@ class AriaTextConfig(PreTrainedConfig):
     model_type = "aria_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",
diff --git a/src/transformers/models/aria/modular_aria.py b/src/transformers/models/aria/modular_aria.py
index 2ae9aa68136c..6e1484f6af03 100644
--- a/src/transformers/models/aria/modular_aria.py
+++ b/src/transformers/models/aria/modular_aria.py
@@ -81,7 +81,6 @@ class AriaTextConfig(LlamaConfig):
     model_type = "aria_text"
     base_config_key = "text_config"
     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",
diff --git a/src/transformers/models/axk1/configuration_axk1.py b/src/transformers/models/axk1/configuration_axk1.py
index db38fe7bd0ba..134ee88b01a6 100644
--- a/src/transformers/models/axk1/configuration_axk1.py
+++ b/src/transformers/models/axk1/configuration_axk1.py
@@ -53,7 +53,6 @@ class AXK1Config(PreTrainedConfig):
     model_type = "axk1"
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.mlp.experts.gate_up_proj": "packed_colwise",
         "layers.*.mlp.experts.down_proj": "rowwise",
         "layers.*.mlp.experts": "moe_tp_experts",
diff --git a/src/transformers/models/axk2/configuration_axk2.py b/src/transformers/models/axk2/configuration_axk2.py
index 4343ab85ba24..a5f57ff65048 100644
--- a/src/transformers/models/axk2/configuration_axk2.py
+++ b/src/transformers/models/axk2/configuration_axk2.py
@@ -62,7 +62,6 @@ class AXK2Config(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
 
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_gate_proj": "colwise",
         "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
         "layers.*.self_attn.kv_b_proj": "colwise",
diff --git a/src/transformers/models/axk2/modular_axk2.py b/src/transformers/models/axk2/modular_axk2.py
index 125220f402b6..d18dd84f4d7d 100644
--- a/src/transformers/models/axk2/modular_axk2.py
+++ b/src/transformers/models/axk2/modular_axk2.py
@@ -123,7 +123,6 @@ class AXK2Config(DeepseekV32Config):
     mlp_bias = AttributeError()
 
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_gate_proj": "colwise",
         "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
         "layers.*.self_attn.kv_b_proj": "colwise",
diff --git a/src/transformers/models/cohere/configuration_cohere.py b/src/transformers/models/cohere/configuration_cohere.py
index 2909d12972ec..44dbb4c01c20 100644
--- a/src/transformers/models/cohere/configuration_cohere.py
+++ b/src/transformers/models/cohere/configuration_cohere.py
@@ -50,7 +50,6 @@ class CohereConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     default_theta = 500000.0
     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",
diff --git a/src/transformers/models/cohere2/configuration_cohere2.py b/src/transformers/models/cohere2/configuration_cohere2.py
index 4a6b1c750580..48c2df360354 100644
--- a/src/transformers/models/cohere2/configuration_cohere2.py
+++ b/src/transformers/models/cohere2/configuration_cohere2.py
@@ -49,7 +49,6 @@ class Cohere2Config(PreTrainedConfig):
     model_type = "cohere2"
     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",
diff --git a/src/transformers/models/cohere2/modular_cohere2.py b/src/transformers/models/cohere2/modular_cohere2.py
index 0c5f9dd6ac0d..9fd81c926971 100644
--- a/src/transformers/models/cohere2/modular_cohere2.py
+++ b/src/transformers/models/cohere2/modular_cohere2.py
@@ -69,7 +69,6 @@ class Cohere2Config(PreTrainedConfig):
     model_type = "cohere2"
     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",
diff --git a/src/transformers/models/cohere2_moe/configuration_cohere2_moe.py b/src/transformers/models/cohere2_moe/configuration_cohere2_moe.py
index 0ed5318195a4..e94bed7ac32c 100644
--- a/src/transformers/models/cohere2_moe/configuration_cohere2_moe.py
+++ b/src/transformers/models/cohere2_moe/configuration_cohere2_moe.py
@@ -71,7 +71,6 @@ class Cohere2MoeConfig(PreTrainedConfig):
     model_type = "cohere2_moe"
     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",
diff --git a/src/transformers/models/cohere_compass/configuration_cohere_compass.py b/src/transformers/models/cohere_compass/configuration_cohere_compass.py
index 68bd2d7ed0b9..74459f7a8ea2 100644
--- a/src/transformers/models/cohere_compass/configuration_cohere_compass.py
+++ b/src/transformers/models/cohere_compass/configuration_cohere_compass.py
@@ -70,7 +70,6 @@ class CohereCompassTextConfig(PreTrainedConfig):
     model_type = "cohere_compass_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",
diff --git a/src/transformers/models/cosmos3_edge/configuration_cosmos3_edge.py b/src/transformers/models/cosmos3_edge/configuration_cosmos3_edge.py
index 838775bc868c..76790b3d135b 100644
--- a/src/transformers/models/cosmos3_edge/configuration_cosmos3_edge.py
+++ b/src/transformers/models/cosmos3_edge/configuration_cosmos3_edge.py
@@ -43,7 +43,6 @@ class Cosmos3EdgeTextConfig(PreTrainedConfig):
     model_type = "cosmos3_edge_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",
diff --git a/src/transformers/models/cosmos3_edge/modular_cosmos3_edge.py b/src/transformers/models/cosmos3_edge/modular_cosmos3_edge.py
index 5172b29876aa..def47bc67e31 100644
--- a/src/transformers/models/cosmos3_edge/modular_cosmos3_edge.py
+++ b/src/transformers/models/cosmos3_edge/modular_cosmos3_edge.py
@@ -89,7 +89,6 @@ class Cosmos3EdgeTextConfig(LlamaConfig):
     default_theta = 100_000_000.0
     ignore_keys_at_rope_validation = {"mrope_section"}
     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",
diff --git a/src/transformers/models/cwm/configuration_cwm.py b/src/transformers/models/cwm/configuration_cwm.py
index 6a8c6c563ada..ecc3743da19d 100644
--- a/src/transformers/models/cwm/configuration_cwm.py
+++ b/src/transformers/models/cwm/configuration_cwm.py
@@ -46,7 +46,6 @@ class CwmConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `CwmModel`
     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",
diff --git a/src/transformers/models/dbrx/configuration_dbrx.py b/src/transformers/models/dbrx/configuration_dbrx.py
index 41ea4fdb234b..17c57149d545 100644
--- a/src/transformers/models/dbrx/configuration_dbrx.py
+++ b/src/transformers/models/dbrx/configuration_dbrx.py
@@ -111,7 +111,6 @@ class DbrxConfig(PreTrainedConfig):
     """
 
     model_type = "dbrx"
-    base_model_tp_plan = {"wte": "embedding_rowwise"}
     sub_configs = {"attn_config": DbrxAttentionConfig, "ffn_config": DbrxFFNConfig}
     attribute_map = {
         "num_attention_heads": "n_heads",
diff --git a/src/transformers/models/deepseek_ocr2/configuration_deepseek_ocr2.py b/src/transformers/models/deepseek_ocr2/configuration_deepseek_ocr2.py
index 02a9c5a6509d..7adb7745d895 100644
--- a/src/transformers/models/deepseek_ocr2/configuration_deepseek_ocr2.py
+++ b/src/transformers/models/deepseek_ocr2/configuration_deepseek_ocr2.py
@@ -91,7 +91,6 @@ class DeepseekOcr2VisionEncoderConfig(PreTrainedConfig):
 
     # Default tensor parallel plan for base model `DeepseekOcr2VisionEncoder`
     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",
@@ -183,7 +182,6 @@ class DeepseekOcr2TextConfig(PreTrainedConfig):
 
     # Override DeepseekV2's MLA TP plan with standard MHA projections
     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",
diff --git a/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py b/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py
index 03ac37bdda64..f06bb355700e 100644
--- a/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py
+++ b/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py
@@ -571,7 +571,6 @@ class DeepseekOcr2TextConfig(DeepseekV2Config):
 
     # Override DeepseekV2's MLA TP plan with standard MHA projections
     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",
diff --git a/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py b/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py
index f12ea083878b..0b9093626489 100644
--- a/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py
+++ b/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py
@@ -52,7 +52,6 @@ class DeepseekV2Config(PreTrainedConfig):
     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.q_b_proj": "colwise",
         "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
diff --git a/src/transformers/models/deepseek_v2/modular_deepseek_v2.py b/src/transformers/models/deepseek_v2/modular_deepseek_v2.py
index b1381bf921ab..f31d8d4528ca 100644
--- a/src/transformers/models/deepseek_v2/modular_deepseek_v2.py
+++ b/src/transformers/models/deepseek_v2/modular_deepseek_v2.py
@@ -70,7 +70,6 @@ class DeepseekV2Config(LlamaConfig):
     """
 
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_proj": "colwise",
         "layers.*.self_attn.q_b_proj": "colwise",
         "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
diff --git a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py
index 25c71936f5bb..a030c1db5d03 100644
--- a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py
+++ b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py
@@ -53,7 +53,6 @@ class DeepseekV3Config(PreTrainedConfig):
     model_type = "deepseek_v3"
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.mlp.experts.gate_up_proj": "packed_colwise",
         "layers.*.mlp.experts.down_proj": "rowwise",
         "layers.*.mlp.experts": "moe_tp_experts",
diff --git a/src/transformers/models/deepseek_v32/configuration_deepseek_v32.py b/src/transformers/models/deepseek_v32/configuration_deepseek_v32.py
index a01ca4607fcf..48eb2047edcb 100644
--- a/src/transformers/models/deepseek_v32/configuration_deepseek_v32.py
+++ b/src/transformers/models/deepseek_v32/configuration_deepseek_v32.py
@@ -57,7 +57,6 @@ class DeepseekV32Config(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
 
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_b_proj": "colwise",
         "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
         "layers.*.self_attn.kv_b_proj": "colwise",
diff --git a/src/transformers/models/deepseek_v32/modular_deepseek_v32.py b/src/transformers/models/deepseek_v32/modular_deepseek_v32.py
index bb829b5342d1..45e979579c8b 100644
--- a/src/transformers/models/deepseek_v32/modular_deepseek_v32.py
+++ b/src/transformers/models/deepseek_v32/modular_deepseek_v32.py
@@ -85,7 +85,6 @@ class DeepseekV32Config(Glm4MoeLiteConfig):
     ```"""
 
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_b_proj": "colwise",
         "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
         "layers.*.self_attn.kv_b_proj": "colwise",
diff --git a/src/transformers/models/diffusion_gemma/configuration_diffusion_gemma.py b/src/transformers/models/diffusion_gemma/configuration_diffusion_gemma.py
index 78a79c32b83d..64460b97be3b 100644
--- a/src/transformers/models/diffusion_gemma/configuration_diffusion_gemma.py
+++ b/src/transformers/models/diffusion_gemma/configuration_diffusion_gemma.py
@@ -47,7 +47,6 @@ class DiffusionGemmaTextConfig(PreTrainedConfig):
     model_type = "diffusion_gemma_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",
diff --git a/src/transformers/models/dots1/configuration_dots1.py b/src/transformers/models/dots1/configuration_dots1.py
index 944f06dd88c5..f04022e58df4 100644
--- a/src/transformers/models/dots1/configuration_dots1.py
+++ b/src/transformers/models/dots1/configuration_dots1.py
@@ -48,7 +48,6 @@ class Dots1Config(PreTrainedConfig):
     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",
diff --git a/src/transformers/models/dots1/modular_dots1.py b/src/transformers/models/dots1/modular_dots1.py
index 47beb2f3e264..ea3a83ec51ed 100644
--- a/src/transformers/models/dots1/modular_dots1.py
+++ b/src/transformers/models/dots1/modular_dots1.py
@@ -65,7 +65,6 @@ class Dots1Config(PreTrainedConfig):
     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",
diff --git a/src/transformers/models/embedding_gemma2/configuration_embedding_gemma2.py b/src/transformers/models/embedding_gemma2/configuration_embedding_gemma2.py
index e0be111fce05..b3c173cced61 100644
--- a/src/transformers/models/embedding_gemma2/configuration_embedding_gemma2.py
+++ b/src/transformers/models/embedding_gemma2/configuration_embedding_gemma2.py
@@ -51,7 +51,6 @@ class EmbeddingGemma2TextConfig(PreTrainedConfig):
 
     model_type = "embedding_gemma2_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",
diff --git a/src/transformers/models/embedding_gemma2/modular_embedding_gemma2.py b/src/transformers/models/embedding_gemma2/modular_embedding_gemma2.py
index c3f2f2e4620b..fe8fefbaf443 100644
--- a/src/transformers/models/embedding_gemma2/modular_embedding_gemma2.py
+++ b/src/transformers/models/embedding_gemma2/modular_embedding_gemma2.py
@@ -83,7 +83,6 @@ class EmbeddingGemma2TextConfig(PreTrainedConfig):
 
     model_type = "embedding_gemma2_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",
diff --git a/src/transformers/models/ernie4_5/configuration_ernie4_5.py b/src/transformers/models/ernie4_5/configuration_ernie4_5.py
index ec422f957085..896fb0b99402 100644
--- a/src/transformers/models/ernie4_5/configuration_ernie4_5.py
+++ b/src/transformers/models/ernie4_5/configuration_ernie4_5.py
@@ -47,7 +47,6 @@ class Ernie4_5Config(PreTrainedConfig):
     default_theta = 500000.0
     # Default tensor parallel plan for base model `Ernie4_5Model`
     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",
diff --git a/src/transformers/models/ernie4_5_moe/configuration_ernie4_5_moe.py b/src/transformers/models/ernie4_5_moe/configuration_ernie4_5_moe.py
index f536a6b66390..142ff7a2b25c 100644
--- a/src/transformers/models/ernie4_5_moe/configuration_ernie4_5_moe.py
+++ b/src/transformers/models/ernie4_5_moe/configuration_ernie4_5_moe.py
@@ -63,7 +63,6 @@ class Ernie4_5_MoeConfig(PreTrainedConfig):
 
     # Default tensor parallel plan for base model `Ernie4_5_MoE`
     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",
diff --git a/src/transformers/models/ernie4_5_vl_moe/configuration_ernie4_5_vl_moe.py b/src/transformers/models/ernie4_5_vl_moe/configuration_ernie4_5_vl_moe.py
index 8cfb60726553..3cb2029656c8 100644
--- a/src/transformers/models/ernie4_5_vl_moe/configuration_ernie4_5_vl_moe.py
+++ b/src/transformers/models/ernie4_5_vl_moe/configuration_ernie4_5_vl_moe.py
@@ -86,7 +86,6 @@ class Ernie4_5_VLMoeTextConfig(PreTrainedConfig):
     default_theta = 500000.0
 
     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",
diff --git a/src/transformers/models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py b/src/transformers/models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py
index e823b4ef7ae7..883b09855edf 100644
--- a/src/transformers/models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py
+++ b/src/transformers/models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py
@@ -120,7 +120,6 @@ class Ernie4_5_VLMoeTextConfig(Ernie4_5_MoeConfig):
     base_config_key = "text_config"
 
     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",
diff --git a/src/transformers/models/esmc/configuration_esmc.py b/src/transformers/models/esmc/configuration_esmc.py
index 038561995f99..219091d5001f 100644
--- a/src/transformers/models/esmc/configuration_esmc.py
+++ b/src/transformers/models/esmc/configuration_esmc.py
@@ -53,7 +53,6 @@ class EsmcConfig(PreTrainedConfig):
     model_type = "esmc"
     # Default tensor parallel plan for base model `EsmcModel`
     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",
diff --git a/src/transformers/models/eurobert/configuration_eurobert.py b/src/transformers/models/eurobert/configuration_eurobert.py
index 189a2354852c..f64c4f7e5a11 100644
--- a/src/transformers/models/eurobert/configuration_eurobert.py
+++ b/src/transformers/models/eurobert/configuration_eurobert.py
@@ -51,7 +51,6 @@ class EuroBertConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `EuroBertModel`
     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",
diff --git a/src/transformers/models/exaone4/configuration_exaone4.py b/src/transformers/models/exaone4/configuration_exaone4.py
index 0550473cf7a9..f29cab8dd8ea 100644
--- a/src/transformers/models/exaone4/configuration_exaone4.py
+++ b/src/transformers/models/exaone4/configuration_exaone4.py
@@ -60,7 +60,6 @@ class Exaone4Config(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `LlamaModel`
     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",
diff --git a/src/transformers/models/exaone4/modular_exaone4.py b/src/transformers/models/exaone4/modular_exaone4.py
index 1bd405e9dc92..b20b18717d2d 100644
--- a/src/transformers/models/exaone4/modular_exaone4.py
+++ b/src/transformers/models/exaone4/modular_exaone4.py
@@ -89,7 +89,6 @@ class Exaone4Config(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `LlamaModel`
     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",
diff --git a/src/transformers/models/exaone_moe/configuration_exaone_moe.py b/src/transformers/models/exaone_moe/configuration_exaone_moe.py
index 13561a4b3325..9cdab1097f8e 100644
--- a/src/transformers/models/exaone_moe/configuration_exaone_moe.py
+++ b/src/transformers/models/exaone_moe/configuration_exaone_moe.py
@@ -66,7 +66,6 @@ class ExaoneMoeConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `LlamaModel`
     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",
diff --git a/src/transformers/models/flex_olmo/configuration_flex_olmo.py b/src/transformers/models/flex_olmo/configuration_flex_olmo.py
index 82cd1b44717a..921ef2b2597c 100644
--- a/src/transformers/models/flex_olmo/configuration_flex_olmo.py
+++ b/src/transformers/models/flex_olmo/configuration_flex_olmo.py
@@ -50,7 +50,6 @@ class FlexOlmoConfig(PreTrainedConfig):
     attribute_map = {"num_local_experts": "num_experts"}
     default_theta = 500000.0
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
         "layers.*.self_attn.k_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
         "layers.*.self_attn.v_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
diff --git a/src/transformers/models/flex_olmo/modular_flex_olmo.py b/src/transformers/models/flex_olmo/modular_flex_olmo.py
index b623367445b9..956bea08670c 100644
--- a/src/transformers/models/flex_olmo/modular_flex_olmo.py
+++ b/src/transformers/models/flex_olmo/modular_flex_olmo.py
@@ -60,7 +60,6 @@ class FlexOlmoConfig(PreTrainedConfig):
     attribute_map = {"num_local_experts": "num_experts"}
     default_theta = 500000.0
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
         "layers.*.self_attn.k_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
         "layers.*.self_attn.v_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
diff --git a/src/transformers/models/gemma/configuration_gemma.py b/src/transformers/models/gemma/configuration_gemma.py
index ce0f63e14c75..67f8ae5f143d 100644
--- a/src/transformers/models/gemma/configuration_gemma.py
+++ b/src/transformers/models/gemma/configuration_gemma.py
@@ -50,7 +50,6 @@ class GemmaConfig(PreTrainedConfig):
     model_type = "gemma"
     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",
diff --git a/src/transformers/models/gemma/modular_gemma.py b/src/transformers/models/gemma/modular_gemma.py
index b9e314c1bf60..272d32a864f7 100644
--- a/src/transformers/models/gemma/modular_gemma.py
+++ b/src/transformers/models/gemma/modular_gemma.py
@@ -66,7 +66,6 @@ class GemmaConfig(PreTrainedConfig):
     model_type = "gemma"
     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",
diff --git a/src/transformers/models/gemma2/configuration_gemma2.py b/src/transformers/models/gemma2/configuration_gemma2.py
index 3cf0a00ad208..11d7b012099f 100644
--- a/src/transformers/models/gemma2/configuration_gemma2.py
+++ b/src/transformers/models/gemma2/configuration_gemma2.py
@@ -51,7 +51,6 @@ class Gemma2Config(PreTrainedConfig):
     model_type = "gemma2"
     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",
diff --git a/src/transformers/models/gemma2/modular_gemma2.py b/src/transformers/models/gemma2/modular_gemma2.py
index d49df566b5e0..ede2b994b568 100644
--- a/src/transformers/models/gemma2/modular_gemma2.py
+++ b/src/transformers/models/gemma2/modular_gemma2.py
@@ -77,7 +77,6 @@ class Gemma2Config(PreTrainedConfig):
     model_type = "gemma2"
     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",
diff --git a/src/transformers/models/gemma3/configuration_gemma3.py b/src/transformers/models/gemma3/configuration_gemma3.py
index e1f933f1e76a..e160c8e9b0eb 100644
--- a/src/transformers/models/gemma3/configuration_gemma3.py
+++ b/src/transformers/models/gemma3/configuration_gemma3.py
@@ -55,7 +55,6 @@ class Gemma3TextConfig(PreTrainedConfig):
     model_type = "gemma3_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",
diff --git a/src/transformers/models/gemma3/modular_gemma3.py b/src/transformers/models/gemma3/modular_gemma3.py
index 2710829160d8..180a0bbea9cb 100644
--- a/src/transformers/models/gemma3/modular_gemma3.py
+++ b/src/transformers/models/gemma3/modular_gemma3.py
@@ -91,7 +91,6 @@ class Gemma3TextConfig(Gemma2Config, PreTrainedConfig):
 
     model_type = "gemma3_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",
diff --git a/src/transformers/models/gemma3n/configuration_gemma3n.py b/src/transformers/models/gemma3n/configuration_gemma3n.py
index 21d189fed03b..adb25b426292 100644
--- a/src/transformers/models/gemma3n/configuration_gemma3n.py
+++ b/src/transformers/models/gemma3n/configuration_gemma3n.py
@@ -76,7 +76,6 @@ class Gemma3nTextConfig(PreTrainedConfig):
     model_type = "gemma3n_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",
diff --git a/src/transformers/models/gemma3n/modular_gemma3n.py b/src/transformers/models/gemma3n/modular_gemma3n.py
index 5eb022580332..565afe3b07cb 100644
--- a/src/transformers/models/gemma3n/modular_gemma3n.py
+++ b/src/transformers/models/gemma3n/modular_gemma3n.py
@@ -116,7 +116,6 @@ class Gemma3nTextConfig(Gemma3TextConfig):
 
     model_type = "gemma3n_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",
diff --git a/src/transformers/models/gemma4/configuration_gemma4.py b/src/transformers/models/gemma4/configuration_gemma4.py
index b60a664ba5d3..6cc0e2a95fd3 100644
--- a/src/transformers/models/gemma4/configuration_gemma4.py
+++ b/src/transformers/models/gemma4/configuration_gemma4.py
@@ -120,7 +120,6 @@ class Gemma4TextConfig(PreTrainedConfig):
     model_type = "gemma4_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",
diff --git a/src/transformers/models/gemma4_unified/configuration_gemma4_unified.py b/src/transformers/models/gemma4_unified/configuration_gemma4_unified.py
index d5a7e5afdcbe..fa8538e7fe7a 100644
--- a/src/transformers/models/gemma4_unified/configuration_gemma4_unified.py
+++ b/src/transformers/models/gemma4_unified/configuration_gemma4_unified.py
@@ -104,7 +104,6 @@ class Gemma4UnifiedTextConfig(PreTrainedConfig):
     model_type = "gemma4_unified_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",
diff --git a/src/transformers/models/gemma4_unified/modular_gemma4_unified.py b/src/transformers/models/gemma4_unified/modular_gemma4_unified.py
index 6ae414a44c69..8adcae22652f 100644
--- a/src/transformers/models/gemma4_unified/modular_gemma4_unified.py
+++ b/src/transformers/models/gemma4_unified/modular_gemma4_unified.py
@@ -444,7 +444,6 @@ class Gemma4UnifiedTextConfig(Gemma4TextConfig):
 
     model_type = "gemma4_unified_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",
diff --git a/src/transformers/models/glm/configuration_glm.py b/src/transformers/models/glm/configuration_glm.py
index 16e54eea0722..98525012a23b 100644
--- a/src/transformers/models/glm/configuration_glm.py
+++ b/src/transformers/models/glm/configuration_glm.py
@@ -40,7 +40,6 @@ class GlmConfig(PreTrainedConfig):
     model_type = "glm"
     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",
diff --git a/src/transformers/models/glm4/configuration_glm4.py b/src/transformers/models/glm4/configuration_glm4.py
index c17bb3f84063..f33129607fec 100644
--- a/src/transformers/models/glm4/configuration_glm4.py
+++ b/src/transformers/models/glm4/configuration_glm4.py
@@ -40,7 +40,6 @@ class Glm4Config(PreTrainedConfig):
     model_type = "glm4"
     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",
diff --git a/src/transformers/models/glm4_moe/configuration_glm4_moe.py b/src/transformers/models/glm4_moe/configuration_glm4_moe.py
index be4f01978aab..884066913209 100644
--- a/src/transformers/models/glm4_moe/configuration_glm4_moe.py
+++ b/src/transformers/models/glm4_moe/configuration_glm4_moe.py
@@ -58,7 +58,6 @@ class Glm4MoeConfig(PreTrainedConfig):
 
     # Default tensor parallel plan for base model `Glm4Moe`
     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",
diff --git a/src/transformers/models/glm4_moe/modular_glm4_moe.py b/src/transformers/models/glm4_moe/modular_glm4_moe.py
index a7c757863297..7bef7d9ca2da 100644
--- a/src/transformers/models/glm4_moe/modular_glm4_moe.py
+++ b/src/transformers/models/glm4_moe/modular_glm4_moe.py
@@ -71,7 +71,6 @@ class Glm4MoeConfig(PreTrainedConfig):
 
     # Default tensor parallel plan for base model `Glm4Moe`
     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",
diff --git a/src/transformers/models/glm4_moe_lite/configuration_glm4_moe_lite.py b/src/transformers/models/glm4_moe_lite/configuration_glm4_moe_lite.py
index 4bd74a037457..b02df0b43b53 100644
--- a/src/transformers/models/glm4_moe_lite/configuration_glm4_moe_lite.py
+++ b/src/transformers/models/glm4_moe_lite/configuration_glm4_moe_lite.py
@@ -52,7 +52,6 @@ class Glm4MoeLiteConfig(PreTrainedConfig):
     model_type = "glm4_moe_lite"
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_b_proj": "colwise",
         "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
         "layers.*.self_attn.kv_b_proj": "colwise",
diff --git a/src/transformers/models/glm4_moe_lite/modular_glm4_moe_lite.py b/src/transformers/models/glm4_moe_lite/modular_glm4_moe_lite.py
index 09acb163f1d5..db05ee5e7b6f 100644
--- a/src/transformers/models/glm4_moe_lite/modular_glm4_moe_lite.py
+++ b/src/transformers/models/glm4_moe_lite/modular_glm4_moe_lite.py
@@ -60,7 +60,6 @@ class Glm4MoeLiteConfig(PreTrainedConfig):
     model_type = "glm4_moe_lite"
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_b_proj": "colwise",
         "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
         "layers.*.self_attn.kv_b_proj": "colwise",
diff --git a/src/transformers/models/glm4v/configuration_glm4v.py b/src/transformers/models/glm4v/configuration_glm4v.py
index dc6bca78ac05..594503efc61a 100644
--- a/src/transformers/models/glm4v/configuration_glm4v.py
+++ b/src/transformers/models/glm4v/configuration_glm4v.py
@@ -93,7 +93,6 @@ class Glm4vTextConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `Glm4v`
     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",
diff --git a/src/transformers/models/glm4v/modular_glm4v.py b/src/transformers/models/glm4v/modular_glm4v.py
index ebcb93b23834..dcfd68bdb28e 100644
--- a/src/transformers/models/glm4v/modular_glm4v.py
+++ b/src/transformers/models/glm4v/modular_glm4v.py
@@ -133,7 +133,6 @@ class Glm4vTextConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `Glm4v`
     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",
diff --git a/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py b/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py
index 4fad96436b1e..f1b8ce70ea63 100644
--- a/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py
+++ b/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py
@@ -53,7 +53,6 @@ class Glm4vMoeTextConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `Glm4vMoe`
     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",
diff --git a/src/transformers/models/glm4v_moe/modular_glm4v_moe.py b/src/transformers/models/glm4v_moe/modular_glm4v_moe.py
index aa75bd70b05f..1143f430c4fb 100644
--- a/src/transformers/models/glm4v_moe/modular_glm4v_moe.py
+++ b/src/transformers/models/glm4v_moe/modular_glm4v_moe.py
@@ -86,7 +86,6 @@ class Glm4vMoeTextConfig(Glm4MoeConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `Glm4vMoe`
     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",
diff --git a/src/transformers/models/glm5_next/configuration_glm5_next.py b/src/transformers/models/glm5_next/configuration_glm5_next.py
index 07decbe6730f..c239123a5b6d 100644
--- a/src/transformers/models/glm5_next/configuration_glm5_next.py
+++ b/src/transformers/models/glm5_next/configuration_glm5_next.py
@@ -70,7 +70,6 @@ class Glm5NextTextConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
 
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_b_proj": "colwise",
         "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
         "layers.*.self_attn.kv_b_proj": "colwise",
diff --git a/src/transformers/models/glm_image/configuration_glm_image.py b/src/transformers/models/glm_image/configuration_glm_image.py
index 42339bcf6d4f..7305ad3d0b3c 100644
--- a/src/transformers/models/glm_image/configuration_glm_image.py
+++ b/src/transformers/models/glm_image/configuration_glm_image.py
@@ -105,7 +105,6 @@ class GlmImageTextConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `GlmImage`
     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",
diff --git a/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py
index d7cc943eb1c8..7f6f457469ed 100644
--- a/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py
+++ b/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py
@@ -62,7 +62,6 @@ class GlmMoeDsaConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
 
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_b_proj": "colwise",
         "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
         "layers.*.self_attn.kv_b_proj": "colwise",
diff --git a/src/transformers/models/glm_ocr/configuration_glm_ocr.py b/src/transformers/models/glm_ocr/configuration_glm_ocr.py
index 624c1edeb90f..3aec6d0a41cb 100644
--- a/src/transformers/models/glm_ocr/configuration_glm_ocr.py
+++ b/src/transformers/models/glm_ocr/configuration_glm_ocr.py
@@ -94,7 +94,6 @@ class GlmOcrTextConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `GlmOcr`
     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",
diff --git a/src/transformers/models/gpt_oss/configuration_gpt_oss.py b/src/transformers/models/gpt_oss/configuration_gpt_oss.py
index 53edd0253c37..ebe7caedbdf6 100644
--- a/src/transformers/models/gpt_oss/configuration_gpt_oss.py
+++ b/src/transformers/models/gpt_oss/configuration_gpt_oss.py
@@ -30,7 +30,6 @@ class GptOssConfig(PreTrainedConfig):
     """
 
     model_type = "gpt_oss"
-    base_model_tp_plan = {"embed_tokens": "embedding_rowwise"}
     attribute_map = {
         "num_experts": "num_local_experts",
     }
diff --git a/src/transformers/models/granite/configuration_granite.py b/src/transformers/models/granite/configuration_granite.py
index f1e10d0e7062..e026cbbe5ff3 100644
--- a/src/transformers/models/granite/configuration_granite.py
+++ b/src/transformers/models/granite/configuration_granite.py
@@ -47,7 +47,6 @@ class GraniteConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `GraniteModel`
     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",
diff --git a/src/transformers/models/granite4_vision/configuration_granite4_vision.py b/src/transformers/models/granite4_vision/configuration_granite4_vision.py
index 24ae0410a941..6f8a26af2119 100644
--- a/src/transformers/models/granite4_vision/configuration_granite4_vision.py
+++ b/src/transformers/models/granite4_vision/configuration_granite4_vision.py
@@ -50,7 +50,6 @@ class Granite4VisionTextConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `Granite4VisionTextModel`
     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",
diff --git a/src/transformers/models/granite_swa/configuration_granite_swa.py b/src/transformers/models/granite_swa/configuration_granite_swa.py
index 7cce18b6c25e..678edf63afac 100644
--- a/src/transformers/models/granite_swa/configuration_granite_swa.py
+++ b/src/transformers/models/granite_swa/configuration_granite_swa.py
@@ -58,7 +58,6 @@ class GraniteSWAConfig(PreTrainedConfig):
     # Extends Granite's plan with the learnable per-head `sinks`, sharded across heads (colwise)
     # to match the q/k/v head-sharding so TP keeps each rank's sink slice aligned with its heads.
     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",
diff --git a/src/transformers/models/granite_swa/modular_granite_swa.py b/src/transformers/models/granite_swa/modular_granite_swa.py
index e0e34f1c9287..042d33c47f25 100644
--- a/src/transformers/models/granite_swa/modular_granite_swa.py
+++ b/src/transformers/models/granite_swa/modular_granite_swa.py
@@ -88,7 +88,6 @@ class GraniteSWAConfig(GraniteConfig):
     # Extends Granite's plan with the learnable per-head `sinks`, sharded across heads (colwise)
     # to match the q/k/v head-sharding so TP keeps each rank's sink slice aligned with its heads.
     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",
diff --git a/src/transformers/models/granitemoe_swa/configuration_granitemoe_swa.py b/src/transformers/models/granitemoe_swa/configuration_granitemoe_swa.py
index df16a882cd7c..578ab17db32f 100644
--- a/src/transformers/models/granitemoe_swa/configuration_granitemoe_swa.py
+++ b/src/transformers/models/granitemoe_swa/configuration_granitemoe_swa.py
@@ -90,7 +90,6 @@ class GraniteMoeSWAConfig(PreTrainedConfig):
     # with the router replicated. The optional shared expert (`shared_mlp`, off by default) is left
     # replicated -- it is small and its full output sums consistently with the all-reduced MoE output.
     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",
diff --git a/src/transformers/models/granitemoe_swa/modular_granitemoe_swa.py b/src/transformers/models/granitemoe_swa/modular_granitemoe_swa.py
index 1857feae135e..63cc0679cbb8 100644
--- a/src/transformers/models/granitemoe_swa/modular_granitemoe_swa.py
+++ b/src/transformers/models/granitemoe_swa/modular_granitemoe_swa.py
@@ -87,7 +87,6 @@ class GraniteMoeSWAConfig(GraniteMoeSharedConfig):
     # with the router replicated. The optional shared expert (`shared_mlp`, off by default) is left
     # replicated -- it is small and its full output sums consistently with the all-reduced MoE output.
     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",
diff --git a/src/transformers/models/helium/configuration_helium.py b/src/transformers/models/helium/configuration_helium.py
index fe6271f0eee8..caa966b58a9d 100644
--- a/src/transformers/models/helium/configuration_helium.py
+++ b/src/transformers/models/helium/configuration_helium.py
@@ -41,7 +41,6 @@ class HeliumConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     default_theta = 100000.0
     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",
diff --git a/src/transformers/models/higgs_audio_v2/configuration_higgs_audio_v2.py b/src/transformers/models/higgs_audio_v2/configuration_higgs_audio_v2.py
index 495806a9d154..97823eb79576 100644
--- a/src/transformers/models/higgs_audio_v2/configuration_higgs_audio_v2.py
+++ b/src/transformers/models/higgs_audio_v2/configuration_higgs_audio_v2.py
@@ -59,7 +59,6 @@ class HiggsAudioV2Config(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `HiggsAudioV2Model`
     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",
diff --git a/src/transformers/models/hrm_text/configuration_hrm_text.py b/src/transformers/models/hrm_text/configuration_hrm_text.py
index d86a155ac4b0..61d29a8311b0 100644
--- a/src/transformers/models/hrm_text/configuration_hrm_text.py
+++ b/src/transformers/models/hrm_text/configuration_hrm_text.py
@@ -56,7 +56,6 @@ class HrmTextConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
 
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         **{f"{stack}.layers.*.self_attn.q_proj": "colwise" for stack in ("L_module", "H_module")},
         **{f"{stack}.layers.*.self_attn.k_proj": "colwise" for stack in ("L_module", "H_module")},
         **{f"{stack}.layers.*.self_attn.v_proj": "colwise" for stack in ("L_module", "H_module")},
diff --git a/src/transformers/models/hrm_text/modular_hrm_text.py b/src/transformers/models/hrm_text/modular_hrm_text.py
index 35975aa2b5b9..274f9d3291c0 100644
--- a/src/transformers/models/hrm_text/modular_hrm_text.py
+++ b/src/transformers/models/hrm_text/modular_hrm_text.py
@@ -73,7 +73,6 @@ class HrmTextConfig(LlamaConfig):
     model_type = "hrm_text"
 
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         **{f"{stack}.layers.*.self_attn.q_proj": "colwise" for stack in ("L_module", "H_module")},
         **{f"{stack}.layers.*.self_attn.k_proj": "colwise" for stack in ("L_module", "H_module")},
         **{f"{stack}.layers.*.self_attn.v_proj": "colwise" for stack in ("L_module", "H_module")},
diff --git a/src/transformers/models/hunyuan_v1_dense/configuration_hunyuan_v1_dense.py b/src/transformers/models/hunyuan_v1_dense/configuration_hunyuan_v1_dense.py
index 52fdc8f05a89..7ac120cec14f 100644
--- a/src/transformers/models/hunyuan_v1_dense/configuration_hunyuan_v1_dense.py
+++ b/src/transformers/models/hunyuan_v1_dense/configuration_hunyuan_v1_dense.py
@@ -31,7 +31,6 @@ class HunYuanDenseV1Config(PreTrainedConfig):
     """
 
     model_type = "hunyuan_v1_dense"
-    base_model_tp_plan = {"embed_tokens": "embedding_rowwise"}
     keys_to_ignore_at_inference = ["past_key_values"]
 
     vocab_size: int = 290943
diff --git a/src/transformers/models/hunyuan_v1_moe/configuration_hunyuan_v1_moe.py b/src/transformers/models/hunyuan_v1_moe/configuration_hunyuan_v1_moe.py
index 5ad0d0bb1244..f20b240507b8 100644
--- a/src/transformers/models/hunyuan_v1_moe/configuration_hunyuan_v1_moe.py
+++ b/src/transformers/models/hunyuan_v1_moe/configuration_hunyuan_v1_moe.py
@@ -32,7 +32,6 @@ class HunYuanMoEV1Config(PreTrainedConfig):
     """
 
     model_type = "hunyuan_v1_moe"
-    base_model_tp_plan = {"embed_tokens": "embedding_rowwise"}
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_ep_plan = {
         "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
diff --git a/src/transformers/models/hunyuan_vl/configuration_hunyuan_vl.py b/src/transformers/models/hunyuan_vl/configuration_hunyuan_vl.py
index b79e2b2cf427..6a6c2d339a16 100644
--- a/src/transformers/models/hunyuan_vl/configuration_hunyuan_vl.py
+++ b/src/transformers/models/hunyuan_vl/configuration_hunyuan_vl.py
@@ -119,7 +119,6 @@ class HunYuanVLTextConfig(PreTrainedConfig):
     """
 
     model_type = "hunyuan_vl_text"
-    base_model_tp_plan = {"embed_tokens": "embedding_rowwise"}
     keys_to_ignore_at_inference = ["past_key_values"]
 
     vocab_size: int = 290943
diff --git a/src/transformers/models/hy_v3/configuration_hy_v3.py b/src/transformers/models/hy_v3/configuration_hy_v3.py
index 807346e23e61..6266f301095c 100644
--- a/src/transformers/models/hy_v3/configuration_hy_v3.py
+++ b/src/transformers/models/hy_v3/configuration_hy_v3.py
@@ -51,7 +51,6 @@ class HYV3Config(PreTrainedConfig):
         "num_local_experts": "num_experts",
     }
     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",
diff --git a/src/transformers/models/hy_v3/modular_hy_v3.py b/src/transformers/models/hy_v3/modular_hy_v3.py
index 62058204fe6b..7aa409123b40 100644
--- a/src/transformers/models/hy_v3/modular_hy_v3.py
+++ b/src/transformers/models/hy_v3/modular_hy_v3.py
@@ -76,7 +76,6 @@ class HYV3Config(PreTrainedConfig):
         "num_local_experts": "num_experts",
     }
     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",
diff --git a/src/transformers/models/hy_v4/configuration_hy_v4.py b/src/transformers/models/hy_v4/configuration_hy_v4.py
index 2131005a3cb5..ac7750fd9cb3 100644
--- a/src/transformers/models/hy_v4/configuration_hy_v4.py
+++ b/src/transformers/models/hy_v4/configuration_hy_v4.py
@@ -62,7 +62,6 @@ class HYV4Config(PreTrainedConfig):
         "num_local_experts": "n_routed_experts",
     }
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_b_proj": "colwise",
         "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
         "layers.*.self_attn.kv_b_proj": "colwise",
diff --git a/src/transformers/models/hy_v4/modular_hy_v4.py b/src/transformers/models/hy_v4/modular_hy_v4.py
index f8a6fa313ca3..85b1aef0d3bb 100644
--- a/src/transformers/models/hy_v4/modular_hy_v4.py
+++ b/src/transformers/models/hy_v4/modular_hy_v4.py
@@ -91,7 +91,6 @@ class HYV4Config(PreTrainedConfig):
         "num_local_experts": "n_routed_experts",
     }
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_b_proj": "colwise",
         "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
         "layers.*.self_attn.kv_b_proj": "colwise",
diff --git a/src/transformers/models/hyperclovax/configuration_hyperclovax.py b/src/transformers/models/hyperclovax/configuration_hyperclovax.py
index 299e3ccab5e8..430a56bf0249 100644
--- a/src/transformers/models/hyperclovax/configuration_hyperclovax.py
+++ b/src/transformers/models/hyperclovax/configuration_hyperclovax.py
@@ -62,7 +62,6 @@ class HyperCLOVAXConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `HyperCLOVAXModel`
     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",
diff --git a/src/transformers/models/inkling/configuration_inkling.py b/src/transformers/models/inkling/configuration_inkling.py
index 12e7c519efa9..2415f9b6638c 100644
--- a/src/transformers/models/inkling/configuration_inkling.py
+++ b/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",
diff --git a/src/transformers/models/inkling/modular_inkling.py b/src/transformers/models/inkling/modular_inkling.py
index 17352fb4c5f1..438fe86d0d10 100644
--- a/src/transformers/models/inkling/modular_inkling.py
+++ b/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",
diff --git a/src/transformers/models/jais2/configuration_jais2.py b/src/transformers/models/jais2/configuration_jais2.py
index d4bc5bba7aaa..7139661d7575 100644
--- a/src/transformers/models/jais2/configuration_jais2.py
+++ b/src/transformers/models/jais2/configuration_jais2.py
@@ -47,7 +47,6 @@ class Jais2Config(PreTrainedConfig):
     model_type = "jais2"
     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",
diff --git a/src/transformers/models/jais2/modular_jais2.py b/src/transformers/models/jais2/modular_jais2.py
index 0fd270cdce53..3fbdf2c8cd46 100644
--- a/src/transformers/models/jais2/modular_jais2.py
+++ b/src/transformers/models/jais2/modular_jais2.py
@@ -31,7 +31,6 @@
 @strict
 class Jais2Config(LlamaConfig):
     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",
diff --git a/src/transformers/models/kimi_linear/configuration_kimi_linear.py b/src/transformers/models/kimi_linear/configuration_kimi_linear.py
index 0ce8cb988c86..90fbf51356ec 100644
--- a/src/transformers/models/kimi_linear/configuration_kimi_linear.py
+++ b/src/transformers/models/kimi_linear/configuration_kimi_linear.py
@@ -44,7 +44,6 @@ class KimiLinearConfig(PreTrainedConfig):
     model_type = "kimi_linear"
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.mlp.experts.gate_up_proj": "packed_colwise",
         "layers.*.mlp.experts.down_proj": "rowwise",
         "layers.*.mlp.experts": "moe_tp_experts",
diff --git a/src/transformers/models/laguna/configuration_laguna.py b/src/transformers/models/laguna/configuration_laguna.py
index 3e09f4107bc2..024937999e4b 100644
--- a/src/transformers/models/laguna/configuration_laguna.py
+++ b/src/transformers/models/laguna/configuration_laguna.py
@@ -61,7 +61,6 @@ class LagunaConfig(PreTrainedConfig):
     model_type = "laguna"
     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",
diff --git a/src/transformers/models/laguna/modular_laguna.py b/src/transformers/models/laguna/modular_laguna.py
index f787576473a7..79b7796c5ee3 100644
--- a/src/transformers/models/laguna/modular_laguna.py
+++ b/src/transformers/models/laguna/modular_laguna.py
@@ -79,7 +79,6 @@ class LagunaConfig(Qwen2MoeConfig):
 
     model_type = "laguna"
     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",
diff --git a/src/transformers/models/lfm2/configuration_lfm2.py b/src/transformers/models/lfm2/configuration_lfm2.py
index 6f291a0f1c7d..f58ce6270d31 100644
--- a/src/transformers/models/lfm2/configuration_lfm2.py
+++ b/src/transformers/models/lfm2/configuration_lfm2.py
@@ -52,7 +52,6 @@ class Lfm2Config(PreTrainedConfig):
     """
 
     model_type = "lfm2"
-    base_model_tp_plan = {"embed_tokens": "embedding_rowwise"}
     keys_to_ignore_at_inference = ["past_key_values"]
     default_theta = 1000000.0
 
diff --git a/src/transformers/models/lfm2_moe/configuration_lfm2_moe.py b/src/transformers/models/lfm2_moe/configuration_lfm2_moe.py
index f1a7c4759f47..0e09fdd607cc 100644
--- a/src/transformers/models/lfm2_moe/configuration_lfm2_moe.py
+++ b/src/transformers/models/lfm2_moe/configuration_lfm2_moe.py
@@ -47,7 +47,6 @@ class Lfm2MoeConfig(PreTrainedConfig):
     """
 
     model_type = "lfm2_moe"
-    base_model_tp_plan = {"embed_tokens": "embedding_rowwise"}
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_ep_plan = {
         "layers.*.feed_forward.experts.gate_up_proj": "grouped_gemm",
diff --git a/src/transformers/models/llama/configuration_llama.py b/src/transformers/models/llama/configuration_llama.py
index f6a08f49e37a..6960a6970592 100644
--- a/src/transformers/models/llama/configuration_llama.py
+++ b/src/transformers/models/llama/configuration_llama.py
@@ -47,7 +47,6 @@ class LlamaConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `LlamaModel`
     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",
diff --git a/src/transformers/models/longcat_flash/configuration_longcat_flash.py b/src/transformers/models/longcat_flash/configuration_longcat_flash.py
index 2befa78aaefe..5a7101897c3e 100644
--- a/src/transformers/models/longcat_flash/configuration_longcat_flash.py
+++ b/src/transformers/models/longcat_flash/configuration_longcat_flash.py
@@ -63,7 +63,6 @@ class LongcatFlashConfig(PreTrainedConfig):
     }
     default_theta = 10000000.0
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.*.q_b_proj": "colwise",
         "layers.*.self_attn.*.kv_a_proj_with_mqa": "mla_kv_a_proj",
         "layers.*.self_attn.*.kv_b_proj": "colwise",
diff --git a/src/transformers/models/mellum/configuration_mellum.py b/src/transformers/models/mellum/configuration_mellum.py
index 8c892961485c..9b3374607297 100644
--- a/src/transformers/models/mellum/configuration_mellum.py
+++ b/src/transformers/models/mellum/configuration_mellum.py
@@ -51,7 +51,6 @@ class MellumConfig(PreTrainedConfig):
 
     # Default tensor parallel plan for base model `Mellum`
     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",
diff --git a/src/transformers/models/mimo_v2_flash/configuration_mimo_v2_flash.py b/src/transformers/models/mimo_v2_flash/configuration_mimo_v2_flash.py
index 86a525f45bf3..e40cc863153a 100644
--- a/src/transformers/models/mimo_v2_flash/configuration_mimo_v2_flash.py
+++ b/src/transformers/models/mimo_v2_flash/configuration_mimo_v2_flash.py
@@ -49,7 +49,6 @@ class MiMoV2FlashConfig(PreTrainedConfig):
     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",
diff --git a/src/transformers/models/mimo_v2_flash/modular_mimo_v2_flash.py b/src/transformers/models/mimo_v2_flash/modular_mimo_v2_flash.py
index 674a1d3a2a76..4c5fd1a9483f 100644
--- a/src/transformers/models/mimo_v2_flash/modular_mimo_v2_flash.py
+++ b/src/transformers/models/mimo_v2_flash/modular_mimo_v2_flash.py
@@ -65,7 +65,6 @@ class MiMoV2FlashConfig(Glm4MoeConfig):
     attribute_map = {"num_local_experts": "n_routed_experts"}
 
     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",
diff --git a/src/transformers/models/minicpm3/configuration_minicpm3.py b/src/transformers/models/minicpm3/configuration_minicpm3.py
index 7283613254da..5367c5497a96 100644
--- a/src/transformers/models/minicpm3/configuration_minicpm3.py
+++ b/src/transformers/models/minicpm3/configuration_minicpm3.py
@@ -68,7 +68,6 @@ class MiniCPM3Config(PreTrainedConfig):
     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.q_b_proj": "colwise",
         "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
diff --git a/src/transformers/models/minicpm3/modular_minicpm3.py b/src/transformers/models/minicpm3/modular_minicpm3.py
index 2b5fe15e68d5..0e229a39c613 100644
--- a/src/transformers/models/minicpm3/modular_minicpm3.py
+++ b/src/transformers/models/minicpm3/modular_minicpm3.py
@@ -82,7 +82,6 @@ class MiniCPM3Config(LlamaConfig):
     """
 
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_proj": "colwise",
         "layers.*.self_attn.q_b_proj": "colwise",
         "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
diff --git a/src/transformers/models/minimax/configuration_minimax.py b/src/transformers/models/minimax/configuration_minimax.py
index 3ac1f05b8074..a7fd614a1d62 100644
--- a/src/transformers/models/minimax/configuration_minimax.py
+++ b/src/transformers/models/minimax/configuration_minimax.py
@@ -62,7 +62,6 @@ class MiniMaxConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     default_theta = 1000000.0
     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",
diff --git a/src/transformers/models/minimax/modular_minimax.py b/src/transformers/models/minimax/modular_minimax.py
index 5c72d2dae97c..37c41a63c126 100644
--- a/src/transformers/models/minimax/modular_minimax.py
+++ b/src/transformers/models/minimax/modular_minimax.py
@@ -89,7 +89,6 @@ class MiniMaxConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     default_theta = 1000000.0
     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",
diff --git a/src/transformers/models/minimax_m2/configuration_minimax_m2.py b/src/transformers/models/minimax_m2/configuration_minimax_m2.py
index 000f8844dae2..3e52eb9f736e 100644
--- a/src/transformers/models/minimax_m2/configuration_minimax_m2.py
+++ b/src/transformers/models/minimax_m2/configuration_minimax_m2.py
@@ -48,7 +48,6 @@ class MiniMaxM2Config(PreTrainedConfig):
     model_type = "minimax_m2"
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_proj": "colwise_gather_output",
         "layers.*.self_attn.k_proj": "colwise_gather_output",
         "layers.*.self_attn.v_proj": "colwise_gather_output",
diff --git a/src/transformers/models/minimax_m2/modular_minimax_m2.py b/src/transformers/models/minimax_m2/modular_minimax_m2.py
index e29fbc5b82e8..57090e31a927 100644
--- a/src/transformers/models/minimax_m2/modular_minimax_m2.py
+++ b/src/transformers/models/minimax_m2/modular_minimax_m2.py
@@ -67,7 +67,6 @@ class MiniMaxM2Config(PreTrainedConfig):
     model_type = "minimax_m2"
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_proj": "colwise_gather_output",
         "layers.*.self_attn.k_proj": "colwise_gather_output",
         "layers.*.self_attn.v_proj": "colwise_gather_output",
diff --git a/src/transformers/models/minimax_m3_vl/configuration_minimax_m3_vl.py b/src/transformers/models/minimax_m3_vl/configuration_minimax_m3_vl.py
index e42b7a975da1..4069a570f96d 100644
--- a/src/transformers/models/minimax_m3_vl/configuration_minimax_m3_vl.py
+++ b/src/transformers/models/minimax_m3_vl/configuration_minimax_m3_vl.py
@@ -55,7 +55,6 @@ class MiniMaxM3VLTextConfig(PreTrainedConfig):
     model_type = "minimax_m3_vl_text"
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_proj": "colwise_gather_output",
         "layers.*.self_attn.k_proj": "colwise_gather_output",
         "layers.*.self_attn.v_proj": "colwise_gather_output",
diff --git a/src/transformers/models/ministral/configuration_ministral.py b/src/transformers/models/ministral/configuration_ministral.py
index 706968c2872c..4ff445e6808a 100644
--- a/src/transformers/models/ministral/configuration_ministral.py
+++ b/src/transformers/models/ministral/configuration_ministral.py
@@ -49,7 +49,6 @@ class MinistralConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `MinistralModel`
     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",
diff --git a/src/transformers/models/ministral3/configuration_ministral3.py b/src/transformers/models/ministral3/configuration_ministral3.py
index 0597792a9ac5..ff0ca8610280 100644
--- a/src/transformers/models/ministral3/configuration_ministral3.py
+++ b/src/transformers/models/ministral3/configuration_ministral3.py
@@ -55,7 +55,6 @@ class Ministral3Config(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `MistralModel`
     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",
diff --git a/src/transformers/models/mistral/configuration_mistral.py b/src/transformers/models/mistral/configuration_mistral.py
index a4d8e020f50d..c57193f58d7b 100644
--- a/src/transformers/models/mistral/configuration_mistral.py
+++ b/src/transformers/models/mistral/configuration_mistral.py
@@ -46,7 +46,6 @@ class MistralConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `MistralModel`
     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",
diff --git a/src/transformers/models/mistral4/configuration_mistral4.py b/src/transformers/models/mistral4/configuration_mistral4.py
index 0e04e7fac8df..eb3c2f3f787d 100644
--- a/src/transformers/models/mistral4/configuration_mistral4.py
+++ b/src/transformers/models/mistral4/configuration_mistral4.py
@@ -47,7 +47,6 @@ class Mistral4Config(PreTrainedConfig):
     model_type = "mistral4"
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.mlp.experts.gate_up_proj": "packed_colwise",
         "layers.*.mlp.experts.down_proj": "rowwise",
         "layers.*.mlp.experts": "moe_tp_experts",
diff --git a/src/transformers/models/mixtral/configuration_mixtral.py b/src/transformers/models/mixtral/configuration_mixtral.py
index 4e57ab005c80..c1158852cf55 100644
--- a/src/transformers/models/mixtral/configuration_mixtral.py
+++ b/src/transformers/models/mixtral/configuration_mixtral.py
@@ -43,7 +43,6 @@ class MixtralConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     default_theta = 1000000.0
     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",
diff --git a/src/transformers/models/muse_glimmer/configuration_muse_glimmer.py b/src/transformers/models/muse_glimmer/configuration_muse_glimmer.py
index ed7454bd9d06..8c423cf8fff2 100644
--- a/src/transformers/models/muse_glimmer/configuration_muse_glimmer.py
+++ b/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",
diff --git a/src/transformers/models/muse_glimmer/modular_muse_glimmer.py b/src/transformers/models/muse_glimmer/modular_muse_glimmer.py
index d070a5317854..23564bc514fd 100644
--- a/src/transformers/models/muse_glimmer/modular_muse_glimmer.py
+++ b/src/transformers/models/muse_glimmer/modular_muse_glimmer.py
@@ -581,7 +581,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",
diff --git a/src/transformers/models/nanochat/configuration_nanochat.py b/src/transformers/models/nanochat/configuration_nanochat.py
index af18cf077772..24a0ab7b6d09 100644
--- a/src/transformers/models/nanochat/configuration_nanochat.py
+++ b/src/transformers/models/nanochat/configuration_nanochat.py
@@ -43,7 +43,6 @@ class NanoChatConfig(PreTrainedConfig):
     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",
diff --git a/src/transformers/models/neomme/configuration_neomme.py b/src/transformers/models/neomme/configuration_neomme.py
index 5418d551fd64..6a1506dd1e5f 100644
--- a/src/transformers/models/neomme/configuration_neomme.py
+++ b/src/transformers/models/neomme/configuration_neomme.py
@@ -59,7 +59,6 @@ class NeoMMEConfig(PreTrainedConfig):
 
     model_type = "neomme"
     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",
diff --git a/src/transformers/models/olmo/configuration_olmo.py b/src/transformers/models/olmo/configuration_olmo.py
index 7d13bc90b3aa..186cc3a704fb 100644
--- a/src/transformers/models/olmo/configuration_olmo.py
+++ b/src/transformers/models/olmo/configuration_olmo.py
@@ -50,7 +50,6 @@ class OlmoConfig(PreTrainedConfig):
     model_type = "olmo"
     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",
diff --git a/src/transformers/models/olmo2/configuration_olmo2.py b/src/transformers/models/olmo2/configuration_olmo2.py
index 3b4754bf8a24..a99a49194305 100644
--- a/src/transformers/models/olmo2/configuration_olmo2.py
+++ b/src/transformers/models/olmo2/configuration_olmo2.py
@@ -53,7 +53,6 @@ class Olmo2Config(PreTrainedConfig):
     model_type = "olmo2"
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
         "layers.*.self_attn.k_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
         "layers.*.self_attn.v_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
diff --git a/src/transformers/models/olmo2/modular_olmo2.py b/src/transformers/models/olmo2/modular_olmo2.py
index d0fbe9213788..4f8dcffeaae0 100644
--- a/src/transformers/models/olmo2/modular_olmo2.py
+++ b/src/transformers/models/olmo2/modular_olmo2.py
@@ -67,7 +67,6 @@ class Olmo2Config(OlmoConfig):
 
     model_type = "olmo2"
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
         "layers.*.self_attn.k_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
         "layers.*.self_attn.v_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
diff --git a/src/transformers/models/olmo3/configuration_olmo3.py b/src/transformers/models/olmo3/configuration_olmo3.py
index d527e08eccf2..f43d5c02b9c5 100644
--- a/src/transformers/models/olmo3/configuration_olmo3.py
+++ b/src/transformers/models/olmo3/configuration_olmo3.py
@@ -48,7 +48,6 @@ class Olmo3Config(PreTrainedConfig):
     model_type = "olmo3"
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
         "layers.*.self_attn.k_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
         "layers.*.self_attn.v_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
diff --git a/src/transformers/models/olmo3/modular_olmo3.py b/src/transformers/models/olmo3/modular_olmo3.py
index b5292f9873e8..015a3461a073 100644
--- a/src/transformers/models/olmo3/modular_olmo3.py
+++ b/src/transformers/models/olmo3/modular_olmo3.py
@@ -66,7 +66,6 @@ class Olmo3Config(Olmo2Config):
     default_theta = 500000.0
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
         "layers.*.self_attn.k_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
         "layers.*.self_attn.v_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
diff --git a/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py
index 5d113e458764..0f7c32e8799e 100644
--- a/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py
+++ b/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py
@@ -73,7 +73,6 @@ class OlmoHybridConfig(PreTrainedConfig):
     model_type = "olmo_hybrid"
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
         "layers.*.self_attn.k_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
         "layers.*.self_attn.v_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
diff --git a/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py
index b93f0cec3865..d0981294b01a 100644
--- a/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py
+++ b/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py
@@ -107,7 +107,6 @@ class OlmoHybridConfig(LlamaConfig):
 
     model_type = "olmo_hybrid"
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
         "layers.*.self_attn.k_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
         "layers.*.self_attn.v_proj": "colwise_gather_output",  # we need to replicate here due to the added norm on q and k
diff --git a/src/transformers/models/olmoe/configuration_olmoe.py b/src/transformers/models/olmoe/configuration_olmoe.py
index 442df02573c4..699703b0d94d 100644
--- a/src/transformers/models/olmoe/configuration_olmoe.py
+++ b/src/transformers/models/olmoe/configuration_olmoe.py
@@ -46,7 +46,6 @@ class OlmoeConfig(PreTrainedConfig):
 
     # Default tensor parallel plan for base model `Olmoe`
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_proj": "colwise_gather_output",  # due to the norm, we have to gather
         "layers.*.self_attn.k_proj": "colwise_gather_output",  # due to the norm, we have to gather
         "layers.*.self_attn.v_proj": "colwise_gather_output",  # due to the norm, we have to gather
diff --git a/src/transformers/models/openai_privacy_filter/configuration_openai_privacy_filter.py b/src/transformers/models/openai_privacy_filter/configuration_openai_privacy_filter.py
index efc0c81291ed..b256f8aadf22 100644
--- a/src/transformers/models/openai_privacy_filter/configuration_openai_privacy_filter.py
+++ b/src/transformers/models/openai_privacy_filter/configuration_openai_privacy_filter.py
@@ -54,7 +54,6 @@ class OpenAIPrivacyFilterConfig(PreTrainedConfig):
     """
 
     model_type = "openai_privacy_filter"
-    base_model_tp_plan = {"embed_tokens": "embedding_rowwise"}
     attribute_map = {
         "num_experts": "num_local_experts",
     }
diff --git a/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py b/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py
index ecc13765d4d3..efcbf9fe7cd2 100644
--- a/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py
+++ b/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py
@@ -97,7 +97,6 @@ class PaddleOCRTextConfig(PreTrainedConfig):
     default_theta = 500000.0
     # Default tensor parallel plan for base model `PaddleOCRTextModel`
     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",
diff --git a/src/transformers/models/phi/configuration_phi.py b/src/transformers/models/phi/configuration_phi.py
index 3720214a1247..2a65f0f16aec 100644
--- a/src/transformers/models/phi/configuration_phi.py
+++ b/src/transformers/models/phi/configuration_phi.py
@@ -46,7 +46,6 @@ class PhiConfig(PreTrainedConfig):
     model_type = "phi"
     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",
diff --git a/src/transformers/models/phi3/configuration_phi3.py b/src/transformers/models/phi3/configuration_phi3.py
index 725c2fb232a2..f85502f8205d 100644
--- a/src/transformers/models/phi3/configuration_phi3.py
+++ b/src/transformers/models/phi3/configuration_phi3.py
@@ -47,7 +47,6 @@ class Phi3Config(PreTrainedConfig):
     model_type = "phi3"
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.qkv_proj": "colwise_gather_output",  # we need to replicate here due to the slicing of qkv
         "layers.*.self_attn.o_proj": "rowwise_split_input",  # input is replicated due to the slicing of qkv
         "layers.*.mlp.gate_up_proj": "colwise_gather_output",  # we need to replicate here due to the `chunk` operation
diff --git a/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py b/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py
index f4c3dbe4dd50..a6f47121057d 100644
--- a/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py
+++ b/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py
@@ -181,7 +181,6 @@ class Phi4MultimodalConfig(PreTrainedConfig):
     model_type = "phi4_multimodal"
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.qkv_proj": "colwise_gather_output",  # we need to replicate here due to the slicing of qkv
         "layers.*.self_attn.o_proj": "rowwise_split_input",  # input is replicated due to the slicing of qkv
         "layers.*.mlp.gate_up_proj": "colwise_gather_output",  # we need to replicate here due to the `chunk` operation
diff --git a/src/transformers/models/phimoe/configuration_phimoe.py b/src/transformers/models/phimoe/configuration_phimoe.py
index 7d83152ce250..89afec64ce3d 100644
--- a/src/transformers/models/phimoe/configuration_phimoe.py
+++ b/src/transformers/models/phimoe/configuration_phimoe.py
@@ -45,7 +45,6 @@ class PhimoeConfig(PreTrainedConfig):
     ```"""
 
     model_type = "phimoe"
-    base_model_tp_plan = {"embed_tokens": "embedding_rowwise"}
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_ep_plan = {
         "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
diff --git a/src/transformers/models/qwen2/configuration_qwen2.py b/src/transformers/models/qwen2/configuration_qwen2.py
index 0815d2ca0c64..ae41af7b211f 100644
--- a/src/transformers/models/qwen2/configuration_qwen2.py
+++ b/src/transformers/models/qwen2/configuration_qwen2.py
@@ -44,7 +44,6 @@ class Qwen2Config(PreTrainedConfig):
 
     # Default tensor parallel plan for base model `Qwen2`
     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",
diff --git a/src/transformers/models/qwen2_5_omni/configuration_qwen2_5_omni.py b/src/transformers/models/qwen2_5_omni/configuration_qwen2_5_omni.py
index 3481bbd99b2f..c4d597e7a44b 100644
--- a/src/transformers/models/qwen2_5_omni/configuration_qwen2_5_omni.py
+++ b/src/transformers/models/qwen2_5_omni/configuration_qwen2_5_omni.py
@@ -154,7 +154,6 @@ class Qwen2_5OmniTextConfig(PreTrainedConfig):
 
     # Default tensor parallel plan for base model `Qwen25OmniText`
     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",
diff --git a/src/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py b/src/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py
index 6592c85fd5d1..3dbdad07c981 100644
--- a/src/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py
+++ b/src/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py
@@ -265,7 +265,6 @@ class Qwen2_5OmniTextConfig(PreTrainedConfig):
 
     # Default tensor parallel plan for base model `Qwen25OmniText`
     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",
diff --git a/src/transformers/models/qwen2_5_vl/configuration_qwen2_5_vl.py b/src/transformers/models/qwen2_5_vl/configuration_qwen2_5_vl.py
index edef654058eb..d73087ba5cce 100644
--- a/src/transformers/models/qwen2_5_vl/configuration_qwen2_5_vl.py
+++ b/src/transformers/models/qwen2_5_vl/configuration_qwen2_5_vl.py
@@ -91,7 +91,6 @@ class Qwen2_5_VLTextConfig(PreTrainedConfig):
     default_theta = 1000000.0
     # Default tensor parallel plan for base model `Qwen2_5_VL`
     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",
diff --git a/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py b/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py
index ef1c003266b7..6f034ef7dde6 100644
--- a/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py
+++ b/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py
@@ -54,7 +54,6 @@ class Qwen2MoeConfig(PreTrainedConfig):
 
     # Default tensor parallel plan for base model `Qwen2Moe`
     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",
diff --git a/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py b/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py
index ac73eae70fbf..6a6fea8893e6 100644
--- a/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py
+++ b/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py
@@ -68,7 +68,6 @@ class Qwen2VLTextConfig(PreTrainedConfig):
     default_theta = 1000000.0
     # Default tensor parallel plan for base model `Qwen2VL`
     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",
diff --git a/src/transformers/models/qwen3/configuration_qwen3.py b/src/transformers/models/qwen3/configuration_qwen3.py
index 9f79fef44d18..07ad0bb24b33 100644
--- a/src/transformers/models/qwen3/configuration_qwen3.py
+++ b/src/transformers/models/qwen3/configuration_qwen3.py
@@ -43,7 +43,6 @@ class Qwen3Config(PreTrainedConfig):
 
     # Default tensor parallel plan for base model `Qwen3`
     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",
diff --git a/src/transformers/models/qwen3_5/configuration_qwen3_5.py b/src/transformers/models/qwen3_5/configuration_qwen3_5.py
index cbda226d8194..ee41a3d3b961 100644
--- a/src/transformers/models/qwen3_5/configuration_qwen3_5.py
+++ b/src/transformers/models/qwen3_5/configuration_qwen3_5.py
@@ -57,7 +57,6 @@ class Qwen3_5TextConfig(PreTrainedConfig):
     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",
diff --git a/src/transformers/models/qwen3_5/modular_qwen3_5.py b/src/transformers/models/qwen3_5/modular_qwen3_5.py
index 92ad3d17757f..5202822a69fd 100644
--- a/src/transformers/models/qwen3_5/modular_qwen3_5.py
+++ b/src/transformers/models/qwen3_5/modular_qwen3_5.py
@@ -101,7 +101,6 @@ class Qwen3_5TextConfig(Qwen3NextConfig):
     base_config_key = "text_config"
 
     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",
diff --git a/src/transformers/models/qwen3_5_moe/configuration_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/configuration_qwen3_5_moe.py
index 0dfc76911558..84739697ed85 100644
--- a/src/transformers/models/qwen3_5_moe/configuration_qwen3_5_moe.py
+++ b/src/transformers/models/qwen3_5_moe/configuration_qwen3_5_moe.py
@@ -57,7 +57,6 @@ class Qwen3_5MoeTextConfig(PreTrainedConfig):
     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",
diff --git a/src/transformers/models/qwen3_5_moe/modular_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/modular_qwen3_5_moe.py
index 7030ae93cb4f..3e9cab7916c1 100644
--- a/src/transformers/models/qwen3_5_moe/modular_qwen3_5_moe.py
+++ b/src/transformers/models/qwen3_5_moe/modular_qwen3_5_moe.py
@@ -93,7 +93,6 @@ class Qwen3_5MoeTextConfig(Qwen3NextConfig):
     base_config_key = "text_config"
 
     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",
diff --git a/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py b/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py
index f3136f021c84..62cf801edfcf 100644
--- a/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py
+++ b/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py
@@ -54,7 +54,6 @@ class Qwen3MoeConfig(PreTrainedConfig):
 
     # Default tensor parallel plan for base model `Qwen3Moe`
     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",
diff --git a/src/transformers/models/qwen3_next/configuration_qwen3_next.py b/src/transformers/models/qwen3_next/configuration_qwen3_next.py
index f418d8ff733e..bcba1a251bea 100644
--- a/src/transformers/models/qwen3_next/configuration_qwen3_next.py
+++ b/src/transformers/models/qwen3_next/configuration_qwen3_next.py
@@ -59,7 +59,6 @@ class Qwen3NextConfig(PreTrainedConfig):
     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",
diff --git a/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py b/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py
index 5c020fe4483b..c358d4d913bb 100644
--- a/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py
+++ b/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py
@@ -134,7 +134,6 @@ class Qwen3OmniMoeTextConfig(PreTrainedConfig):
 
     # Default tensor parallel plan for base model `Qwen3OmniMoeText`
     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",
@@ -247,7 +246,6 @@ class Qwen3OmniMoeTalkerCodePredictorConfig(PreTrainedConfig):
 
     # Default tensor parallel plan for base model `Qwen3OmniMoeTalkerCodePredictor`
     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",
@@ -337,7 +335,6 @@ class Qwen3OmniMoeTalkerTextConfig(PreTrainedConfig):
 
     # Default tensor parallel plan for base model `Qwen3OmniMoeTalkerText`
     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",
diff --git a/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py b/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py
index 5dfc0fde4b45..3d29469afb0c 100644
--- a/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py
+++ b/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py
@@ -285,7 +285,6 @@ class Qwen3OmniMoeTextConfig(PreTrainedConfig):
 
     # Default tensor parallel plan for base model `Qwen3OmniMoeText`
     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",
diff --git a/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py b/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py
index 15e03e4c2878..def956bb83a1 100644
--- a/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py
+++ b/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py
@@ -56,7 +56,6 @@ class Qwen3VLMoeTextConfig(PreTrainedConfig):
     }
     # Default tensor parallel plan for base model `Qwen3VLMoe`
     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",
diff --git a/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py b/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py
index 1f61a19dcb35..0902fd903091 100644
--- a/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py
+++ b/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py
@@ -83,7 +83,6 @@ class Qwen3VLMoeTextConfig(Qwen3MoeConfig):
     default_theta = 500000.0
     # Default tensor parallel plan for base model `Qwen3VLMoe`
     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",
diff --git a/src/transformers/models/qwen4_exp/configuration_qwen4_exp.py b/src/transformers/models/qwen4_exp/configuration_qwen4_exp.py
index 3cb9d97978a4..9a7e85300baa 100644
--- a/src/transformers/models/qwen4_exp/configuration_qwen4_exp.py
+++ b/src/transformers/models/qwen4_exp/configuration_qwen4_exp.py
@@ -97,7 +97,6 @@ class Qwen4ExpTextConfig(PreTrainedConfig):
         # 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": "embedding_colwise",
-        "embed_tokens": "embedding_rowwise",
     }
     base_model_pp_plan = None
     base_model_ep_plan = {
diff --git a/src/transformers/models/qwen4_exp/modular_qwen4_exp.py b/src/transformers/models/qwen4_exp/modular_qwen4_exp.py
index 7a6efa3558ae..e23ca547f566 100644
--- a/src/transformers/models/qwen4_exp/modular_qwen4_exp.py
+++ b/src/transformers/models/qwen4_exp/modular_qwen4_exp.py
@@ -142,7 +142,6 @@ class Qwen4ExpTextConfig(Qwen3_5MoeTextConfig):
         # 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": "embedding_colwise",
-        "embed_tokens": "embedding_rowwise",
     }
     base_model_fsdp_plan = {
         "embed_tokens": "free_full_weight",
diff --git a/src/transformers/models/seed_oss/configuration_seed_oss.py b/src/transformers/models/seed_oss/configuration_seed_oss.py
index e4d42a0ea5b8..b1221fcf53ce 100644
--- a/src/transformers/models/seed_oss/configuration_seed_oss.py
+++ b/src/transformers/models/seed_oss/configuration_seed_oss.py
@@ -45,7 +45,6 @@ class SeedOssConfig(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `SeedOssModel`
     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",
diff --git a/src/transformers/models/smollm3/configuration_smollm3.py b/src/transformers/models/smollm3/configuration_smollm3.py
index 9053ada91fce..f48c979a1dd3 100644
--- a/src/transformers/models/smollm3/configuration_smollm3.py
+++ b/src/transformers/models/smollm3/configuration_smollm3.py
@@ -55,7 +55,6 @@ class SmolLM3Config(PreTrainedConfig):
     default_theta = 2000000.0
 
     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",
diff --git a/src/transformers/models/smollm3/modular_smollm3.py b/src/transformers/models/smollm3/modular_smollm3.py
index fe7a04fddc9d..f75017ad2645 100644
--- a/src/transformers/models/smollm3/modular_smollm3.py
+++ b/src/transformers/models/smollm3/modular_smollm3.py
@@ -71,7 +71,6 @@ class SmolLM3Config(PreTrainedConfig):
     default_theta = 2000000.0
 
     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",
diff --git a/src/transformers/models/solar_open/configuration_solar_open.py b/src/transformers/models/solar_open/configuration_solar_open.py
index ef49b7bbd3a8..5d7b88efc59b 100644
--- a/src/transformers/models/solar_open/configuration_solar_open.py
+++ b/src/transformers/models/solar_open/configuration_solar_open.py
@@ -38,7 +38,6 @@ class SolarOpenConfig(PreTrainedConfig):
 
     # Default tensor parallel plan for base model `SolarOpenModel`
     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",
diff --git a/src/transformers/models/solar_open/modular_solar_open.py b/src/transformers/models/solar_open/modular_solar_open.py
index 6a50edbac46d..74236a466411 100644
--- a/src/transformers/models/solar_open/modular_solar_open.py
+++ b/src/transformers/models/solar_open/modular_solar_open.py
@@ -44,7 +44,6 @@ class SolarOpenConfig(Glm4MoeConfig):
 
     # Default tensor parallel plan for base model `SolarOpenModel`
     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",
diff --git a/src/transformers/models/starcoder2/configuration_starcoder2.py b/src/transformers/models/starcoder2/configuration_starcoder2.py
index 19fef828c593..59efa94fc5f4 100644
--- a/src/transformers/models/starcoder2/configuration_starcoder2.py
+++ b/src/transformers/models/starcoder2/configuration_starcoder2.py
@@ -45,7 +45,6 @@ class Starcoder2Config(PreTrainedConfig):
     keys_to_ignore_at_inference = ["past_key_values"]
     # Default tensor parallel plan for base model `Starcoder2`
     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",
diff --git a/src/transformers/models/step3p7/configuration_step3p7.py b/src/transformers/models/step3p7/configuration_step3p7.py
index 89728fc16178..56c747be1dfa 100644
--- a/src/transformers/models/step3p7/configuration_step3p7.py
+++ b/src/transformers/models/step3p7/configuration_step3p7.py
@@ -113,7 +113,6 @@ class Step3p7TextConfig(PreTrainedConfig):
     # gates their gathered output). Spelled out in full, not `{**MiniMaxM3VLTextConfig.base_model_tp_plan, ...}`:
     # generated files have no cross-model imports, so that name wouldn't resolve there.
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_proj": "colwise_gather_output",
         "layers.*.self_attn.k_proj": "colwise_gather_output",
         "layers.*.self_attn.v_proj": "colwise_gather_output",
diff --git a/src/transformers/models/step3p7/modular_step3p7.py b/src/transformers/models/step3p7/modular_step3p7.py
index fb659c197f01..384f2cbfd5cf 100644
--- a/src/transformers/models/step3p7/modular_step3p7.py
+++ b/src/transformers/models/step3p7/modular_step3p7.py
@@ -166,7 +166,6 @@ class Step3p7TextConfig(MiniMaxM3VLTextConfig):
     # gates their gathered output). Spelled out in full, not `{**MiniMaxM3VLTextConfig.base_model_tp_plan, ...}`:
     # generated files have no cross-model imports, so that name wouldn't resolve there.
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.self_attn.q_proj": "colwise_gather_output",
         "layers.*.self_attn.k_proj": "colwise_gather_output",
         "layers.*.self_attn.v_proj": "colwise_gather_output",
diff --git a/src/transformers/models/t5gemma/configuration_t5gemma.py b/src/transformers/models/t5gemma/configuration_t5gemma.py
index a2051edaf06d..9509b7df9d09 100644
--- a/src/transformers/models/t5gemma/configuration_t5gemma.py
+++ b/src/transformers/models/t5gemma/configuration_t5gemma.py
@@ -51,7 +51,6 @@ class T5GemmaModuleConfig(PreTrainedConfig):
     model_type = "t5_gemma_module"
     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",
diff --git a/src/transformers/models/t5gemma2/configuration_t5gemma2.py b/src/transformers/models/t5gemma2/configuration_t5gemma2.py
index 692f87e36ac6..65bcc1a43ac4 100644
--- a/src/transformers/models/t5gemma2/configuration_t5gemma2.py
+++ b/src/transformers/models/t5gemma2/configuration_t5gemma2.py
@@ -42,7 +42,6 @@ class T5Gemma2TextConfig(PreTrainedConfig):
     model_type = "t5gemma2_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",
@@ -211,7 +210,6 @@ class T5Gemma2DecoderConfig(PreTrainedConfig):
     model_type = "t5gemma2_decoder"
     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",
diff --git a/src/transformers/models/vaultgemma/configuration_vaultgemma.py b/src/transformers/models/vaultgemma/configuration_vaultgemma.py
index 0cc086235448..a60b7e8edc0c 100644
--- a/src/transformers/models/vaultgemma/configuration_vaultgemma.py
+++ b/src/transformers/models/vaultgemma/configuration_vaultgemma.py
@@ -50,7 +50,6 @@ class VaultGemmaConfig(PreTrainedConfig):
     model_type = "vaultgemma"
     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",
diff --git a/src/transformers/models/voxtral_realtime/configuration_voxtral_realtime.py b/src/transformers/models/voxtral_realtime/configuration_voxtral_realtime.py
index 82ff27fb55e6..3ba420f94ade 100644
--- a/src/transformers/models/voxtral_realtime/configuration_voxtral_realtime.py
+++ b/src/transformers/models/voxtral_realtime/configuration_voxtral_realtime.py
@@ -27,7 +27,6 @@ class VoxtralRealtimeTextConfig(PreTrainedConfig):
     model_type = "voxtral_realtime_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",
diff --git a/src/transformers/models/youtu/configuration_youtu.py b/src/transformers/models/youtu/configuration_youtu.py
index 75c9989333fe..6d9f2cef1f96 100644
--- a/src/transformers/models/youtu/configuration_youtu.py
+++ b/src/transformers/models/youtu/configuration_youtu.py
@@ -51,7 +51,6 @@ class YoutuConfig(PreTrainedConfig):
     model_type = "youtu"
     keys_to_ignore_at_inference = ["past_key_values"]
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.mlp.gate_proj": "colwise",
         "layers.*.mlp.up_proj": "colwise",
         "layers.*.mlp.down_proj": "rowwise",
diff --git a/src/transformers/models/youtu/modular_youtu.py b/src/transformers/models/youtu/modular_youtu.py
index 0f33d53ee937..c5ba1f03da53 100644
--- a/src/transformers/models/youtu/modular_youtu.py
+++ b/src/transformers/models/youtu/modular_youtu.py
@@ -60,7 +60,6 @@ class YoutuConfig(DeepseekV3Config):
 
     model_type = "youtu"
     base_model_tp_plan = {
-        "embed_tokens": "embedding_rowwise",
         "layers.*.mlp.gate_proj": "colwise",
         "layers.*.mlp.up_proj": "colwise",
         "layers.*.mlp.down_proj": "rowwise",
diff --git a/tests/models/dbrx/test_modeling_dbrx.py b/tests/models/dbrx/test_modeling_dbrx.py
index d2458c67af53..63e10bb41388 100644
--- a/tests/models/dbrx/test_modeling_dbrx.py
+++ b/tests/models/dbrx/test_modeling_dbrx.py
@@ -95,10 +95,6 @@ def test_model_from_pretrained(self):
 
     # Offload does not work with Dbrx models because of the forward of DbrxExperts where we chunk the experts.
     # The issue is that the offloaded weights of the mlp layer are still on meta device (w1_chunked, v1_chunked, w2_chunked)
-    @unittest.skip(reason="DBRX doesn't support tied word embeddings")
-    def test_tp_forward_1(self):
-        pass
-
     @unittest.skip(reason="Dbrx models do not work with offload")
     def test_cpu_offload(self):
         pass
diff --git a/tests/models/phimoe/test_modeling_phimoe.py b/tests/models/phimoe/test_modeling_phimoe.py
index 0be32266462d..7598d74c4dc7 100644
--- a/tests/models/phimoe/test_modeling_phimoe.py
+++ b/tests/models/phimoe/test_modeling_phimoe.py
@@ -101,10 +101,6 @@ def is_pipeline_test_to_skip(
     ):
         return True
 
-    @unittest.skip("PhiMoE routes by Gumbel sampling in training, so TP and reference backward pick different experts")
-    def test_tp_backward(self):
-        pass
-
     @unittest.skip("PhiMoE's RoPE has custom parameterization")
     def test_model_rope_scaling_frequencies(self):
         pass
diff --git a/tests/test_tensor_parallel_mixin.py b/tests/test_tensor_parallel_mixin.py
index af559bfcea5c..fc99db413a81 100644
--- a/tests/test_tensor_parallel_mixin.py
+++ b/tests/test_tensor_parallel_mixin.py
@@ -583,12 +583,13 @@ def _skip_if_not_supported(self, expert_parallel: bool = False):
         #     self.skipTest("VLM models are not yet supported in TP tests")
 
     def test_tp_plan_shards_input_embeddings(self):
-        """A model with a TP plan must shard its input embeddings: replicated, they cost `vocab x hidden` on every
-        rank, and a tied `lm_head` sharded on the vocabulary would not match the embedding weight it shares."""
+        """A model whose TP plan shards its decoder must shard its input embeddings too: replicated, they cost
+        `vocab x hidden` on every rank, and a tied `lm_head` sharded on the vocabulary would not match the embedding
+        weight it shares. A plan that only covers the `lm_head` is not a decoder plan."""
         config, _ = self.model_tester.prepare_config_and_inputs_for_common()
         for model_class in self.all_model_classes:
             model = model_class(copy.deepcopy(config))
-            if not model.tp_plan:
+            if not any(key.rsplit(".", 1)[-1] != "lm_head" for key in model.tp_plan):
                 continue
             embeddings = model.get_input_embeddings()
             if not isinstance(embeddings, torch.nn.Embedding):

From 1fa6155199aea838af2c84d9ac0165630e365fe1 Mon Sep 17 00:00:00 2001
From: IlyasMoutawwakil <moutawwakil.ilyas.tsi@gmail.com>
Date: Wed, 14 Oct 2026 11:47:33 +0200
Subject: [PATCH 09/13] TP: add the input-embedding default in
 `init_parallel_plans`

The config-level default added `embed_tokens` to every plan, including vision towers and models whose embedding has another name, which `test_tp_plan_matches_params` rejects. The model now adds `embed_tokens: embedding_rowwise` only when its base model has an `nn.Embedding` of that name; composite models pick it up from their text sub-model's plan. The tied-embeddings config default is left as on main for backward compatibility.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---
 src/transformers/configuration_utils.py | 7 +++++--
 src/transformers/distributed/mixin.py   | 5 +++++
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/src/transformers/configuration_utils.py b/src/transformers/configuration_utils.py
index 786bbbf4a2b9..3f32558c4ad6 100755
--- a/src/transformers/configuration_utils.py
+++ b/src/transformers/configuration_utils.py
@@ -507,8 +507,11 @@ def __post_init__(self, **kwargs):
             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 self.base_model_tp_plan is not None:
-            self.base_model_tp_plan = {"embed_tokens": "embedding_rowwise", **self.base_model_tp_plan}
+        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)
diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py
index cb02ae43ab6a..6e5ff9d09912 100644
--- a/src/transformers/distributed/mixin.py
+++ b/src/transformers/distributed/mixin.py
@@ -18,6 +18,8 @@
 import warnings
 from typing import TYPE_CHECKING
 
+import torch
+
 from ..utils import is_torch_greater_or_equal, logging
 from ..utils.hub import create_and_tag_model_card
 from .configuration_utils import DistributedConfig
@@ -82,6 +84,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):

From 0b178707347e16afa4f13c565d4a30bb1dc5d751 Mon Sep 17 00:00:00 2001
From: IlyasMoutawwakil <moutawwakil.ilyas.tsi@gmail.com>
Date: Wed, 14 Oct 2026 11:56:31 +0200
Subject: [PATCH 10/13] TP: drop the tied-embeddings config default

`init_parallel_plans` now adds the same `embed_tokens` entry, tied or not, so the config-level one only duplicated it (and overrode a plan's own `embed_tokens` style).

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---
 src/transformers/configuration_utils.py | 7 -------
 1 file changed, 7 deletions(-)

diff --git a/src/transformers/configuration_utils.py b/src/transformers/configuration_utils.py
index 3f32558c4ad6..0c9e00f57a12 100755
--- a/src/transformers/configuration_utils.py
+++ b/src/transformers/configuration_utils.py
@@ -506,13 +506,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)
 

From ad65de51df0e7f2099916fb303687e9797ef3506 Mon Sep 17 00:00:00 2001
From: Ilyas Moutawwakil <57442720+IlyasMoutawwakil@users.noreply.github.com>
Date: Mon, 19 Oct 2026 09:17:18 +0200
Subject: [PATCH 11/13] Apply suggestion from @IlyasMoutawwakil

---
 src/transformers/distributed/tensor_parallel.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py
index 93b477c93dde..6cbb7aaabe97 100644
--- a/src/transformers/distributed/tensor_parallel.py
+++ b/src/transformers/distributed/tensor_parallel.py
@@ -172,8 +172,9 @@ def transform_output_post_forward(self, module, output, mesh):
 
     @staticmethod
     def _from_local_last_dim_shard(local, mesh, size):
-        """Wrap this rank's slice of a tensor sharded on its last dim, whose full width is `size`. `DTensor.from_local`
-        would assume every rank holds an equal slice, which isn't so when `size` doesn't divide by the TP 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()

From 549aae4ee62be84c0ae52bbe55d26509b888f29c Mon Sep 17 00:00:00 2001
From: Ilyas Moutawwakil <57442720+IlyasMoutawwakil@users.noreply.github.com>
Date: Mon, 19 Oct 2026 09:20:56 +0200
Subject: [PATCH 12/13] Apply suggestion from @IlyasMoutawwakil

---
 src/transformers/distributed/tensor_parallel.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py
index 8edfe336581f..b4e6500ee50a 100644
--- a/src/transformers/distributed/tensor_parallel.py
+++ b/src/transformers/distributed/tensor_parallel.py
@@ -172,9 +172,8 @@ def transform_output_post_forward(self, module, output, mesh):
 
     @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.
-    """
+        """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()

From 73370f957b57fc72bc99164566ab73afc101fffd Mon Sep 17 00:00:00 2001
From: Ilyas Moutawwakil <ilyas@huggingface.co>
Date: Mon, 19 Oct 2026 00:47:30 -0700
Subject: [PATCH 13/13] fix nemotron H

---
 src/transformers/models/nemotron_h/configuration_nemotron_h.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/transformers/models/nemotron_h/configuration_nemotron_h.py b/src/transformers/models/nemotron_h/configuration_nemotron_h.py
index 175cf54d2334..8727a80dd49e 100644
--- a/src/transformers/models/nemotron_h/configuration_nemotron_h.py
+++ b/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",