🚨 build and expand once the tp/ep_plan/fsdp_plan

@@ -43,6 +43,10 @@ class DistributedConfig:             does not change `tp_size` or `fsdp_size`.         fsdp_size (`int`, *optional*):             Number of devices for FSDP (data parallelism). If `None` and `tp_size` is set, defaults to 1.+        fsdp_plan (`dict[str, str]`, *optional*):+            FSDP2 sharding plan. Leave as `None` to use the model's predefined `base_model_fsdp_plan` and+            class-level `_fsdp_plan`. Pass a dictionary to override individual rules; unspecified rules are kept.+            Applied only when `fsdp_size > 1`.         fsdp_cpu_offload (`bool`, *optional*, defaults to `False`):             Whether to enable CPU offloading for FSDP2.         fsdp_mixed_precision (`bool`, *optional*, defaults to `False`):@@ -65,6 +69,7 @@ class DistributedConfig:     enable_sequence_parallel: bool = False     enable_expert_parallel: bool = False     fsdp_size: int | None = None+    fsdp_plan: dict[str, str] | None = None     fsdp_cpu_offload: bool = False     fsdp_mixed_precision: bool = False     pp_size: int | None = None@@ -142,6 +147,8 @@ def _validate_mesh_config(self):         """Validate mesh sizes before the model's expert plan is available."""         if self.ep_plan is not None and not isinstance(self.ep_plan, dict):             raise ValueError("`ep_plan` must be a dictionary or None.")+        if self.fsdp_plan is not None and not isinstance(self.fsdp_plan, dict):+            raise ValueError("`fsdp_plan` must be a dictionary or None.")          if self.ep_size > 1:             if self.ep_size % self.tp_size:

@@ -22,7 +22,6 @@ from ..utils import logging from ..utils.import_utils import is_torch_available, is_torch_distributed_available from .parallel_utils import ALL_PARALLEL_STYLES, _use_local_dtensor_params-from .tensor_parallel import _get_parameter_plan   if TYPE_CHECKING:@@ -385,20 +384,22 @@ def transform_output_post_forward(self, module, output, mesh, source=None):   def apply_expert_parallelism(model: nn.Module, ep_mesh: DeviceMesh, tp_mesh: DeviceMesh, plan: dict[str, str]):-    """Shard experts on EP; use TP to split shared tokens and reconstruct outputs around dispatch."""-    for name, module in model.named_modules():-        for p_name, _ in list(module.named_parameters(recurse=False)):-            full = f"{name}.{p_name}" if name else p_name-            style_name = _get_parameter_plan(parameter_name=full, plan=plan, is_weight=True)-            if style_name is not None and style_name in ALL_PARALLEL_STYLES:-                style = ALL_PARALLEL_STYLES[style_name]-                style.validate_param(module, p_name, ep_mesh, parameter_name=full)-                style.shard_param(module, p_name, ep_mesh)--        # Dispatch hooks need both meshes to redistribute tokens between TP and EP ranks.-        style_name = _get_parameter_plan(parameter_name=name, plan=plan, is_weight=False)-        if style_name is not None and style_name in ALL_PARALLEL_STYLES:-            style = ALL_PARALLEL_STYLES[style_name]+    """+    Shard experts on EP; use TP to split shared tokens and reconstruct outputs around dispatch.+    """+    modules = dict(model.named_modules())+    parameters = dict(model.named_parameters())++    for name, style_name in plan.items():+        style = ALL_PARALLEL_STYLES[style_name]+        if name in parameters:+            module_name, _, param_name = name.rpartition(".")+            module = modules[module_name]+            style.validate_param(module, param_name, ep_mesh, parameter_name=name)+            style.shard_param(module, param_name, ep_mesh)+        else:+            module = modules[name]+            # Dispatch hooks need both meshes to redistribute tokens between TP and EP ranks.             if style_name == "ep_dispatch_experts":                 style.install_forward(module, ep_mesh=ep_mesh, tp_mesh=tp_mesh)             else:

@@ -19,15 +19,15 @@  from ..utils import is_torch_available, is_torch_distributed_available, logging, strtobool from ..utils.quantization_config import QuantizationMethod-from .tensor_parallel import _get_parameter_plan, replace_layer_number_by_wildcard-from .utils import _is_torch_distributed_initialized+from .parallel_utils import replace_layer_number_by_wildcard+from .utils import _is_torch_distributed_initialized, _warning_once_on_rank_zero   if TYPE_CHECKING:     import torch.nn as nn      from .configuration_utils import DistributedConfig-    from .utils import MeshManager+    from .utils import TransformersDeviceMesh  if is_torch_available():     import torch@@ -38,6 +38,9 @@  logger = logging.get_logger(__name__) +# FSDP plan styles: whether a module frees its full weight after forward (and re-gathers it for backward) or keeps it.+ALL_FSDP_STYLES = ("free_full_weight", "keep_full_weight")+  def is_fsdp_enabled() -> bool:     """Check if FSDP is active via Accelerate (env var based) — covers FSDP1 only."""@@ -143,29 +146,14 @@ def _resolve_tied_embed_lm_head_plan(     return adapted_plan  -def expand_fsdp_plan(-    model: nn.Module,-    fsdp_plan: dict[str, str],-) -> tuple[list[tuple[str, nn.Module]], list[tuple[str, nn.Module]]]:-    """Expand plan keys into reshard and no-reshard ``(module_name, module)`` shard targets."""-    reshard_targets: list[tuple[str, nn.Module]] = []-    no_reshard_targets: list[tuple[str, nn.Module]] = []--    for module_name, module in model.named_modules():-        plan_key = module_name if module_name in fsdp_plan else replace_layer_number_by_wildcard(module_name)-        if plan_key in fsdp_plan:-            if fsdp_plan[plan_key] == "keep_full_weight":-                no_reshard_targets.append((module_name, module))-            else:-                reshard_targets.append((module_name, module))--    return reshard_targets, no_reshard_targets-- def verify_fsdp_plan(module_names: list[str], fsdp_plan: dict[str, str] | None) -> None:     """     Verify the FSDP plan of the model, log a warning if plan keys were not applied or strategies are invalid.     """+    _warning_once_on_rank_zero(+        "`verify_fsdp_plan` is deprecated and will be removed in v5.23: the model validates its plans itself when "+        "they are set and in `_resolve_parallel_plans`."+    )     if not fsdp_plan:         return @@ -185,33 +173,26 @@ def verify_fsdp_plan(module_names: list[str], fsdp_plan: dict[str, str] | None)         logger.warning(f"The following FSDP rules were not applied to any module: {unused_rules}")  -def apply_fully_sharded_data_parallelism(model: nn.Module, mesh_manager: MeshManager) -> nn.Module:+def apply_fully_sharded_data_parallelism(+    model: nn.Module,+    mesh_manager: TransformersDeviceMesh,+    fsdp_plan: dict[str, str],+    ep_plan: dict[str, str],+) -> nn.Module:     """     Apply FSDP2 (fully_shard) to a model: dispatched experts on `efsdp` and the rest of the modules on `fsdp`.      Torch availability, distributed initialization and the version requirement     are asserted upstream by `initialize_distributed_mesh`.     """-    distributed_config = model.config.distributed_config-    fsdp_mesh = mesh_manager.get_mesh("fsdp")-    fsdp_plan = dict(getattr(model, "_fsdp_plan", None) or {})-    if not fsdp_plan:-        raise ValueError(-            f"{type(model).__name__} does not have a FSDP2 plan declared. Set "-            "`base_model_fsdp_plan` on the config and `_fsdp_plan` on the head class."-        )-     distributed_config = getattr(model.config, "distributed_config", None)+    fsdp_mesh = mesh_manager.get_mesh("fsdp")     fsdp_policy_kwargs = _get_fsdp_policy_kwargs(distributed_config) -    adapted_fsdp_plan = _resolve_tied_embed_lm_head_plan(fsdp_plan, model)-    reshard_targets, no_reshard_targets = expand_fsdp_plan(model, adapted_fsdp_plan)--    fsdp_policy_kwargs = _get_fsdp_policy_kwargs(distributed_config)-    if distributed_config.ep_size > 1 and "ep_dispatch_experts" in model.ep_plan.values():+    if distributed_config is not None and distributed_config.ep_size > 1:         expert_mesh = mesh_manager.get_mesh("efsdp")         for module_name, module in model.named_modules():-            if _get_parameter_plan(module_name, model.ep_plan, is_weight=False) == "ep_dispatch_experts":+            if ep_plan.get(module_name) == "ep_dispatch_experts":                 fully_shard(module, mesh=expert_mesh, reshard_after_forward=True, **fsdp_policy_kwargs)                 # an EP group holds ep_size / tp_size distinct batches                 # the efsdp reduce the sums on efsdp_size copies of that expert.@@ -222,18 +203,22 @@ def apply_fully_sharded_data_parallelism(model: nn.Module, mesh_manager: MeshMan                 # reject and NCCL zeroes in bfloat16 up to torch 2.13 (pytorch/pytorch#190747): sum, divide locally.                 module.set_force_sum_reduction_for_comms(True) -    for module_name, module in reshard_targets:+    reshard_targets, no_reshard_targets = [], []+    for name, style in fsdp_plan.items():+        if style == "keep_full_weight":+            no_reshard_targets.append((name, model.get_submodule(name)))+        else:+            reshard_targets.append((name, model.get_submodule(name)))++    for name, module in reshard_targets:         fully_shard(module, mesh=fsdp_mesh, reshard_after_forward=True, **fsdp_policy_kwargs)-        logger.debug(f"Applied fully_shard to {module_name} (reshard=True)")+        logger.debug(f"Applied fully_shard to {name} (reshard=True)")      # Optimization: when the keep buffer is exactly the (final_norm, lm_head/embed)     # tail pair, bundle them into one fully_shard so that we dont need to do all-gather during backward pass.     if is_norm_and_head_pair(no_reshard_targets, model):-        names, modules = [], []-        for name, module in no_reshard_targets:-            names.append(name)-            modules.append(module)-        fully_shard(modules, mesh=fsdp_mesh, reshard_after_forward=False, **fsdp_policy_kwargs)+        names, modules = zip(*no_reshard_targets)+        fully_shard(list(modules), mesh=fsdp_mesh, reshard_after_forward=False, **fsdp_policy_kwargs)         logger.debug(f"Grouped tail {names} (reshard=False)")     else:         for name, module in no_reshard_targets:

@@ -16,7 +16,7 @@ import os import re import shutil-import warnings+from itertools import chain from typing import TYPE_CHECKING  import torch@@ -32,14 +32,10 @@ ) from .configuration_utils import DistributedConfig from .expert_parallel import apply_expert_parallelism-from .fsdp import apply_fully_sharded_data_parallelism+from .fsdp import ALL_FSDP_STYLES, _resolve_tied_embed_lm_head_plan, apply_fully_sharded_data_parallelism+from .parallel_utils import ALL_PARALLEL_STYLES, _get_parameter_plan, replace_layer_number_by_wildcard from .pipeline_parallel import apply_pipeline_parallelism-from .tensor_parallel import (-    _validate_parallel_plan_styles,-    apply_tensor_parallelism,-    gather_state_dict_for_save,-    resolve_parallel_plans,-)+from .tensor_parallel import apply_tensor_parallelism, gather_state_dict_for_save from .utils import (     TransformersDeviceMesh,     _check_distributed_checkpointing_available,@@ -109,12 +105,12 @@ def init_parallel_plans(self) -> None:     @property     def tp_plan(self) -> dict[str, str]:         """The full tensor parallel plan for the model's modules."""-        return self._tp_plan if self._tp_plan is not None else {}+        return self._tp_plan      @property     def ep_plan(self) -> dict[str, str]:         """The full expert parallel plan for the model's modules, kept separate from `tp_plan`."""-        return self._ep_plan if self._ep_plan is not None else {}+        return self._ep_plan      @property     def fsdp_plan(self) -> dict[str, str]:@@ -132,22 +128,7 @@ def tp_plan(self, plan: dict[str, str] | None):         if not isinstance(plan, dict):             raise ValueError("Can only set a dictionary as `tp_plan`") -        _validate_parallel_plan_styles(plan)--        model_param_names = [name for name, _ in self.named_parameters()]-        for layer_pattern in plan.keys():-            regex_pattern = layer_pattern.replace("*", r"\d+")-            pattern_matched = False-            for param_name in model_param_names:-                if re.match(regex_pattern, param_name):-                    pattern_matched = True-                    break-            if not pattern_matched:-                warnings.warn(-                    f"Layer pattern '{layer_pattern}' does not match any parameters in the model. This rule may not "-                    "be applied during tensor parallelization, or may lead to dimension mismatches"-                )-+        self._validate_parallel_plan(plan, ALL_PARALLEL_STYLES, "tp_plan")         self._tp_plan = plan      @ep_plan.setter@@ -158,9 +139,20 @@ def ep_plan(self, plan: dict[str, str] | None):         if not isinstance(plan, dict):             raise ValueError("Can only set a dictionary as `ep_plan`") -        _validate_parallel_plan_styles(plan)+        self._validate_parallel_plan(plan, ALL_PARALLEL_STYLES, "ep_plan")         self._ep_plan = plan +    @fsdp_plan.setter+    def fsdp_plan(self, plan: dict[str, str] | None):+        if plan is None:+            self._fsdp_plan = {}+            return+        if not isinstance(plan, dict):+            raise ValueError("Can only set a dictionary as `fsdp_plan`")++        self._validate_parallel_plan(plan, ALL_FSDP_STYLES, "fsdp_plan")+        self._fsdp_plan = plan+     @pp_plan.setter     def pp_plan(self, plan: dict[str, tuple[str, str]] | None):         if plan is None:@@ -195,6 +187,101 @@ def prepare_distribute_model(          return distributed_config, device_map, mesh_manager +    def _validate_parallel_plan(+        self, plan: dict[str, str], supported_styles, plan_name: str, names: set[str] | None = None+    ) -> None:+        """Raise on a style that `supported_styles` does not know, warn on a rule that names nothing in the model."""+        unsupported_styles = {style for style in plan.values() if style not in supported_styles}+        if unsupported_styles:+            raise ValueError(+                f"Unsupported `{plan_name}` styles: {unsupported_styles}. Supported styles are {list(supported_styles)}"+            )++        if names is None:+            names = {+                replace_layer_number_by_wildcard(n) for n, _ in chain(self.named_modules(), self.named_parameters())+            }+        if unused := plan.keys() - names:+            if _get_torch_distributed_rank() == 0:+                logger.warning(+                    f"`{plan_name}` rules {sorted(unused)} match nothing in {type(self).__name__} and will not be applied."+                )++    def _expand_parallel_plan(self, plan: dict[str, str]) -> dict[str, str]:+        """Replace the `*` of the plan's rules with the layer numbers of the model, one entry per module or parameter."""+        expanded_plan = {}++        for module_name, module in self.named_modules():+            module_style = _get_parameter_plan(parameter_name=module_name, plan=plan, is_weight=False)++            if module_style in ALL_FSDP_STYLES:+                # FSDP: list the module to wrap.+                expanded_plan[module_name] = module_style+            else:+                # TP/EP: list the module's own parameters (to shard) before the module (to hook).+                for param_name, _ in module.named_parameters(recurse=False):+                    name = f"{module_name}.{param_name}" if module_name else param_name+                    param_style = _get_parameter_plan(parameter_name=name, plan=plan, is_weight=True)+                    if param_style is not None:+                        expanded_plan[name] = param_style+                if module_style is not None:+                    expanded_plan[module_name] = module_style++        return expanded_plan++    def _resolve_parallel_plans(+        self, distributed_config: DistributedConfig+    ) -> tuple[dict[str, str], dict[str, str], dict[str, str]]:+        """Build the final TP, EP and FSDP plans: apply the config overrides, then replace `*` with the layer numbers.++        Each returned plan maps one concrete module or parameter name to a style, e.g. `model.layers.3.mlp.up_proj`+        instead of `model.layers.*.mlp.up_proj`. A module that has an EP rule keeps no TP rule.+        """+        names = {replace_layer_number_by_wildcard(n) for n, _ in chain(self.named_modules(), self.named_parameters())}+        for plan_name in ("tp_plan", "ep_plan", "fsdp_plan"):+            override = getattr(distributed_config, plan_name)+            if isinstance(override, dict):+                plan = getattr(self, plan_name)+                if unknown := override.keys() - names - plan.keys():+                    raise ValueError(f"`{plan_name}` keys {sorted(unknown)} match nothing in {type(self).__name__}.")+                setattr(self, f"_{plan_name}", {} if plan_name == "tp_plan" and not override else plan | override)++        tp_plan = dict(self.tp_plan) if distributed_config.tp_size > 1 else {}+        ep_plan = dict(self.ep_plan) if distributed_config.ep_size > 1 else {}+        fsdp_plan = dict(self.fsdp_plan) if distributed_config.fsdp_size > 1 else {}++        for kind, plan in (("tp", tp_plan), ("ep", ep_plan), ("fsdp", fsdp_plan)):+            size = getattr(distributed_config, f"{kind}_size")+            if size > 1 and not plan:+                raise ValueError(+                    f"`{kind}_size={size}` was requested, but `{type(self).__name__}` has no `{kind}_plan`. Pass "+                    f"`{kind}_plan` in `DistributedConfig`, add `base_model_{kind}_plan` to the model's config, "+                    f"or set `{kind}_size=1`."+                )++        if "ep_dispatch_experts" in ep_plan.values() and "ep_router" in ep_plan.values():+            raise ValueError(+                "`ep_dispatch_experts` routes tokens itself; remove the `ep_router` rules from `ep_plan`."+            )++        # EP rules take precedence: drop TP rules on EP modules and their children.+        is_expert = re.compile(rf"(?:{'|'.join(map(re.escape, ep_plan))})(?:\..+)?").fullmatch+        tp_plan = {name: style for name, style in tp_plan.items() if not is_expert(name)}++        self._validate_parallel_plan(tp_plan, ALL_PARALLEL_STYLES, "tp_plan", names)+        self._validate_parallel_plan(ep_plan, ALL_PARALLEL_STYLES, "ep_plan", names)+        self._validate_parallel_plan(fsdp_plan, ALL_FSDP_STYLES, "fsdp_plan", names)+        distributed_config._validate_resolved_ep_plan(ep_plan)++        # Tied embed/lm_head weights are wrapped once: fold the head's FSDP rule into the embedding's.+        fsdp_plan = _resolve_tied_embed_lm_head_plan(fsdp_plan, self)++        tp_plan = self._expand_parallel_plan(tp_plan)+        ep_plan = self._expand_parallel_plan(ep_plan)+        fsdp_plan = self._expand_parallel_plan(fsdp_plan)++        return tp_plan, ep_plan, fsdp_plan+     @classmethod     def maybe_distribute_model(         cls,@@ -213,18 +300,15 @@ def maybe_distribute_model(         model._tp_size = distributed_config.tp_size         model._fsdp_size = distributed_config.fsdp_size -        # Resolve both plans before sharding anything: overrides are merged into `model.tp_plan` / `model.ep_plan`,-        # and the experts named by the EP plan are removed from the TP plan so they are sharded once.-        tp_plan, ep_plan = resolve_parallel_plans(model, distributed_config)-        distributed_config._validate_resolved_ep_plan(ep_plan)+        tp_plan, ep_plan, fsdp_plan = model._resolve_parallel_plans(distributed_config)          if distributed_config.pp_size > 1:             model = apply_pipeline_parallelism(model, mesh_manager.get_mesh("pp")) -        if tp_plan:+        if distributed_config.tp_size > 1:             model = apply_tensor_parallelism(model, mesh_manager.get_mesh("tp"), tp_plan) -        if ep_plan:+        if distributed_config.ep_size > 1:             tp_mesh = mesh_manager.get_mesh("tp")             ep_mesh = mesh_manager.get_mesh("ep") @@ -238,7 +322,7 @@ def maybe_distribute_model(                 model = apply_expert_parallelism(model, ep_mesh, tp_mesh, ep_plan)          if distributed_config.fsdp_size > 1:-            model = apply_fully_sharded_data_parallelism(model, mesh_manager)+            model = apply_fully_sharded_data_parallelism(model, mesh_manager, fsdp_plan, ep_plan)          return model 

@@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import contextlib+import re  from ..utils.generic import GeneralInterface from ..utils.import_utils import is_torch_available, is_torch_distributed_available@@ -35,6 +36,34 @@ class ParallelInterface(GeneralInterface): ALL_PARALLEL_STYLES: ParallelInterface = ParallelInterface()  +def replace_layer_number_by_wildcard(name: str) -> str:+    """+    Replace the numbers in the `name` by wildcards, only if they are in-between dots (`.`) or if they are between+    a dot (`.`) and the end of the string.+    This matches how modules are named/numbered when using a nn.ModuleList or nn.Sequential, but will NOT match+    numbers in a parameter name itself, e.g. if the param is named `"w1"` or `"w2"`.+    """+    return re.sub(r"\.\d+(\.|$)", lambda m: ".*" + m.group(1), name)+++def _get_parameter_plan(parameter_name: str, plan: dict[str, str], is_weight=True) -> str | None:+    """+    Get the parallel style for a parameter or module from a TP or EP plan.++    The plan is a dictionary that maps parameter or module names to parallel styles.+    The parameter name can be a generic name with wildcards (e.g. "*.weight") or a specific name (e.g. "layer_1.weight").++    The `is_weight` is important because for weights, we want to support `.weights` and `.bias` cases seamlessly! but+    not parent classes for `post_init` calls+    """+    generic_param_name = replace_layer_number_by_wildcard(parameter_name)+    if generic_param_name in plan:+        return plan[generic_param_name]+    elif is_weight and "." in generic_param_name and (module_name := generic_param_name.rsplit(".", 1)[0]) in plan:+        return plan[module_name]+    return None++ @contextlib.contextmanager def _use_local_dtensor_params(module):     # Kernels as DeepGEMM require local tensors rather than DTensors.

@@ -16,7 +16,6 @@ import contextlib import re import warnings-from itertools import chain from typing import TYPE_CHECKING  from ..utils import logging@@ -25,11 +24,15 @@     ALL_PARALLEL_STYLES,     ParallelInterface,  # noqa: F401     _use_local_dtensor_params,+    replace_layer_number_by_wildcard, )-from .utils import _get_torch_distributed_rank+from .utils import _get_torch_distributed_rank, _warning_once_on_rank_zero   if TYPE_CHECKING:+    from torch import nn+    from torch.distributed.device_mesh import DeviceMesh+     from ..modeling_utils import PreTrainedModel     from .configuration_utils import DistributedConfig @@ -47,16 +50,6 @@     from .parallel_utils import _AllReduceBackward, _AllReduceForward  -def replace_layer_number_by_wildcard(name: str) -> str:-    """-    Replace the numbers in the `name` by wildcards, only if they are in-between dots (`.`) or if they are between-    a dot (`.`) and the end of the string.-    This matches how modules are named/numbered when using a nn.ModuleList or nn.Sequential, but will NOT match-    numbers in a parameter name itself, e.g. if the param is named `"w1"` or `"w2"`.-    """-    return re.sub(r"\.\d+(\.|$)", lambda m: ".*" + m.group(1), name)-- def verify_tp_plan(expected_keys: list[str], tp_plan: dict[str, str] | None):     """     Verify the TP plan of the model, log a warning if the layers that were not sharded and the rules that were not applied.@@ -86,24 +79,6 @@ def verify_tp_plan(expected_keys: list[str], tp_plan: dict[str, str] | None):         logger.warning(f"The following layers were not sharded: {', '.join(unsharded_layers)}")  -def _get_parameter_plan(parameter_name: str, plan: dict[str, str], is_weight=True) -> str | None:-    """-    Get the parallel style for a parameter or module from a TP or EP plan.--    The plan is a dictionary that maps parameter or module names to parallel styles.-    The parameter name can be a generic name with wildcards (e.g. "*.weight") or a specific name (e.g. "layer_1.weight").--    The `is_weight` is important because for weights, we want to support `.weights` and `.bias` cases seamlessly! but-    not parent classes for `post_init` calls-    """-    generic_param_name = replace_layer_number_by_wildcard(parameter_name)-    if generic_param_name in plan:-        return plan[generic_param_name]-    elif is_weight and "." in generic_param_name and (module_name := generic_param_name.rsplit(".", 1)[0]) in plan:-        return plan[module_name]-    return None-- def get_kv_heads_per_rank(tp_plan: dict[str, str], num_heads: int | list[int], tp_size: int) -> int | list[int] | None:     """Returns the KV heads each rank holds under `tp_plan`, or `None` if they cannot be evenly sharded."""     # Projections whose output holds the attention keys, fused or not@@ -670,77 +645,41 @@ def transform_inputs_pre_forward(self, module, args, kwargs, mesh):     ALL_PARALLEL_STYLES.register("replicate", TensorParallelLayer())  -def _validate_parallel_plan_styles(plan: dict[str, str] | None) -> None:-    unsupported_styles = {style for style in (plan or {}).values() if style not in ALL_PARALLEL_STYLES}-    if unsupported_styles:-        raise ValueError(-            f"Unsupported parallel styles: {unsupported_styles}. Supported styles are {list(ALL_PARALLEL_STYLES.keys())}"-        )-- def resolve_parallel_plans(     model: PreTrainedModel, distributed_config: DistributedConfig ) -> tuple[dict[str, str], dict[str, str]]:-    """Merge the `DistributedConfig` overrides into the model's plans and split them between TP and EP.--    Returns the TP plan to apply to the dense modules and the EP plan to apply to the experts. Each plan is empty-    when its parallel size is 1. EP owns every module it names, so TP rules for those modules and their children-    are dropped: expert weights are sharded once, by the EP plan.     """-    # Reject invalid paths before merging, e.g. "layers.*" when the model uses "model.layers.*".-    names = {replace_layer_number_by_wildcard(n) for n, _ in chain(model.named_modules(), model.named_parameters())}-    for plan_name in ("tp_plan", "ep_plan"):-        override = getattr(distributed_config, plan_name)-        if isinstance(override, dict):-            plan = getattr(model, plan_name)-            if unknown := override.keys() - names - plan.keys():-                raise ValueError(f"`{plan_name}` keys {sorted(unknown)} match nothing in {type(model).__name__}.")-            setattr(model, f"_{plan_name}", {} if plan_name == "tp_plan" and not override else plan | override)--    tp_plan = dict(model.tp_plan) if distributed_config.tp_size > 1 else {}-    ep_plan = dict(model.ep_plan) if distributed_config.ep_size > 1 else {}-    if distributed_config.ep_size > 1 and not ep_plan:-        raise ValueError(-            f"Expert parallelism was requested (`ep_size={distributed_config.ep_size}`), but `{type(model).__name__}` "-            "does not define an expert-parallel plan. Pass `ep_plan` in `DistributedConfig`, add a "-            "`base_model_ep_plan` to the model's config, or disable expert parallelism."-        )--    if "ep_dispatch_experts" in ep_plan.values() and "ep_router" in ep_plan.values():-        raise ValueError("`ep_dispatch_experts` routes tokens itself; remove the `ep_router` rules from `ep_plan`.")--    # EP rules take precedence: drop TP rules on EP modules and their children.-    is_expert = re.compile(rf"(?:{'|'.join(map(re.escape, ep_plan))})(?:\..+)?").fullmatch-    tp_plan = {name: style for name, style in tp_plan.items() if not is_expert(name)}--    _validate_parallel_plan_styles(tp_plan)-    _validate_parallel_plan_styles(ep_plan)+    Merge the `DistributedConfig` overrides into the model's plans and split them between TP and EP.+    """+    _warning_once_on_rank_zero(+        "`resolve_parallel_plans` is deprecated and will be removed in v5.23: the model resolves its plans itself "+        "in `_resolve_parallel_plans`, which also returns the FSDP plan."+    )+    tp_plan, ep_plan, _ = model._resolve_parallel_plans(distributed_config)     return tp_plan, ep_plan  -def apply_tensor_parallelism(model, tp_mesh, tp_plan=None):-    """Apply parameter sharding and forward hooks on the TP mesh."""-    tp_plan = model.tp_plan if tp_plan is None else tp_plan-    _validate_parallel_plan_styles(tp_plan)--    for name, module in model.named_modules():-        # Create DTensor placeholders so the loader knows which shard belongs to this rank.-        for p_name, _ in list(module.named_parameters(recurse=False)):-            full = f"{name}.{p_name}" if name else p_name-            style_name = _get_parameter_plan(parameter_name=full, plan=tp_plan, is_weight=True)-            if style_name is not None and style_name in ALL_PARALLEL_STYLES:-                style = ALL_PARALLEL_STYLES[style_name]-                style.validate_param(module, p_name, tp_mesh, parameter_name=full)-                style.shard_param(module, p_name, tp_mesh)--        # Install the input/output transforms required by this module's TP style.-        style_name = _get_parameter_plan(parameter_name=name, plan=tp_plan, is_weight=False)-        if style_name is not None and style_name in ALL_PARALLEL_STYLES:+def apply_tensor_parallelism(model: nn.Module, tp_mesh: DeviceMesh, tp_plan: dict[str, str]):+    """+    Apply parameter sharding and forward hooks on the TP mesh.+    """+    modules = dict(model.named_modules())+    parameters = dict(model.named_parameters())++    for name, style_name in tp_plan.items():+        style = ALL_PARALLEL_STYLES[style_name]+        if name in parameters:+            module_name, _, param_name = name.rpartition(".")+            module = modules[module_name]+            style.validate_param(module, param_name, tp_mesh, parameter_name=name)+            style.shard_param(module, param_name, tp_mesh)+        else:+            module = modules[name]             if style_name == "mla_kv_a_proj":                 # MLA needs to know the qk_rope_head_dim to split the projection output into KV and RoPE parts.                 # TODO: Store qk_rope_head_dim on MLA projection modules when the models initialize them.                 module.config = model.config.get_text_config()-            ALL_PARALLEL_STYLES[style_name].install_forward(module, tp_mesh)+            style.install_forward(module, tp_mesh)         module._is_hooked = True      return model

@@ -65,6 +65,12 @@ def _get_torch_distributed_rank() -> int:     return torch.distributed.get_rank()  +def _warning_once_on_rank_zero(message: str) -> None:+    """Log a warning once per process, on rank 0 only."""+    if _get_torch_distributed_rank() == 0:+        logger.warning_once(message)++ def _get_torch_distributed_world_size() -> int:     if not _is_torch_distributed_initialized():         return 1

@@ -27,6 +27,7 @@     MoeTensorParalellMegaMoeExperts,     RouterParallelMegaMoe, )+from ..distributed.parallel_utils import replace_layer_number_by_wildcard from ..distributed.tensor_parallel import (     ALL_PARALLEL_STYLES,     AllReduceParallel,@@ -43,7 +44,6 @@     TensorParallelLayer,     apply_tensor_parallelism,     gather_state_dict_for_save,-    replace_layer_number_by_wildcard,     verify_tp_plan, ) from ..distributed.utils import initialize_tensor_parallelism

@@ -167,7 +167,7 @@ def update_parallel_plan(self, model, plan: dict[str, str]):         """         Updates in place `model.tp_plan` or `model.ep_plan` (the full plans, merged from the config and all         sub-configs) to take into account the quantizer's custom keys like 'inv_scale'"""-        from ..distributed.tensor_parallel import replace_layer_number_by_wildcard+        from ..distributed.parallel_utils import replace_layer_number_by_wildcard          params = {replace_layer_number_by_wildcard(name) for name, _ in model.named_parameters()}         suffixes = {n.removeprefix(p) for p, names in self.serialized_names.items() for n in names}

@@ -141,7 +141,7 @@ def _process_model_after_weight_loading(self, model, **kwargs):         return model      def update_parallel_plan(self, model, plan):-        from ..distributed.tensor_parallel import _get_parameter_plan, replace_layer_number_by_wildcard+        from ..distributed.parallel_utils import _get_parameter_plan, replace_layer_number_by_wildcard         from ..integrations.finegrained_fp8 import FP8Experts          super().update_parallel_plan(model, plan)

@@ -63,7 +63,7 @@ def _process_model_before_weight_loading(self, model, **kwargs):         model._keys_to_ignore_on_load_unexpected = ignored  # type: ignore[unresolved-attribute]      def update_parallel_plan(self, model, plan):-        from ..distributed.tensor_parallel import _get_parameter_plan, replace_layer_number_by_wildcard+        from ..distributed.parallel_utils import _get_parameter_plan, replace_layer_number_by_wildcard          super().update_parallel_plan(model, plan)         # Every rank keeps the full scalar activation scales, and the per-output-channel `weight_scale` of a rowwise

@@ -45,7 +45,7 @@         Qwen4ExpVisionConfig,         StaticCache,     )-    from transformers.distributed.fsdp import verify_fsdp_plan+    from transformers.distributed.fsdp import ALL_FSDP_STYLES     from transformers.integrations import FP8Embedding     from transformers.quantizers.quantizer_finegrained_fp8 import FineGrainedFP8HfQuantizer @@ -526,8 +526,8 @@ def test_model_forward_default_config_values(self):     def test_fsdp_plan_has_no_unused_rules(self):         with torch.device("meta"):             model = Qwen4ExpForConditionalGeneration(self.model_tester.get_config())-        with self.assertNoLogs("transformers.distributed.fsdp", level="WARNING"):-            verify_fsdp_plan([name for name, _ in model.named_modules()], model._fsdp_plan)+        with self.assertNoLogs("transformers.distributed.mixin", level="WARNING"):+            model._validate_parallel_plan(model.fsdp_plan, ALL_FSDP_STYLES, "fsdp_plan")      def test_generate_with_ple_and_inputs_embeds(self):         config = self.model_tester.get_config()

@@ -11,7 +11,6 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License.-import warnings from unittest.mock import patch  import torch@@ -77,14 +76,15 @@ def test_tp_plan_validation_invalid_style(self):         self.assertIn("Supported styles are", error_message)      def test_apply_tensor_parallelism_reports_all_invalid_styles(self):-        model = torch.nn.Module()-        model.tp_plan = {+        model_id = "hf-internal-testing/tiny-random-LlamaForCausalLM"+        model = AutoModelForCausalLM.from_pretrained(model_id, dtype="auto")+        tp_plan = {             "first_layer": "invalid_style",             "second_layer": "another_invalid_style",         }          with self.assertRaises(ValueError) as context:-            tensor_parallel.apply_tensor_parallelism(model, tp_mesh=None)+            model._validate_parallel_plan(tp_plan, ALL_PARALLEL_STYLES, "tp_plan")          error_message = str(context.exception)         self.assertIn("'invalid_style'", error_message)@@ -96,15 +96,9 @@ def test_tp_plan_validation_nonexistent_layer_warning(self):         model_id = "hf-internal-testing/tiny-random-LlamaForCausalLM"         model = AutoModelForCausalLM.from_pretrained(model_id, dtype="auto") -        # Test warning for non-existent layer pattern-        with warnings.catch_warnings(record=True) as w:-            warnings.simplefilter("always")+        with self.assertLogs("transformers.distributed.mixin", level="WARNING") as logs:             model.tp_plan = {"nonexistent.*.layer": "colwise"}--            # Check that a warning was issued-            self.assertTrue(len(w) > 0)-            warning_message = str(w[0].message)-            self.assertIn("Layer pattern 'nonexistent.*.layer' does not match any parameters", warning_message)+        self.assertIn("`tp_plan` rules ['nonexistent.*.layer'] match nothing in LlamaForCausalLM", logs.output[0])      def test_tp_plan_valid_layer_patterns(self):         """Test that valid layer patterns are accepted without warnings."""@@ -119,25 +113,9 @@ def test_tp_plan_valid_layer_patterns(self):         ]          for plan in valid_plans:-            with warnings.catch_warnings(record=True) as w:-                warnings.simplefilter("always")+            with self.assertNoLogs("transformers.distributed.mixin", level="WARNING"):                 model.tp_plan = plan -                # Filter out any warnings that are not about layer patterns-                layer_warnings = [-                    warning-                    for warning in w-                    if "Layer pattern" in str(warning.message)-                    and "does not match any parameters" in str(warning.message)-                ]--                # Should not have layer pattern warnings for valid patterns-                self.assertEqual(-                    len(layer_warnings),-                    0,-                    f"Unexpected warning for valid pattern {plan}: {[str(w.message) for w in layer_warnings]}",-                )-         # Verify the final plan was set correctly         self.assertEqual(model.tp_plan, valid_plans[-1]) 

@@ -53,7 +53,7 @@      from transformers.distributed import DistributedConfig     from transformers.distributed.checkpoint import load_optimizer_distributed, save_optimizer_distributed-    from transformers.distributed.fsdp import _resolve_tied_embed_lm_head_plan, expand_fsdp_plan+    from transformers.distributed.fsdp import _resolve_tied_embed_lm_head_plan     from transformers.distributed.utils import gather_full_state_dict  @@ -423,9 +423,8 @@ def _test_fsdp2_sharding_structure_impl(rank, config_class, config_dict, tie_wor         model = AutoModelForCausalLM.from_pretrained(init_dir, distributed_config=distributed_config)         dist.barrier() -        adapted_fsdp_plan = _resolve_tied_embed_lm_head_plan(model._fsdp_plan, model)-        reshard_targets, no_reshard_targets = expand_fsdp_plan(model, adapted_fsdp_plan)-        expected_targets = {""} | {name for name, _ in reshard_targets + no_reshard_targets}+        fsdp_plan = _resolve_tied_embed_lm_head_plan(model._fsdp_plan, model)+        expected_targets = {""} | model._expand_parallel_plan(fsdp_plan).keys()         actual_targets = {name for name, module in model.named_modules() if type(module).__name__.startswith("FSDP")}          if rank == 0:

@@ -137,7 +137,7 @@     from torch import nn      from transformers import MODEL_MAPPING-    from transformers.distributed.tensor_parallel import _get_parameter_plan+    from transformers.distributed.parallel_utils import _get_parameter_plan     from transformers.integrations.accelerate import compute_module_sizes     from transformers.modeling_utils import load_state_dict     from transformers.pytorch_utils import id_tensor_storage

@@ -21,7 +21,7 @@  from transformers import TorchAoConfig, set_seed from transformers.distributed.configuration_utils import DistributedConfig-from transformers.distributed.tensor_parallel import _get_parameter_plan+from transformers.distributed.parallel_utils import _get_parameter_plan from transformers.testing_utils import (     SINGLE_PROCESS_DEVICES,     is_tensor_parallel_test,