🚨 Pass tp_plan from lm_heads

src/transformers/integrations/tensor_parallel.py

@@ -674,6 +674,10 @@ def shard_tensor(
     ) -> torch.Tensor:
         raise NotImplementedError
 
+    def validate_module(self, module: nn.Module, device_mesh, layer_name: str = ""):
+        """Raise if the module cannot be sharded with this style on the given mesh."""
+        pass
+
     def prepare_module_tp(self, module: nn.Module, device_mesh, **kwargs) -> nn.Module:
         distribute_module(
             module,
@@ -719,6 +723,16 @@ def __init__(self, gather_output: bool = False, **kwargs):
         super().__init__(**kwargs)
         self.gather_output = gather_output
 
+    def validate_module(self, module: nn.Module, device_mesh, layer_name: str = ""):
+        out_features = getattr(module, "out_features", None)
+        if self.gather_output and out_features is not None and out_features % device_mesh.size() != 0:
+            raise ValueError(
+                f"`{layer_name}` ({type(module).__name__} with out_features={out_features}) is sharded with "
+                f"'colwise_gather_output', which requires out_features to be divisible by the number of ranks "
+                f"({device_mesh.size()}) to all-gather equal-size shards. Resize the weight (e.g. "
+                f"`model.resize_token_embeddings` for LM heads) or override this module's entry in the tp_plan."
+            )
+
     def _prepare_input_fn(self, mod, inputs, device_mesh):
         input_tensor = inputs[0] if inputs else inputs
         return all_reduce_backward(input_tensor, device_mesh)
@@ -1512,6 +1526,7 @@ def add_tensor_parallel_hooks_to_module(
     """
     if current_module_plan is not None:
         tp_layer = ALL_PARALLEL_STYLES[current_module_plan]
+        tp_layer.validate_module(module, device_mesh, layer_name)
         try:
             tp_layer.prepare_module_tp(module, device_mesh, config=model.config)
         except NotImplementedError as e:

src/transformers/modeling_utils.py

@@ -1268,6 +1268,9 @@ class PreTrainedModel(nn.Module, EmbeddingAccessMixin, ModuleUtilsMixin, PushToH
     # models, this attribute is currently defined in respective model code. For base models, it comes from
     # `config.base_model_pp_plan` during `post_init`.
     _pp_plan: dict[str, tuple[str, str]] = None
+    # An expert parallel plan used instead of `_tp_plan` when expert parallelism is enabled. For base models, it comes
+    # from `config.base_model_ep_plan` during `post_init`.
+    _ep_plan: dict[str, str] = None
     # FSDP2 sharding plan of the form `{"layers.*": "free_full_weight"}`. For top-level models, this attribute is
     # defined on the head class (e.g. `*ForCausalLM`). For base models, it comes from `config.base_model_fsdp_plan`
     # during `post_init`.
@@ -1415,15 +1418,18 @@ def post_init(self):
         """
         # Attach the different parallel plans and tied weight keys to the top-most model, so that everything is
         # easily available.
-        self._tp_plan, self._ep_plan, self._pp_plan, self._fsdp_plan = {}, {}, {}, {}
+        # Start from the class-level plans (e.g. `{"lm_head": "colwise_rep"}` on `...ForCausalLM` classes), copying
+        # them as they are mutated below and would otherwise contaminate the class attribute shared by all instances
+        self._tp_plan = dict(self._tp_plan or {})
+        self._ep_plan = dict(self._ep_plan or {})
+        self._pp_plan = dict(self._pp_plan or {})
+        self._fsdp_plan = dict(self._fsdp_plan or {})
         # If current model is a base model, attach `base_model_tp_plan` and `base_model_pp_plan` from config
         if self.base_model is self:
-            self._pp_plan = self.config.base_model_pp_plan.copy() if self.config.base_model_pp_plan is not None else {}
-            self._tp_plan = self.config.base_model_tp_plan.copy() if self.config.base_model_tp_plan is not None else {}
-            self._ep_plan = self.config.base_model_ep_plan.copy() if self.config.base_model_ep_plan is not None else {}
-            self._fsdp_plan = (
-                self.config.base_model_fsdp_plan.copy() if self.config.base_model_fsdp_plan is not None else {}
-            )
+            self._pp_plan.update(self.config.base_model_pp_plan or {})
+            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 {})
         # Current submodel should register its tied weights
         self.all_tied_weights_keys = self.get_expanded_tied_weights_keys(all_submodels=False)
         # Current submodel should register its `_keep_in_fp32_modules`

tests/tensor_parallel/test_tensor_parallel.py

@@ -25,6 +25,7 @@
     PackedColwiseParallel,
     PackedRowwiseParallel,
     RowwiseParallel,
+    add_tensor_parallel_hooks_to_module,
     get_packed_weights,
     repack_weights,
 )
@@ -166,6 +167,17 @@ def test_tp_plan_none_handling(self):
         model.tp_plan = {"model.layers.*.self_attn.q_proj": "colwise"}
         self.assertEqual(model.tp_plan, {"model.layers.*.self_attn.q_proj": "colwise"})
 
+    def test_post_init_keeps_class_level_plans(self):
+        """Class-level plans (e.g. `lm_head` on ForCausalLM classes) must survive post_init alongside the base model plan."""
+        model_id = "hf-internal-testing/tiny-random-LlamaForCausalLM"
+        model = AutoModelForCausalLM.from_pretrained(model_id, dtype="auto")
+
+        self.assertIn("lm_head", model._tp_plan)
+        self.assertIn("model.layers.*.self_attn.q_proj", model._tp_plan)
+        self.assertIn("lm_head", model._pp_plan)
+        # The merge must not have mutated the class attribute shared by all instances
+        self.assertEqual(set(type(model)._tp_plan), {"lm_head"})
+
 
 @is_tensor_parallel_test
 class TestTensorParallelLayer(TestCasePlus):
@@ -181,6 +193,21 @@ def size(self):
         def get_local_rank(self):
             return self.rank
 
+    def test_colwise_gather_output_rejects_indivisible_out_features(self):
+        device_mesh = self.MockDeviceMesh(world_size=2, rank=0)
+
+        with self.assertRaises(ValueError) as context:
+            add_tensor_parallel_hooks_to_module(
+                model=SimpleNamespace(config=None),
+                module=torch.nn.Linear(8, 99),
+                current_module_plan="colwise_gather_output",
+                layer_name="lm_head",
+                device_mesh=device_mesh,
+            )
+
+        self.assertIn("lm_head", str(context.exception))
+        self.assertIn("divisible", str(context.exception))
+
     def test_colwise_get_expected_sharded_shape(self):
         world_size = 3
         size = 10  # not divisible by world_size to test edge case

tests/test_tensor_parallel_mixin.py

@@ -491,6 +491,15 @@ def _get_tp_model_class(self):
             return self.model_tester.causal_lm_class
         return self.all_model_classes[0]
 
+    def _get_tp_config(self):
+        """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."""
+        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
+        return config
+
     def _skip_if_not_supported(self, expert_parallel: bool = False):
         """Check and skip the test if tensor/expert parallel is not supported for this model/environment."""
         parallelism = "Expert" if expert_parallel else "Tensor"
@@ -539,7 +548,7 @@ def _skip_if_not_supported(self, expert_parallel: bool = False):
     def test_tp_forward(self):
         self._skip_if_not_supported()
 
-        config = self.model_tester.get_config()
+        config = self._get_tp_config()
         model_class = self._get_tp_model_class()
         atol = self.tensor_parallel_atol
         rtol = self.tensor_parallel_rtol
@@ -555,7 +564,7 @@ def test_tp_forward(self):
     def test_tp_backward(self):
         self._skip_if_not_supported()
 
-        config = self.model_tester.get_config()
+        config = self._get_tp_config()
         model_class = self._get_tp_model_class()
         atol = self.tensor_parallel_atol
         rtol = self.tensor_parallel_rtol
@@ -572,7 +581,7 @@ def test_tp_generation(self):
         # Test TP generation: unfused checkpoint → conversion mapping (if needed) → TP sharding → model → generate
         self._skip_if_not_supported()
 
-        config = self.model_tester.get_config()
+        config = self._get_tp_config()
 
         model_class = self._get_tp_model_class()
         atol = self.tensor_parallel_atol
@@ -594,7 +603,7 @@ def test_tp_generation_quantized(self):
         if not is_torchao_available():
             self.skipTest("Test requires torchao")
 
-        config = self.model_tester.get_config()
+        config = self._get_tp_config()
         model_class = self._get_tp_model_class()
         max_new_tokens = 25
 
@@ -611,7 +620,7 @@ def test_tp_generation_quantized(self):
     def test_ep_forward(self):
         self._skip_if_not_supported(expert_parallel=True)
 
-        config = self.model_tester.get_config()
+        config = self._get_tp_config()
         model_class = self._get_tp_model_class()
         atol = self.tensor_parallel_atol
         rtol = self.tensor_parallel_rtol
@@ -627,7 +636,7 @@ def test_ep_forward(self):
     def test_ep_backward(self):
         self._skip_if_not_supported(expert_parallel=True)
 
-        config = self.model_tester.get_config()
+        config = self._get_tp_config()
         model_class = self._get_tp_model_class()
         atol = self.tensor_parallel_atol
         rtol = self.tensor_parallel_rtol