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

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

@@ -507,13 +507,6 @@ def __post_init__(self, **kwargs):         if per_layer_config is not None:             self.per_layer_config = per_layer_config -        # TODO: to support models whose input embedding module is not named `embed_tokens` (e.g. GPT-NeoX's `embed_in`).-        if getattr(self, "tie_word_embeddings", False) and self.base_model_tp_plan is not None:-            self.base_model_tp_plan = {-                **self.base_model_tp_plan,-                "embed_tokens": "embedding_rowwise",-            }-         # Remap layer types if needed         remap_legacy_layer_types(config=self) 

@@ -19,6 +19,8 @@ import warnings from typing import TYPE_CHECKING +import torch+ from ..dynamic_module_utils import custom_object_save from ..utils import is_torch_greater_or_equal, logging from ..utils.hub import create_and_tag_model_card@@ -90,6 +92,9 @@ def init_parallel_plans(self) -> None:             self._tp_plan.update(self.config.base_model_tp_plan or {})             self._ep_plan.update(self.config.base_model_ep_plan or {})             self._fsdp_plan.update(self.config.base_model_fsdp_plan or {})+            # Shard the input embeddings on the vocabulary unless the plan already names them.+            if self._tp_plan and isinstance(getattr(self, "embed_tokens", None), torch.nn.Embedding):+                self._tp_plan = {"embed_tokens": "embedding_rowwise", **self._tp_plan}          for name, module in self.named_children():             if plan := getattr(module, "_ep_plan", None):

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

@@ -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"]),

@@ -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"]),

@@ -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

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

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

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

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

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

@@ -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 = {

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

@@ -234,18 +234,6 @@ def _make_dtensor_shard_op(self, mesh, placement, param_shape, local_shape):         op._axis0_size = param_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)))@@ -274,18 +262,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):+        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):-            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)-+            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)) @@ -314,7 +305,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)},             },         }

@@ -98,14 +98,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@@ -301,10 +302,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):@@ -549,12 +549,10 @@ def _get_tp_model_class(self):         )      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@@ -599,6 +597,24 @@ 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 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 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):+                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`",+            )+     def test_moe_parallel_plans_shard_experts(self):         """An MoE model's expert weights, the bulk of its parameters, must be sharded by every parallel plan it defines.         A plan missing them keeps every expert on every rank without changing any output, so no numerical test sees it.