Refactor EP sharding to apply DTensor wrapping during loading

AmineDiro · AmineDiro · commit 9c712a551ba2 · 2026-04-28T08:56:01.000Z
Move EP parameter DTensor wrapping from post-load model wrapping to
the tensor parallel layer's `post_shard_wrap` method, which applies
during parameter loading. This ensures DTensor wrapping happens at the
right time in the loading pipeline and removes duplicated logic.
diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py
@@ -1077,6 +1077,8 @@ def set_param_for_module(
         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 distributed_operation is not None:
+                param_value = distributed_operation.post_shard_wrap(param_value)
             # super important otherwise _init_weight will re-init the param
             param_value._is_hf_initialized = True
             setattr(module_obj, param_name, param_value)
diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py
@@ -29,6 +29,7 @@
     import torch
     import torch.distributed as dist
     from torch import nn
+    from torch.distributed.tensor import DTensor, Shard
 
     # Cache this result has it's a C FFI call which can be pretty time-consuming
     _torch_distributed_available = torch.distributed.is_available()
@@ -130,6 +131,17 @@ def _get_parameter_tp_plan(parameter_name: str, tp_plan: dict[str, str], is_weig
     return None
 
 
+def get_ep_sharded_param_names(model) -> list[str]:
+    """FQNs of parameters whose data is per-rank unique under EP sharding."""
+    if not getattr(model, "has_ep", False):
+        return []
+    return [
+        name
+        for name, _ in model.named_parameters()
+        if _get_parameter_tp_plan(parameter_name=name, tp_plan=model.tp_plan, is_weight=True) == "grouped_gemm"
+    ]
+
+
 # =============================================================================
 # Tensor Sharding Utilities
 # =============================================================================
@@ -685,6 +697,14 @@ def update_module_attributes(self, module: nn.Module):
         """
         pass
 
+    def post_shard_wrap(self, param: nn.Parameter) -> nn.Parameter:
+        """
+        Optional final wrap applied to a parameter after `shard_tensor` and before it is
+        attached to the module. Default is identity. Subclasses can override to e.g. wrap
+        the local shard as a DTensor.
+        """
+        return param
+
 
 class ColwiseParallel(TensorParallelLayer):
     """
@@ -1078,6 +1098,15 @@ def update_module_attributes(self, module: nn.Module):
         if hasattr(module, "num_experts"):
             module.num_experts = self.get_expected_sharded_shape((self.empty_param.shape[0],))[0]
 
+    def post_shard_wrap(self, param: nn.Parameter) -> nn.Parameter:
+        """
+        Wrap the EP-sharded local tensor as a DTensor on the TP/EP mesh. Without this, the
+        optimizer's foreach ops error with "mixed Tensor and DTensor" against the
+        FSDP-wrapped DTensor params on the rest of the model.
+        """
+        dt = DTensor.from_local(param.data, self.device_mesh, [Shard(0)], run_check=False)
+        return nn.Parameter(dt, requires_grad=param.requires_grad)
+
 
 class RouterParallel(TensorParallelLayer):
     """
@@ -1487,6 +1516,8 @@ def shard_and_distribute_module(
     # otherwise loading is crazy slow
     if not isinstance(param, torch.nn.Parameter):
         param = torch.nn.Parameter(param, requires_grad=empty_param.is_floating_point())
+    if current_shard_plan is not None:
+        param = tp_layer.post_shard_wrap(param)
     setattr(module_to_tp, param_type, param)
     tp_layer.update_module_attributes(module_to_tp)
     return param
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
@@ -39,7 +39,6 @@
 from safetensors.torch import load as _safe_load_bytes
 from safetensors.torch import save_file as safe_save_file
 from torch import Tensor, nn
-from torch.distributed.tensor import DTensor, Shard
 from torch.distributions import constraints
 from torch.utils.checkpoint import checkpoint
 
@@ -1381,18 +1380,6 @@ def has_ep(self) -> bool:
         distributed_config = getattr(getattr(self, "config", None), "distributed_config", None)
         return distributed_config is not None and getattr(distributed_config, "enable_expert_parallel", False)
 
-    @property
-    def ep_sharded_param_names(self) -> list[str]:
-        """FQNs of parameters whose data is per-rank unique under EP sharding."""
-        if not self.has_ep:
-            return []
-        plan = self.tp_plan
-        return [
-            name
-            for name, _ in self.named_parameters()
-            if _get_parameter_tp_plan(parameter_name=name, tp_plan=plan, is_weight=True) == "grouped_gemm"
-        ]
-
     @property
     def tp_plan(self) -> dict[str, str]:
         """
@@ -4236,8 +4223,6 @@ def from_pretrained(
         model.eval()  # Set model in evaluation mode to deactivate Dropout modules by default
         model.set_use_kernels(use_kernels, kernel_config)
 
-        cls._wrap_ep_params_as_dtensor(model, device_mesh)
-
         # If it is a model with generation capabilities, attempt to load generation files (generation config,
         # custom generate function)
         if model.can_generate() and hasattr(model, "adjust_generation_fn") and not gguf_file:
@@ -4376,24 +4361,6 @@ def _load_pretrained_model(
 
         return loading_info, disk_offload_index
 
-    @staticmethod
-    def _wrap_ep_params_as_dtensor(model, device_mesh) -> None:
-        """Wrap EP-sharded params (`grouped_gemm` style) as DTensors in-place.
-
-        Without this, the optimizer's foreach ops error with "mixed Tensor and DTensor"
-        against the FSDP-wrapped DTensor params on the rest of the model.
-        """
-
-        if not model.has_ep:
-            return
-        plan = model.tp_plan
-        for name, p in list(model.named_parameters()):
-            if _get_parameter_tp_plan(parameter_name=name, tp_plan=plan, is_weight=True) != "grouped_gemm":
-                continue
-            parent, attr = get_module_from_name(model, name)
-            dt = DTensor.from_local(p.data, device_mesh, [Shard(0)], run_check=False)
-            setattr(parent, attr, nn.Parameter(dt, requires_grad=p.requires_grad))
-
     @staticmethod
     def _finalize_model_loading(
         model, load_config: LoadStateDictConfig, loading_info: LoadStateDictInfo
diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
@@ -70,6 +70,7 @@
 from .integrations.liger import apply_liger_kernel
 from .integrations.neftune import activate_neftune, deactivate_neftune
 from .integrations.peft import MIN_PEFT_VERSION
+from .integrations.tensor_parallel import get_ep_sharded_param_names
 from .integrations.tpu import save_tpu_checkpoint, tpu_spmd_dataloader, wrap_model_xla_fsdp
 from .modelcard import TrainingSummary
 from .modeling_utils import PreTrainedModel, unwrap_model
@@ -828,9 +829,8 @@ def create_accelerator_and_postprocess(self) -> None:
         # post accelerator creation setup
         if self.is_fsdp_enabled:
             fsdp_plugin = self.accelerator.state.fsdp_plugin
-            # EP-sharded experts must not be re-sharded by FSDP — their params are
-            # already DTensors on the EP mesh.
-            ep_param_names = getattr(self.model, "ep_sharded_param_names", []) or []
+            # EP-sharded experts must not be re-sharded by FSDP,  their params are DTensors on the EP mesh.
+            ep_param_names = get_ep_sharded_param_names(self.model)
             if ep_param_names:
                 module_names = list({n.rsplit(".", 1)[0] for n in ep_param_names})
                 fsdp_plugin.ignored_modules = [self.model.get_submodule(n) for n in module_names]
