transformers 5.0.0rc0__py3-none-any.whl → 5.0.0rc2__py3-none-any.whl

transformers/integrations/tensor_parallel.py

@@ -19,6 +19,9 @@ import os
 import re
 from functools import partial, reduce
 
+from ..distributed import DistributedConfig
+from ..utils import is_torch_greater_or_equal, logging
+from ..utils.generic import GeneralInterface
 from ..utils.import_utils import is_torch_available
 
 
@@ -27,14 +30,6 @@ if is_torch_available():
     import torch.distributed as dist
     from torch import nn
 
-from ..distributed import DistributedConfig
-from ..utils import is_torch_greater_or_equal, logging
-from ..utils.generic import GeneralInterface
-
-
-logger = logging.get_logger(__name__)
-
-if is_torch_available():
     # Cache this result has it's a C FFI call which can be pretty time-consuming
     _torch_distributed_available = torch.distributed.is_available()
 
@@ -42,6 +37,9 @@ if is_torch_available():
         from torch.distributed.tensor import DTensor, Placement, Replicate, Shard
 
 
+logger = logging.get_logger(__name__)
+
+
 def initialize_tensor_parallelism(
     tp_plan: str | dict[str, str] | None, tp_size: int | None = None, device_mesh=None, device_map=None
 ):
@@ -470,7 +468,12 @@ class TensorParallelLayer:
     @staticmethod
     def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_mesh): ...
 
-    def partition_tensor(self, param, empty_param, param_type, param_casting_dtype, to_contiguous, rank, device_mesh):
+    def shard_tensor(
+        self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None
+    ) -> torch.Tensor:
+        raise NotImplementedError
+
+    def partition_tensor(self, param: torch.Tensor, dtype, to_contiguous: bool):
         raise NotImplementedError
 
     def prepare_module_tp(self, module: nn.Module, device_mesh) -> nn.Module:
@@ -519,19 +522,10 @@ class GatherParallel(TensorParallelLayer):
         return outputs
 
     def shard_tensor(
-        self,
-        param,
-        param_type=None,
-        param_casting_dtype=None,
-        to_contiguous=None,
-        rank=None,
-        device_mesh=None,
-        tensor_idx=None,
-    ):
-        shard = [Replicate()]
-        parameter = param[...].to(param_casting_dtype)
-        self.shard = shard
-        return parameter, shard
+        self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None
+    ) -> torch.Tensor:
+        self.shard = [Replicate()]
+        return param[...].to(device=device, dtype=dtype)
 
     def prepare_module_tp(self, module: nn.Module, device_mesh) -> nn.Module:
         distribute_module(
@@ -562,29 +556,20 @@ class IsolatedParallel(TensorParallelLayer):
         return outputs
 
     def shard_tensor(
-        self,
-        param,
-        param_type=None,
-        param_casting_dtype=None,
-        to_contiguous=None,
-        rank=None,
-        device_mesh=None,
-        tensor_idx=None,
-    ):
-        mesh = device_mesh or self.device_mesh
-        parameter = param[...].to(param_casting_dtype)
-        if mesh is not None:
-            parameter = parameter / mesh.size()
+        self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None
+    ) -> torch.Tensor:
+        parameter = param[...].to(device=device, dtype=dtype)
+        if self.device_mesh is not None:
+            parameter = parameter / self.device_mesh.size()
         self.shard = None
-        return parameter, None
+        return parameter
 
-    def partition_tensor(self, param, empty_param, param_type, param_casting_dtype, to_contiguous, rank, device_mesh):
-        param = param[...].to(param_casting_dtype)
+    def partition_tensor(self, param: torch.Tensor, dtype, to_contiguous: bool):
+        parameter = self.shard_tensor(param, dtype=dtype)
         if to_contiguous:
-            param = param.contiguous()
-        param = param / device_mesh.size()  # TODO should be optionable
+            parameter = parameter.contiguous()
         # TODO: assumes parent module will allreduce the output afterwards (e.g rowlinear bias is IsolatedParallel and parent module is GatherParallel)
-        return param
+        return parameter
 
     def prepare_module_tp(self, module: nn.Module, device_mesh) -> nn.Module:
         distribute_module(
@@ -623,31 +608,15 @@ class ReplicateParallel(TensorParallelLayer):
         return outputs.to_local() if use_local_output and isinstance(outputs, DTensor) else outputs
 
     def shard_tensor(
-        self,
-        param,
-        param_type=None,
-        param_casting_dtype=None,
-        to_contiguous=None,
-        rank=None,
-        device_mesh=None,
-        tensor_idx=None,
-    ):
-        parameter = param[...].to(param_casting_dtype)
-        shard = [Replicate()]
-        self.shard = shard
-        return parameter, shard
-
-    def partition_tensor(self, param, empty_param, param_type, param_casting_dtype, to_contiguous, rank, device_mesh):
-        parameter, shard = self.shard_tensor(
-            param,
-            param_type=param_type,
-            param_casting_dtype=param_casting_dtype,
-            to_contiguous=to_contiguous,
-            rank=rank,
-            device_mesh=device_mesh,
-        )
+        self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None
+    ) -> torch.Tensor:
+        self.shard = [Replicate()]
+        return param[...].to(device=device, dtype=dtype)
+
+    def partition_tensor(self, param: torch.Tensor, dtype, to_contiguous: bool):
+        parameter = self.shard_tensor(param, dtype=dtype)
         if self.use_dtensor:
-            parameter = DTensor.from_local(parameter, device_mesh, shard, run_check=False)
+            parameter = DTensor.from_local(parameter, self.device_mesh, self.shard, run_check=False)
         return parameter
 
 
@@ -685,38 +654,34 @@ class ColwiseParallel(TensorParallelLayer):
         return input_tensor
 
     def shard_tensor(
-        self,
-        param,
-        param_type=None,
-        param_casting_dtype=None,
-        to_contiguous=None,
-        rank=None,
-        device_mesh=None,
-        tensor_idx=None,
-    ):
-        device_mesh = self.device_mesh
-        empty_param = self.empty_param
-        rank = self.rank
-        if param_type == "bias":
-            parameter = get_tensor_shard(param, empty_param, device_mesh, rank, -1, tensor_idx)
+        self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None
+    ) -> torch.Tensor:
+        # If only 1 dim, shard this one (usually it's a `bias`)
+        dim = param.dim() if isinstance(param, torch.Tensor) else len(param.get_shape())
+        if dim == 1:
+            parameter = get_tensor_shard(param, self.empty_param, self.device_mesh, self.rank, -1, tensor_idx)
             shard = [Shard(-1)]
         else:
             shard = [Shard(-2)]
-            parameter = get_tensor_shard(param, empty_param, device_mesh, rank, -2, tensor_idx)
-        parameter = parameter.to(param_casting_dtype)
+            parameter = get_tensor_shard(param, self.empty_param, self.device_mesh, self.rank, -2, tensor_idx)
         self.shard = shard
-        return parameter, shard
+        return parameter.to(device=device, dtype=dtype)
 
-    def partition_tensor(self, param, empty_param, param_type, param_casting_dtype, to_contiguous, rank, device_mesh):
+    def partition_tensor(self, param: torch.Tensor, dtype, to_contiguous: bool):
         # colwise shard weight/bias to Shard(0), weight be Shard(-2) (0 if you have 1 dim only)
         # means Colwise as Linear is input * weight^T + bias, where
         # weight would become Shard(1)
-        parameter, shard = self.shard_tensor(param, param_type, param_casting_dtype, to_contiguous, rank, device_mesh)
+        parameter = self.shard_tensor(param, dtype=dtype)
         if to_contiguous:
             parameter = parameter.contiguous()
         if self.use_dtensor:
             parameter = DTensor.from_local(
-                parameter, device_mesh, shard, run_check=False, shape=empty_param.size(), stride=empty_param.stride()
+                parameter,
+                self.device_mesh,
+                self.shard,
+                run_check=False,
+                shape=self.empty_param.size(),
+                stride=self.empty_param.stride(),
             )
         return nn.Parameter(parameter, requires_grad=parameter.is_floating_point())
 
@@ -731,33 +696,41 @@ class ColwiseParallel(TensorParallelLayer):
 
 class PackedColwiseParallel(ColwiseParallel):
     def shard_tensor(
-        self,
-        param,
-        param_type=None,
-        param_casting_dtype=None,
-        to_contiguous=None,
-        rank=None,
-        device_mesh=None,
-        tensor_idx=None,
-    ):
-        device_mesh = device_mesh or self.device_mesh
-        empty_param = self.empty_param
-        rank = rank if rank is not None else self.rank
-        return get_packed_weights(param, empty_param, device_mesh, rank, -2).to(param_casting_dtype), [Shard(-2)]
+        self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None
+    ) -> torch.Tensor:
+        parameter = get_packed_weights(param, self.empty_param, self.device_mesh, self.rank, -2)
+        return parameter.to(device=device, dtype=dtype)
 
-    def partition_tensor(self, param, empty_param, param_type, param_casting_dtype, to_contiguous, rank, device_mesh):
+    def partition_tensor(self, param: torch.Tensor, dtype, to_contiguous: bool):
         # colwise shard weight/bias to Shard(0), weight be Shard(-2) (0 if you have 1 dim only)
         # means Colwise as Linear is input * weight^T + bias, where
         # weight would become Shard(1)
-        parameter = get_packed_weights(param, empty_param, device_mesh, rank, -2)
-        parameter = parameter.to(param_casting_dtype)
+        parameter = self.shard_tensor(param, dtype=dtype)
         if to_contiguous:
             parameter = parameter.contiguous()
         if self.use_dtensor:
-            parameter = DTensor.from_local(parameter, device_mesh, [Shard(-2)], run_check=False)
+            parameter = DTensor.from_local(parameter, self.device_mesh, [Shard(-2)], run_check=False)
         return nn.Parameter(parameter, requires_grad=parameter.is_floating_point())
 
 
+class LocalColwiseParallel(ColwiseParallel):
+    """
+    Colwise parallel with use_dtensor=False for local tensor operations.
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(use_dtensor=False, **kwargs)
+
+
+class ColwiseParallelReplicate(ColwiseParallel):
+    """
+    Colwise parallel with output layouts replicated.
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(output_layouts=Replicate(), **kwargs)
+
+
 class RowwiseParallel(TensorParallelLayer):
     """
     Partition a compatible nn.Module in a row-wise fashion. Currently supports nn.Linear and nn.Embedding.
@@ -782,7 +755,7 @@ class RowwiseParallel(TensorParallelLayer):
         input_layouts: Placement | None = None,
         output_layouts: Placement | None = None,
         use_local_output: bool = True,
-        use_dtensor=True,
+        use_dtensor: bool = True,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -792,45 +765,36 @@ class RowwiseParallel(TensorParallelLayer):
         self.use_dtensor = use_dtensor
 
     def shard_tensor(
-        self,
-        param,
-        param_type=None,
-        param_casting_dtype=None,
-        to_contiguous=None,
-        rank=None,
-        device_mesh=None,
-        tensor_idx=None,
-    ):
-        device_mesh = device_mesh or self.device_mesh
-        empty_param = self.empty_param
-        rank = rank if rank is not None else self.rank
-        if param_type == "bias":
+        self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None
+    ) -> torch.Tensor:
+        # If only 1 dim, it should not be sharded (usually it's a `bias`)
+        dim = param.dim() if isinstance(param, torch.Tensor) else len(param.get_shape())
+        if dim == 1:
             shard = [Replicate()]
             parameter = param[...]
         else:
-            parameter = get_tensor_shard(param, empty_param, device_mesh, rank, -1, tensor_idx=tensor_idx)
+            parameter = get_tensor_shard(
+                param, self.empty_param, self.device_mesh, self.rank, -1, tensor_idx=tensor_idx
+            )
             shard = [Shard(-1)]
-        parameter = parameter.to(param_casting_dtype)
         self.shard = shard
-        return parameter, shard
+        return parameter.to(device=device, dtype=dtype)
 
-    def partition_tensor(self, param, empty_param, param_type, param_casting_dtype, to_contiguous, rank, device_mesh):
+    def partition_tensor(self, param: torch.Tensor, dtype, to_contiguous: bool):
         # Rowwise shard weight to Shard(1), bias to Replicate(), weight be Shard(1)
         # means Rowwise as nn.Linear is input * weight^T + bias, where
         # weight would become Shard(0)
-        if param_type != "bias":
-            parameter = get_tensor_shard(param, empty_param, device_mesh, rank, -1)
-            shard = [Shard(-1)]
-        else:
-            shard = [Replicate()]
-            parameter = param[:]
-
-        parameter = parameter.to(param_casting_dtype)
+        parameter = self.shard_tensor(param, dtype=dtype)
         if to_contiguous:
             parameter = parameter.contiguous()
         if self.use_dtensor:
             parameter = DTensor.from_local(
-                parameter, device_mesh, shard, run_check=False, shape=empty_param.size(), stride=empty_param.stride()
+                parameter,
+                self.device_mesh,
+                self.shard,
+                run_check=False,
+                shape=self.empty_param.size(),
+                stride=self.empty_param.stride(),
             )
         return nn.Parameter(parameter, requires_grad=parameter.is_floating_point())
 
@@ -886,33 +850,50 @@ class RowwiseParallel(TensorParallelLayer):
 
 class PackedRowwiseParallel(RowwiseParallel):
     def shard_tensor(
-        self,
-        param,
-        param_type=None,
-        param_casting_dtype=None,
-        to_contiguous=None,
-        rank=None,
-        device_mesh=None,
-        tensor_idx=None,
-    ):
-        device_mesh = device_mesh or self.device_mesh
-        empty_param = self.empty_param
-        rank = rank if rank is not None else self.rank
-        return get_packed_weights(param, empty_param, device_mesh, rank, -1), [Shard(-1)]
+        self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None
+    ) -> torch.Tensor:
+        parameter = get_packed_weights(param, self.empty_param, self.device_mesh, self.rank, -1)
+        return parameter.to(device=device, dtype=dtype)
 
-    def partition_tensor(self, param, empty_param, param_type, param_casting_dtype, to_contiguous, rank, device_mesh):
+    def partition_tensor(self, param: torch.Tensor, dtype, to_contiguous: bool):
         # colwise shard weight/bias to Shard(0), weight be Shard(-2) (0 if you have 1 dim only)
         # means Colwise as Linear is input * weight^T + bias, where
         # weight would become Shard(1)
-        parameter = get_packed_weights(param, empty_param, device_mesh, rank, -1)
-        parameter = parameter.to(param_casting_dtype)
+        parameter = self.shard_tensor(param, dtype=dtype)
         if to_contiguous:
             parameter = parameter.contiguous()
         if self.use_dtensor:
-            parameter = DTensor.from_local(parameter, device_mesh, [Shard(-1)], run_check=False)
+            parameter = DTensor.from_local(parameter, self.device_mesh, [Shard(-1)], run_check=False)
         return nn.Parameter(parameter, requires_grad=parameter.is_floating_point())
 
 
+class LocalRowwiseParallel(RowwiseParallel):
+    """
+    Rowwise parallel with use_dtensor=False for local tensor operations.
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(use_dtensor=False, **kwargs)
+
+
+class LocalPackedRowwiseParallel(PackedRowwiseParallel):
+    """
+    Packed rowwise parallel with use_dtensor=False for local tensor operations.
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(use_dtensor=False, **kwargs)
+
+
+class RowwiseParallelReplicate(RowwiseParallel):
+    """
+    Rowwise parallel with input layouts replicated.
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(input_layouts=Replicate(), **kwargs)
+
+
 class SequenceParallel(TensorParallelLayer):
     """
     SequenceParallel replicates a compatible ``nn.Module`` parameters and runs the sharded computation with
@@ -970,18 +951,13 @@ class SequenceParallel(TensorParallelLayer):
 
     def shard_tensor(
         self,
-        param,
-        param_type=None,
-        param_casting_dtype=None,
-        to_contiguous=None,
-        rank=None,
-        device_mesh=None,
+        param: torch.Tensor,
         tensor_idx=None,
-    ):
-        parameter = param[...].to(param_casting_dtype)
-        shard = [Replicate()]
-        self.shard = shard
-        return parameter, shard
+        device=None,
+        dtype=None,
+    ) -> torch.Tensor:
+        self.shard = [Replicate()]
+        return param[...].to(device=device, dtype=dtype)
 
     @staticmethod
     def _prepare_input_fn(input_layouts, desired_input_layouts, mod, inputs, device_mesh):
@@ -999,16 +975,15 @@ class SequenceParallel(TensorParallelLayer):
         )  # maybe we have to replicate ? because next layer is not sharded
         return outputs.to_local()  # if use_local_output else outputs
 
-    def partition_tensor(self, param, empty_param, param_type, param_casting_dtype, to_contiguous, rank, device_mesh):
+    def partition_tensor(self, param: torch.Tensor, dtype, to_contiguous: bool):
         # colwise shard weight/bias to Shard(0), weight be Shard(-2) (0 if you have 1 dim only)
         # means Colwise as Linear is input * weight^T + bias, where
         # weight would become Shard(1)
-        parameter = param[...]
-        parameter = parameter.to(param_casting_dtype)
+        parameter = self.shard_tensor(param, dtype=dtype)
         if to_contiguous:
             parameter = parameter.contiguous()
         if self.use_dtensor:
-            parameter = DTensor.from_local(parameter, device_mesh, [Replicate()], run_check=False)
+            parameter = DTensor.from_local(parameter, self.device_mesh, [Replicate()], run_check=False)
         return nn.Parameter(parameter, requires_grad=parameter.is_floating_point())
 
 
@@ -1022,41 +997,23 @@ class GroupedGemmParallel(TensorParallelLayer):
         self.use_dtensor = False
 
     def shard_tensor(
-        self,
-        param,
-        param_type=None,
-        param_casting_dtype=None,
-        to_contiguous=None,
-        rank=None,
-        device_mesh=None,
-        tensor_idx=None,
-    ):
-        empty_param = self.empty_param
-        ep_rank = self.rank
-        device_mesh = self.device_mesh
-
-        global_num_experts = empty_param.shape[0]
-        if global_num_experts % device_mesh.size() != 0:
+        self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None
+    ) -> torch.Tensor:
+        global_num_experts = self.empty_param.shape[0]
+        if global_num_experts % self.device_mesh.size() != 0:
             raise ValueError(
-                f"Global number of experts must be divisible by number of devices: {global_num_experts} % {device_mesh.size()} != 0"
+                f"Global number of experts must be divisible by number of devices: {global_num_experts} % {self.device_mesh.size()} != 0"
             )
-        local_num_experts = global_num_experts // device_mesh.size()
-        parameter = param[ep_rank * local_num_experts : (ep_rank + 1) * local_num_experts].to(param_casting_dtype)
+        local_num_experts = global_num_experts // self.device_mesh.size()
+        parameter = param[self.rank * local_num_experts : (self.rank + 1) * local_num_experts]
         self.shard = None
-        return parameter, None
+        return parameter.to(device=device, dtype=dtype)
 
-    def partition_tensor(self, param, empty_param, param_type, param_casting_dtype, to_contiguous, rank, device_mesh):
-        ep_rank = rank
-        global_num_experts = empty_param.shape[0]
-        if global_num_experts % device_mesh.size() != 0:
-            raise ValueError(
-                f"Global number of experts must be divisible by number of devices: {global_num_experts} % {device_mesh.size()} != 0"
-            )
-        local_num_experts = global_num_experts // device_mesh.size()
-        param = param[ep_rank * local_num_experts : (ep_rank + 1) * local_num_experts].to(param_casting_dtype)
+    def partition_tensor(self, param: torch.Tensor, dtype, to_contiguous: bool):
+        parameter = self.shard_tensor(param, dtype=dtype)
         if to_contiguous:
-            param = param.contiguous()
-        return param
+            parameter = parameter.contiguous()
+        return parameter
 
 
 class RouterParallel(TensorParallelLayer):
@@ -1064,10 +1021,10 @@ class RouterParallel(TensorParallelLayer):
     Allows to reshape the router scores to support running expert parallel.
     """
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, use_dtensor: bool = False, *args, **kwargs):
         super().__init__(**kwargs)
         self.args = args
-        self.use_dtensor = False
+        self.use_dtensor = use_dtensor
 
     @staticmethod
     def _prepare_input_fn(input_layouts, desired_input_layouts, mod, inputs, device_mesh):
@@ -1118,7 +1075,7 @@ class RouterParallel(TensorParallelLayer):
                 f"The number of experts must be divisible by number of ep_size: {mod.num_experts} % {ep_size} != 0"
             )
         num_local_experts = mod.num_experts // ep_size
-        router_scores, router_indices = outputs
+        router_logits, router_scores, router_indices = outputs
         router_scores = router_scores[:, ep_rank * num_local_experts : (ep_rank + 1) * num_local_experts]
         router_indices = router_indices.masked_fill((router_indices // num_local_experts) != ep_rank, -1)
         # As -1 % 1 is 0, we can only use mask fill when num_local_experts is 1
@@ -1129,28 +1086,20 @@ class RouterParallel(TensorParallelLayer):
         router_indices = router_indices.masked_fill(
             router_indices == -1, num_local_experts
         )  # masking class for one hot
-        return router_scores, router_indices
+        return router_logits, router_scores, router_indices
 
     def shard_tensor(
-        self,
-        param,
-        param_type=None,
-        param_casting_dtype=None,
-        to_contiguous=None,
-        rank=None,
-        device_mesh=None,
-        tensor_idx=None,
-    ):
-        parameter = param[...].to(param_casting_dtype)
+        self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None
+    ) -> torch.Tensor:
         self.shard = None
-        return parameter, None
+        return param[...].to(device=device, dtype=dtype)
 
-    def partition_tensor(self, param, empty_param, param_type, param_casting_dtype, to_contiguous, rank, device_mesh):
+    def partition_tensor(self, param: torch.Tensor, dtype, to_contiguous: bool):
         # TODO: i'd like for this to be the default
-        param = param[...].to(param_casting_dtype)
+        parameter = self.shard_tensor(param, dtype=dtype)
         if to_contiguous:
-            param = param.contiguous()
-        return param
+            parameter = parameter.contiguous()
+        return parameter
 
     def prepare_module_tp(self, module: nn.Module, device_mesh) -> nn.Module:
         # TODO: need an abstract Parallel class that is different from TensorParallelLayer
@@ -1169,13 +1118,13 @@ class ParallelInterface(GeneralInterface):
         {
             "colwise": ColwiseParallel(),
             "rowwise": RowwiseParallel(),
-            "colwise_rep": ColwiseParallel(output_layouts=Replicate()),
-            "rowwise_rep": RowwiseParallel(input_layouts=Replicate()),
-            "local_colwise": ColwiseParallel(use_dtensor=False),
-            "local_rowwise": RowwiseParallel(use_dtensor=False),
+            "colwise_rep": ColwiseParallelReplicate(),
+            "rowwise_rep": RowwiseParallelReplicate(),
+            "local_colwise": LocalColwiseParallel(),
+            "local_rowwise": LocalRowwiseParallel(),
             "local": IsolatedParallel(),
             "gather": GatherParallel(),
-            "local_packed_rowwise": PackedRowwiseParallel(use_dtensor=False),
+            "local_packed_rowwise": LocalPackedRowwiseParallel(),
             "sequence_parallel": SequenceParallel(),
             "replicate": ReplicateParallel(),
             "grouped_gemm": GroupedGemmParallel(),
@@ -1286,13 +1235,10 @@ def shard_and_distribute_module(
 
     if current_shard_plan is not None:
         try:
-            tp_layer = ALL_PARALLEL_STYLES[current_shard_plan]
-            tp_layer.empty_param = empty_param
-            tp_layer.device_mesh = device_mesh
-            tp_layer.rank = rank
-            param = tp_layer.partition_tensor(
-                param, empty_param, param_type, param_casting_dtype, is_contiguous, rank, device_mesh
+            tp_layer = ALL_PARALLEL_STYLES[current_shard_plan](
+                empty_param=empty_param, device_mesh=device_mesh, rank=rank
             )
+            param = tp_layer.partition_tensor(param, param_casting_dtype, is_contiguous)
         except NotImplementedError as e:
             print(
                 f"Trying to prepare {parameter_name}, but it's not supported. Corresponding module: {module_to_tp} Fix it's TP plan, current layer: {tp_layer} : {e}"