diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index 02e5d910b589..31493a3a0b3b 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/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: diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 3f14d74979fe..4af362092dd8 100644 --- a/src/transformers/modeling_utils.py +++ b/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` diff --git a/tests/tensor_parallel/test_tensor_parallel.py b/tests/tensor_parallel/test_tensor_parallel.py index 91770b683e45..3d0e644f69cd 100644 --- a/tests/tensor_parallel/test_tensor_parallel.py +++ b/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 diff --git a/tests/test_tensor_parallel_mixin.py b/tests/test_tensor_parallel_mixin.py index 618009e42b02..d965e15ad752 100644 --- a/tests/test_tensor_parallel_mixin.py +++ b/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