Repository: huggingface/transformers
Pull request: [loading][TP] Fix device placement at loading-time, and simplify sharding primitives

# What does this PR do?

As per the title.
Device placement was not correct during loading. It would be loaded on cpu, and only moved thanks to the `param_value = DTensor.from_local(` call in `set_param_for_module` function that would implicitly set it to the correct device. However, this is less efficient than doing it during loading directly.

For llama3 70B and TP 8, it goes from ~9s to ~6.5s in loading time for the following code:

```python
from transformers import AutoModelForCausalLM
import torch
import time
import os

model_id = "meta-llama/Llama-3.3-70B-Instruct"

rank = int(os.environ["RANK"])
world_size = int(os.environ["WORLD_SIZE"])
device = torch.device(f"cuda:{rank}")
# Need to be initialized explicitly to use the `barrier` before loading
torch.distributed.init_process_group(backend="nccl", rank=rank, world_size=world_size, device_id=rank)

torch.cuda.synchronize(device)
torch.distributed.barrier()
t0 = time.time()
model = AutoModelForCausalLM.from_pretrained(model_id, tp_plan="auto")
torch.cuda.synchronize(device)
torch.distributed.barrier()
dt = time.time() - t0
if rank == 0:
    print(f"Took {dt:.2f} s")

torch.distributed.destroy_process_group()
```

So it's quite a big difference.

Also, there was a lot of redundancy and useless args in sharding primitive functions. This PR cleans it.

Changed files:
filename: src/transformers/core_model_loading.py
patch: @@ -712,13 +712,13 @@ def _job():
 
 
 def spawn_tp_materialize(
-    thread_pool: ThreadPoolExecutor | None, tensor: torch.Tensor, sharding_method, tensor_idx, dtype=None
+    thread_pool: ThreadPoolExecutor | None, tensor: torch.Tensor, sharding_method, tensor_idx, device=None, dtype=None
 ) -> Future | Callable:
     """Materialize and shard a tensor (according to the TP-plan) from file asynchronously if `thread_pool` is provided, or
     return a Callable that will load the tensor synchronously when called."""
 
     def _job():
-        return sharding_method.shard_tensor(tensor, param_casting_dtype=dtype, tensor_idx=tensor_idx)[0]
+        return sharding_method.shard_tensor(tensor, tensor_idx=tensor_idx, device=device, dtype=dtype)
 
     if thread_pool is not None:
         return thread_pool.submit(_job)
@@ -796,20 +796,17 @@ def set_param_for_module(
     if ref is None:
         unexpected_keys.add(target_name)
     else:
-        use_dtensor = hasattr(distributed_operation, "use_dtensor") and distributed_operation.use_dtensor
         if not isinstance(param_value, torch.nn.Parameter):
             if distributed_operation is not None:
-                param_value = DTensor.from_local(
-                    param_value,
-                    distributed_operation.device_mesh,
-                    getattr(distributed_operation, "shard", Replicate()),
-                    run_check=False,
-                    shape=ref.size(),
-                    stride=ref.stride(),
-                )
-                if not use_dtensor:
-                    # we convert to local
-                    param_value = param_value.to_local()
+                if getattr(distributed_operation, "use_dtensor", False):
+                    param_value = DTensor.from_local(
+                        param_value,
+                        distributed_operation.device_mesh,
+                        getattr(distributed_operation, "shard", Replicate()),
+                        run_check=False,
+                        shape=ref.size(),
+                        stride=ref.stride(),
+                    )
             if param_name not in module_obj._buffers:
                 param_value = torch.nn.Parameter(param_value, requires_grad=param_value.is_floating_point())
 
@@ -1079,6 +1076,7 @@ def convert_and_load_state_dict_in_model(
                         tensor,
                         mapping.distributed_operation,
                         shard_index,
+                        device_map[""],
                         _dtype,
                     )
 

filename: src/transformers/integrations/tensor_parallel.py
patch: @@ -19,6 +19,9 @@
 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,21 +30,16 @@
     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()
 
     if is_torch_greater_or_equal("2.5") and _torch_distributed_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 @@ def _prepare_input_fn(input_layouts, desired_input_layouts, mod, inputs, device_
     @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 @@ def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_me
         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 @@ def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_me
         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 @@ def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_me
         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 @@ def _prepare_input_fn(input_layouts, desired_input_layouts, mod, inputs, device_
         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,30 +696,20 @@ def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_me
 
 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())
 
 
@@ -810,45 +765,36 @@ def __init__(
         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())
 
@@ -904,30 +850,20 @@ def prepare_module_tp(self, module: nn.Module, device_mesh) -> nn.Module:
 
 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())
 
 
@@ -1015,18 +951,13 @@ def __init__(self, sequence_dim: int = 1, use_local_output: bool = False, use_dt
 
     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):
@@ -1044,16 +975,15 @@ def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_me
         )  # 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())
 
 
@@ -1067,41 +997,23 @@ def __init__(self, **kwargs):
         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):
@@ -1177,25 +1089,17 @@ def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_me
         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
@@ -1331,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}"

filename: src/transformers/modeling_utils.py
patch: @@ -4541,13 +4541,7 @@ def get_total_byte_count(
 
     total_byte_count = defaultdict(lambda: 0)
     tied_param_names = model.all_tied_weights_keys.keys()
-
-    tp_plan = getattr(model, "_tp_plan", []) or []
-    tp_plan_regex = (
-        re.compile("|".join([re.escape(plan) for plan in tp_plan]))
-        if _torch_distributed_available and torch.distributed.is_initialized()
-        else None
-    )
+    tp_plan = model._tp_plan
 
     for param_name, device in accelerator_device_map.items():
         # Skip if the parameter has already been accounted for (tied weights)
@@ -4563,9 +4557,9 @@ def get_total_byte_count(
 
         param_byte_count = param.numel() * dtype_size
 
-        if tp_plan_regex is not None:
-            generic_name = re.sub(r"\.\d+\.", ".*.", param_name)
-            param_byte_count //= torch.distributed.get_world_size() if tp_plan_regex.search(generic_name) else 1
+        if len(tp_plan) > 0:
+            is_part_of_plan = _get_parameter_tp_plan(param_name, tp_plan, is_weight=True) is not None
+            param_byte_count //= torch.distributed.get_world_size() if is_part_of_plan else 1
 
         total_byte_count[device] += param_byte_count
     return total_byte_count