🚨 TP dtensor API inference + training

src/transformers/core_model_loading.py

@@ -33,14 +33,12 @@
 from .distributed.sharding_utils import DtensorShardOperation, _dtensor_from_local_like
 from .distributed.utils import is_dtensor
 from .integrations.accelerate import get_device, offload_weight
-from .integrations.tensor_parallel import ALL_PARALLEL_STYLES
 from .utils import is_env_variable_true
 from .utils.loading_report import LoadStateDictInfo
 from .utils.logging import get_logger, tqdm
 
 
 if TYPE_CHECKING:
-    from .integrations.tensor_parallel import TensorParallelLayer
     from .modeling_utils import LoadStateDictConfig, PreTrainedModel
     from .quantizers import HfQuantizer
 
@@ -1232,9 +1230,9 @@ def spawn_materialize(
 ) -> Future | Callable:
     """Materialize (and optionally shard) a tensor, asynchronously if a thread pool is provided.
 
-    When ``sharding_op`` is given the tensor is sharded (DTensor placement or legacy TP plan);
-    otherwise it is simply copied to *device*/*dtype*. Without a thread pool a deferred
-    callable is returned instead of a Future.
+    When ``sharding_op`` is given the tensor is sharded according to its DTensor placements;
+    otherwise it is simply copied to *device*/*dtype*. Without a thread pool a deferred callable
+    is returned instead of a Future.
     """
 
     def _job():
@@ -1319,7 +1317,6 @@ def set_param_for_module(
     target_name: str,
     param_value: torch.Tensor,
     loading_info: LoadStateDictInfo,
-    distributed_operation: TensorParallelLayer | None,
     hf_quantizer: HfQuantizer,
 ):
     module_path, _, param_name = target_name.rpartition(".")
@@ -1341,26 +1338,21 @@ def set_param_for_module(
         # Remove from missing keys (it's either mismatched, or all good)
         loading_info.missing_keys.discard(target_name)
 
-        # Determine expected shape: for TP/Dtensor, use sharded shape; otherwise, use full shape
-        if distributed_operation is not None:
-            expected_shape = torch.Size(distributed_operation.get_expected_sharded_shape(ref.shape))
-        elif is_dtensor(ref):
-            expected_shape = ref._local_tensor.shape
-        else:
-            expected_shape = ref.shape
+        # For DTensor parameters, compare against the local shard loaded on this rank.
+        expected_shape = ref._local_tensor.shape if is_dtensor(ref) else ref.shape
 
         if ref is not None and param_value.shape != expected_shape and hf_quantizer is None:
             loading_info.mismatched_keys.add((target_name, param_value.shape, expected_shape))
         else:
             if is_dtensor(ref):
                 local_param = param_value.detach() if isinstance(param_value, torch.nn.Parameter) else param_value
                 dtensor_param = _dtensor_from_local_like(local_param, ref)
-                param_value = torch.nn.Parameter(dtensor_param, requires_grad=ref.requires_grad)
+                param_value = torch.nn.Parameter(
+                    dtensor_param, requires_grad=ref.requires_grad and dtensor_param.is_floating_point()
+                )
             # super important otherwise _init_weight will re-init the param
             param_value._is_hf_initialized = True
             setattr(module_obj, param_name, param_value)
-            if distributed_operation is not None:
-                distributed_operation.update_module_attributes(module_obj)
 
 
 def offload_and_maybe_resave_param(
@@ -1466,7 +1458,6 @@ def convert_and_load_state_dict_in_model(
     model: PreTrainedModel,
     state_dict: dict[str, Any],
     load_config: LoadStateDictConfig,
-    tp_plan: dict[str, str] | None,
     disk_offload_index: dict | None = None,
 ):
     r"""
@@ -1556,11 +1547,9 @@ def convert_and_load_state_dict_in_model(
 
     """
     base_model_prefix = model.base_model_prefix
-    tp_plan = tp_plan or {}
     device_map = load_config.device_map or {"": "cpu"}
     hf_quantizer = load_config.hf_quantizer
     dtype = load_config.dtype
-    device_mesh = load_config.device_mesh
     disk_offload_folder = load_config.disk_offload_folder
     offload_buffers = load_config.offload_buffers
     dtype_plan = load_config.dtype_plan or {}
@@ -1595,10 +1584,6 @@ def convert_and_load_state_dict_in_model(
     converters = [entry for entry in weight_mapping if isinstance(entry, WeightConverter)]
     param_name_to_load: dict[str, WeightRenaming | WeightConverter] = {}
 
-    # build '(?P<g0>.*.*\\.block_sparse_moe\\..*)' and group to source {'g0': '*.block_sparse_moe.'}
-    # and target to source {'g0': '*.mlp.'}. This allows us to quickly find which pattern matched.
-    if tp_plan != {}:
-        tp_plan_alt, tp_plan_by_group_name, _ = build_glob_alternation(list(tp_plan.keys()))
     if dtype_plan != {}:
         dtype_policy_alt, dtype_policy_by_group_name, _ = build_glob_alternation(list(dtype_plan.keys()))
 
@@ -1673,23 +1658,13 @@ def convert_and_load_state_dict_in_model(
                 else None
             )
 
-            # 4. Handle TP/Dtensor sharding or device_map placement
+            # 4. Handle DTensor sharding or device_map placement
             param_device = get_device(device_map, renamed_key, valid_torch_device=True)
             sharding_op = None
             materialize_device = param_device
 
             if is_dtensor(empty_param):
                 sharding_op = DtensorShardOperation(empty_param)
-            elif device_mesh and tp_plan:
-                if matched_tp_pattern := tp_plan_alt.search(renamed_key):
-                    matched_tp_pattern = tp_plan_by_group_name[matched_tp_pattern.lastgroup]
-                    if getattr(mapping, "distributed_operation", None) is None:
-                        tp_layer = ALL_PARALLEL_STYLES[model.tp_plan[matched_tp_pattern]].__class__
-                        mapping.distributed_operation = tp_layer(
-                            device_mesh=device_mesh, rank=device_mesh.get_local_rank(), empty_param=empty_param.clone()
-                        )
-                    sharding_op = mapping.distributed_operation
-                    materialize_device = device_map[""]
 
             future_or_tensor = spawn_materialize(
                 thread_pool,
@@ -1732,7 +1707,6 @@ def convert_and_load_state_dict_in_model(
                             target_name,
                             param,
                             loading_info,
-                            mapping.distributed_operation,
                             hf_quantizer,
                         )
 

src/transformers/distributed/fsdp.py

@@ -17,9 +17,9 @@
 import os
 from typing import TYPE_CHECKING, Any
 
-from ..integrations.tensor_parallel import replace_layer_number_by_wildcard
 from ..utils import is_torch_available, is_torch_distributed_available, is_torch_greater_or_equal, logging, strtobool
 from ..utils.quantization_config import QuantizationMethod
+from .tensor_parallel import replace_layer_number_by_wildcard
 from .utils import _is_torch_distributed_initialized
 
 

src/transformers/distributed/mixin.py

@@ -18,16 +18,15 @@
 import warnings
 from typing import TYPE_CHECKING
 
-from ..integrations.tensor_parallel import (
-    ALL_PARALLEL_STYLES,
-    apply_tensor_parallelism,
-    gather_state_dict_for_save,
-    initialize_tensor_parallelism,
-)
 from ..utils import is_torch_greater_or_equal, logging
 from ..utils.hub import create_and_tag_model_card
 from .configuration_utils import DistributedConfig
 from .fsdp import apply_fully_sharded_data_parallelism, is_fsdp_managed_module
+from .tensor_parallel import (
+    ALL_PARALLEL_STYLES,
+    apply_tensor_parallelism,
+    gather_state_dict_for_save,
+)
 from .utils import (
     _distributed_barrier,
     _ensure_torch_distributed,
@@ -36,6 +35,7 @@
     _is_torch_distributed_initialized,
     gather_full_state_dict,
     initialize_fully_sharded_data_parallelism,
+    initialize_tensor_parallelism,
     save_model_checkpoint_distributed,
 )
 
@@ -48,11 +48,7 @@
 
 
 class DistributedMixin:
-    """Distributed orchestration and save/load hooks for [`PreTrainedModel`].
-
-    Stateless heavy lifting stays in `transformers.distributed.*` and
-    `integrations.tensor_parallel`. This mixin owns orchestration and instance state.
-    """
+    """Distributed orchestration and save/load hooks for [`PreTrainedModel`]."""
 
     _device_mesh = None
     _tp_plan: dict[str, str] | None = None
@@ -197,15 +193,13 @@ def maybe_distribute_model(
             model._device_mesh = device_mesh
 
             if distributed_config.tp_size > 1:
-                model = apply_tensor_parallelism(
-                    model,
-                    distributed_config.tp_plan,
-                    distributed_config,
-                    device_mesh,
-                )
+                tp_mesh = device_mesh["tp"] if device_mesh.ndim > 1 else device_mesh
+                model = apply_tensor_parallelism(model, tp_mesh)
+
             elif distributed_config.fsdp_size > 1:
                 fsdp_mesh = device_mesh["fsdp"] if device_mesh.ndim > 1 else device_mesh
                 model = apply_fully_sharded_data_parallelism(model, fsdp_mesh)
+
         return model
 
     def should_save_on_this_rank(self, is_main_process: bool) -> bool:
@@ -267,7 +261,9 @@ def gather_sharded_state_dict_for_save(
             return state_dict
 
         if distributed_config.tp_size > 1:
-            state_dict = gather_state_dict_for_save(state_dict, self._tp_plan, self._device_mesh, self._tp_size)
+            state_dict = gather_state_dict_for_save(
+                state_dict, self._tp_plan, self._device_mesh, distributed_config.tp_size
+            )
             if not save_on_this_rank:
                 state_dict = {}
             return state_dict