diff --git a/.circleci/create_circleci_config.py b/.circleci/create_circleci_config.py index 26f254d141b7..54c69567c159 100644 --- a/.circleci/create_circleci_config.py +++ b/.circleci/create_circleci_config.py @@ -398,6 +398,15 @@ def job_name(self): parallelism=6, ) +fsdp_ci_job = CircleCIJob( + "fsdp_ci", + additional_env={"RUN_FSDP_TESTS": True}, + docker_image=[{"image": "huggingface/transformers-torch-light"}], + install_steps=["uv pip install .", "uv pip install torchao"], + marker="is_fsdp_test", + parallelism=6, +) + # We also include a `dummy.py` file in the files to be doc-tested to prevent edge case failure. Otherwise, the pytest # hangs forever during test collection while showing `collecting 0 items / 21 errors`. (To see this, we have to remove # the bash output redirection.) @@ -429,7 +438,8 @@ def job_name(self): DOC_TESTS = [doc_test_job] TRAINING_CI_TESTS = [training_ci_job] TENSOR_PARALLEL_CI_TESTS = [tensor_parallel_ci_job] -ALL_TESTS = REGULAR_TESTS + EXAMPLES_TESTS + PIPELINE_TESTS + REPO_UTIL_TESTS + DOC_TESTS + [custom_tokenizers_job] + [exotic_models_job] + TRAINING_CI_TESTS + TENSOR_PARALLEL_CI_TESTS # fmt: skip +FSDP_CI_TESTS = [fsdp_ci_job] +ALL_TESTS = REGULAR_TESTS + EXAMPLES_TESTS + PIPELINE_TESTS + REPO_UTIL_TESTS + DOC_TESTS + [custom_tokenizers_job] + [exotic_models_job] + TRAINING_CI_TESTS + TENSOR_PARALLEL_CI_TESTS + FSDP_CI_TESTS # fmt: skip def create_circleci_config(folder=None): diff --git a/src/transformers/__init__.py b/src/transformers/__init__.py index e949386f2dc5..bc71ecafd36b 100755 --- a/src/transformers/__init__.py +++ b/src/transformers/__init__.py @@ -104,6 +104,7 @@ "debug_utils": [], "dependency_versions_check": [], "dependency_versions_table": [], + "distributed": [], "dynamic_module_utils": [], "feature_extraction_sequence_utils": ["SequenceFeatureExtractor"], "feature_extraction_utils": ["BatchFeature", "FeatureExtractionMixin"], diff --git a/src/transformers/configuration_utils.py b/src/transformers/configuration_utils.py index 7ba033c538d8..eb22824db8b0 100755 --- a/src/transformers/configuration_utils.py +++ b/src/transformers/configuration_utils.py @@ -147,6 +147,13 @@ class PreTrainedConfig(PushToHubMixin, RotaryEmbeddingConfigMixin): naming of attributes. - **base_model_tp_plan** (`dict[str, Any]`) -- A dict that maps sub-modules FQNs of a base model to a tensor parallel plan applied to the sub-module when `model.tensor_parallel` is called. + - **base_model_sp_plan** (`dict[str, Any]`) -- A dict that maps sub-modules FQNs of a base model to a sequence + parallel plan, used in place of `base_model_tp_plan` when `distributed_config.enable_sequence_parallel` is set. + Same key/value shape as the TP plan; values are style names registered in `ALL_PARALLEL_STYLES` + (e.g. `"vocab_reduce_scatter"`, `"rowwise_reduce_scatter"`, `"activation"`, `"module_allgather"`). + - **base_model_fsdp_plan** (`dict[Any, str]`) -- A dict that maps sub-modules of a base model to an FSDP2 + sharding strategy (e.g. `"free_full_weight"` / `"keep_full_weight"`). Keys can be wildcard module paths + (e.g. `"layers.*"`) or tuples of paths (grouped into a single `fully_shard` call). - **base_model_pp_plan** (`dict[str, tuple[list[str]]]`) -- A dict that maps child-modules of a base model to a pipeline parallel plan that enables users to place the child-module on the appropriate device. @@ -218,6 +225,8 @@ class PreTrainedConfig(PushToHubMixin, RotaryEmbeddingConfigMixin): keys_to_ignore_at_inference: ClassVar[list[str]] = [] attribute_map: ClassVar[dict[str, str]] = {} base_model_tp_plan: ClassVar[dict[str, Any] | None] = None + base_model_sp_plan: ClassVar[dict[str, Any] | None] = None + base_model_fsdp_plan: ClassVar[dict[Any, str] | None] = None base_model_pp_plan: ClassVar[dict[str, Sequence[list[str]]] | None] = None base_model_ep_plan: ClassVar[dict[str, Sequence[list[str]]] | None] = None _auto_class: ClassVar[str | None] = None @@ -1018,6 +1027,9 @@ def to_dict(self) -> dict[str, Any]: # Pop "kwargs" since they are unpacked and set in the post init output.pop("kwargs", None) + if "distributed_config" in output and hasattr(output["distributed_config"], "to_dict"): + output["distributed_config"] = output["distributed_config"].to_dict() + def to_list(value): if isinstance(value, tuple): value = [to_list(item) for item in value] @@ -1163,6 +1175,7 @@ def _remove_keys_not_serialized(self, d: dict[str, Any]) -> None: "_experts_implementation_internal", "ignore_keys_at_rope_validation", "base_model_tp_plan", + "base_model_sp_plan", "base_model_pp_plan", ]: d.pop(key_to_remove, None) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 682e8a470f13..e15c2314ac31 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -30,8 +30,8 @@ import torch +from .distributed.sharding_utils import DtensorShardOperation, _dtensor_from_local_like 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 @@ -40,10 +40,10 @@ _torch_distributed_available = torch.distributed.is_available() if TYPE_CHECKING: - from .integrations.tensor_parallel import TensorParallelLayer from .modeling_utils import LoadStateDictConfig, PreTrainedModel from .quantizers import HfQuantizer - +elif _torch_distributed_available: + from torch.distributed.tensor import DTensor logger = get_logger(__name__) @@ -384,7 +384,7 @@ def __init__(self): def _apply(self, tensor: torch.Tensor) -> torch.Tensor: dim1, dim2 = tensor.shape - n_heads = self.config.getattr("num_attention_heads", 1) + n_heads = getattr(self.config, "num_attention_heads", 1) tensor = tensor.view(n_heads, dim1 // n_heads // 2, 2, dim2) tensor = tensor.transpose(1, 2).reshape(dim1, dim2) @@ -400,11 +400,10 @@ def convert( **kwargs, ) -> dict[str, list[torch.Tensor]]: self.config = config - output: dict[str, list[torch.Tensor]] = {} + output = {} for key, tensors in input_dict.items(): - if len(tensors) != 1: - raise ValueError("PermuteForRope expects a single tensor per key.") - output[key] = [self._apply(tensors[0])] + tensor = tensors[0] if isinstance(tensors, list) else tensors + output[key] = self._apply(tensor) return output @@ -605,7 +604,7 @@ def __init__(self, source_patterns: str | list[str], target_patterns: str | list self._original_target_patterns = self.target_patterns.copy() # Init fields that will be used during conversion - self.distributed_operation: TensorParallelLayer | None = None + self.distributed_operation: Any = None self.quantization_operation: ConversionOps | None = None self.collected_tensors: dict[str, list[Future]] = defaultdict(list) self.layer_targets: dict[str, set[str]] = defaultdict(set) @@ -750,7 +749,9 @@ def reverse_transform(self) -> WeightTransform: kwargs["operations"] = [op.reverse_op for op in self.operations[::-1]] reverse_transform = self.__class__( - source_patterns=self._original_target_patterns, target_patterns=self._original_source_patterns, **kwargs + source_patterns=self._original_target_patterns, + target_patterns=self._original_source_patterns, + **kwargs, ) reverse_transform.scope_prefix = self.scope_prefix return reverse_transform @@ -775,7 +776,7 @@ def materialize_tensors(self) -> dict[str, list[torch.Tensor]]: tensors = [future.result() for future in tensors if future.result() is not None] # Sync loading elif callable(tensors[0]): - tensors = [func() for func in tensors] + tensors = [tensor for func in tensors if (tensor := func()) is not None] # Add them to the new dictionary collected_tensors[key] = tensors @@ -977,36 +978,26 @@ def spawn_materialize( tensor: torch.Tensor, device=None, dtype=None, + sharding_op: DtensorShardOperation | None = None, + tensor_idx: int | None = None, ) -> Future | Callable: - """Materialize a tensor from file asynchronously if `thread_pool` is provided, or return a Callable that will - load the tensor synchronously when called.""" + """Materialize (and optionally shard) a tensor, asynchronously if a thread pool is provided. - def _job(): - return _materialize_copy(tensor, device, dtype) - - if thread_pool is not None: - return thread_pool.submit(_job) - else: - # Return the Callable here, not the Tensor itself, so we actually delay loading to avoid saturating cpu - # memory during Conversion - return _job - - -def spawn_tp_materialize( - 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.""" + When ``sharding_op`` is given the tensor is sharded according to the DTensor + placement strategy; otherwise it is simply copied to *device*/*dtype*. + Without a thread pool a deferred callable is returned instead of a Future. + """ def _job(): - return sharding_method.shard_tensor(tensor, tensor_idx=tensor_idx, device=device, dtype=dtype) + if sharding_op is not None: + return sharding_op.shard_tensor(tensor, tensor_idx=tensor_idx, device=device, dtype=dtype) + return _materialize_copy(tensor, device, dtype) if thread_pool is not None: return thread_pool.submit(_job) - else: - # Return the Callable here, not the Tensor itself, so we actually delay loading to avoid saturating cpu - # memory during Conversion - return _job + # Return the Callable here, not the Tensor itself, so we actually delay loading + # to avoid saturating cpu memory during Conversion + return _job def dot_natural_key(s: str): @@ -1072,12 +1063,12 @@ def _format_op_name(curr_op: list[ConversionOps] | ConversionOps | None) -> str raise SkipParameters() +@torch.no_grad() def set_param_for_module( model: PreTrainedModel, target_name: str, param_value: torch.Tensor, loading_info: LoadStateDictInfo, - distributed_operation: TensorParallelLayer | None, hf_quantizer: HfQuantizer, ): module_path, _, param_name = target_name.rpartition(".") @@ -1092,27 +1083,25 @@ def set_param_for_module( if ref is None: loading_info.unexpected_keys.add(target_name) else: - if not isinstance(param_value, torch.nn.Parameter): + if not isinstance(param_value, torch.nn.Parameter) and not isinstance(ref, DTensor): if param_name not in module_obj._buffers: param_value = torch.nn.Parameter(param_value, requires_grad=param_value.is_floating_point()) # Remove from missing keys (it's either mismatched, or all good) loading_info.missing_keys.discard(target_name) - # Determine expected shape: for TP, 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)) - else: - expected_shape = ref.shape + expected_shape = ref._local_tensor.shape if isinstance(ref, DTensor) 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 isinstance(ref, DTensor): + 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) # 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( @@ -1294,10 +1283,21 @@ def convert_and_load_state_dict_in_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 + + if load_config.device_map is not None: + device_map = load_config.device_map + elif device_mesh is not None: + if device_mesh.device_type == "cpu": + device_map = {"": torch.device("cpu")} + else: + device_map = { + "": torch.device(device_mesh.device_type, getattr(torch, device_mesh.device_type).current_device()) + } + else: + device_map = {"": "cpu"} disk_offload_folder = load_config.disk_offload_folder offload_buffers = load_config.offload_buffers dtype_plan = load_config.dtype_plan or {} @@ -1332,10 +1332,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.*.*\\.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())) @@ -1399,32 +1395,23 @@ def convert_and_load_state_dict_in_model( elif empty_param is not None and empty_param.dtype != _dtype: _dtype = empty_param.dtype # usually correct when initializing - # 4. Handle TP sharding or device_map placement - future_or_tensor = None - if 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() - ) - shard_index = ( - len(mapping.collected_tensors.get(source_pattern, [])) - if isinstance(mapping, WeightConverter) and isinstance(mapping.operations[0], MergeModulelist) - else None - ) - future_or_tensor = spawn_tp_materialize( - thread_pool, - tensor, - mapping.distributed_operation, - shard_index, - device_map[""], - _dtype, - ) - - if future_or_tensor is None: - param_device = get_device(device_map, renamed_key, valid_torch_device=True) + # 4. Materialize tensor — shard-on-read for DTensor params, plain copy otherwise + param_device = get_device(device_map, renamed_key, valid_torch_device=True) + if isinstance(empty_param, DTensor): + tensor_idx = ( + len(mapping.collected_tensors.get(source_pattern, [])) + if isinstance(mapping, WeightConverter) and isinstance(mapping.operations[0], MergeModulelist) + else None + ) + future_or_tensor = spawn_materialize( + thread_pool, + tensor, + param_device, + _dtype, + sharding_op=DtensorShardOperation(empty_param), + tensor_idx=tensor_idx, + ) + else: future_or_tensor = spawn_materialize(thread_pool, tensor, param_device, _dtype) mapping.add_tensor(renamed_key, original_key, source_pattern, future_or_tensor) @@ -1454,14 +1441,7 @@ def convert_and_load_state_dict_in_model( target_name, param, loading_info, disk_offload_folder, disk_offload_index, mapping ) else: - set_param_for_module( - model, - target_name, - param, - loading_info, - mapping.distributed_operation, - hf_quantizer, - ) + set_param_for_module(model, target_name, param, loading_info, hf_quantizer) # Cleanup all the tensors that were gathered before next iteration del realized_value diff --git a/src/transformers/distributed/__init__.py b/src/transformers/distributed/__init__.py index ba6db8358d2b..09f81832c8e0 100644 --- a/src/transformers/distributed/__init__.py +++ b/src/transformers/distributed/__init__.py @@ -19,6 +19,14 @@ _import_structure = { "configuration_utils": ["DistributedConfig"], + "fsdp": ["is_fsdp_enabled", "is_fsdp_managed_module"], + "sharding_utils": [], + "tensor_parallel": [ + "ALL_PARALLEL_STYLES", + "apply_tensor_parallel", + "verify_tp_plan", + ], + "utils": [], } @@ -26,6 +34,12 @@ from .configuration_utils import ( DistributedConfig, ) + from .fsdp import is_fsdp_enabled, is_fsdp_managed_module + from .tensor_parallel import ( + ALL_PARALLEL_STYLES, + apply_tensor_parallel, + verify_tp_plan, + ) else: import sys diff --git a/src/transformers/distributed/configuration_utils.py b/src/transformers/distributed/configuration_utils.py index 7726d9f3290d..c41dd4978d50 100644 --- a/src/transformers/distributed/configuration_utils.py +++ b/src/transformers/distributed/configuration_utils.py @@ -12,99 +12,70 @@ # See the License for the specific language governing permissions and # limitations under the License. -import copy import json import os -from dataclasses import dataclass -from typing import Any +from dataclasses import asdict, dataclass + +from ..utils import is_torch_available + + +if is_torch_available(): + import torch @dataclass class DistributedConfig: """ - Base class for distributed configs + Configuration for native distributed training (FSDP2 + TP). + + Args: + tp_size (`int`, *optional*): + Number of devices for tensor parallelism. If `None` and `fsdp_size` is set, defaults to 1. + tp_plan (`dict`, *optional*): + Tensor parallel sharding plan. Leave as `None` to use the model's default plan. + fsdp_size (`int`, *optional*): + Number of devices for FSDP (data parallelism). If `None` and `tp_size` is set, defaults to 1. + fsdp_plan (`dict`, *optional*): + FSDP wrapping plan. Leave as `None` to wrap each transformer layer + root. """ + tp_size: int | None = None + tp_plan: dict[str, str] | None = None + enable_sequence_parallel: bool = False enable_expert_parallel: bool = False - # TODO: add tp_plan, pp_plan, device_mesh etc.. + fsdp_size: int | None = None + fsdp_plan: dict | None = None + + def __post_init__(self): + if self.tp_size is None and self.fsdp_size is None: + return + + if self.tp_size is None: + self.tp_size = 1 + if self.fsdp_size is None: + self.fsdp_size = 1 + + if torch.distributed.is_available() and torch.distributed.is_initialized(): + world_size = torch.distributed.get_world_size() + assert self.tp_size * self.fsdp_size == world_size, ( + f"tp_size ({self.tp_size}) * fsdp_size ({self.fsdp_size}) must be equal to world_size ({world_size})" + ) @classmethod - def from_dict(cls, config_dict, **kwargs): - """ - Constructs a DistributedConfig instance from a dictionary of parameters. - Args: - config_dict (Dict[str, Any]): Dictionary containing configuration parameters. - **kwargs: Additional keyword arguments to override dictionary values. - Returns: - DistributedConfig: Instance of DistributedConfig constructed from the dictionary. - """ - config = cls(**config_dict) - to_remove = [] - for key, value in kwargs.items(): - if hasattr(config, key): - setattr(config, key, value) - to_remove.append(key) - for key in to_remove: - kwargs.pop(key, None) - return config - - # Copied from transformers.utils.quantization_config.QuantizationConfigMixin.to_json_file + def from_dict(cls, config_dict: dict, **kwargs) -> "DistributedConfig": + merged = {**config_dict, **kwargs} + valid_keys = {f.name for f in cls.__dataclass_fields__.values()} + return cls(**{k: v for k, v in merged.items() if k in valid_keys}) + + def to_dict(self) -> dict: + return asdict(self) + + def to_json_string(self) -> str: + return json.dumps(self.to_dict(), indent=2) + "\n" + def to_json_file(self, json_file_path: str | os.PathLike): - """ - Save this instance to a JSON file. - Args: - json_file_path (`str` or `os.PathLike`): - Path to the JSON file in which this configuration instance's parameters will be saved. - use_diff (`bool`, *optional*, defaults to `True`): - If set to `True`, only the difference between the config instance and the default - `QuantizationConfig()` is serialized to JSON file. - """ - with open(json_file_path, "w", encoding="utf-8") as writer: - config_dict = self.to_dict() - json_string = json.dumps(config_dict, indent=2, sort_keys=True) + "\n" - - writer.write(json_string) - - def to_dict(self) -> dict[str, Any]: - """ - Serializes this instance to a Python dictionary. Returns: - `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance. - """ - return copy.deepcopy(self.__dict__) - - # Copied from transformers.utils.quantization_config.QuantizationConfigMixin.__iter__ - def __iter__(self): - """allows `dict(obj)` for situations where obj may be a dict or QuantizationConfigMixin""" - yield from copy.deepcopy(self.__dict__).items() - - # Copied from transformers.utils.quantization_config.QuantizationConfigMixin.__repr__ + with open(json_file_path, "w", encoding="utf-8") as f: + f.write(self.to_json_string()) + def __repr__(self): return f"{self.__class__.__name__} {self.to_json_string()}" - - def to_json_string(self): - """ - Serializes this instance to a JSON formatted string. - Returns: - str: JSON formatted string representing the configuration instance. - """ - return json.dumps(self.__dict__, indent=2) + "\n" - - def update(self, **kwargs): - """ - Updates attributes of this class instance with attributes from `kwargs` if they match existing attributes, - returning all the unused kwargs. - Args: - kwargs (`Dict[str, Any]`): - Dictionary of attributes to tentatively update this class. - Returns: - `Dict[str, Any]`: Dictionary containing all the key-value pairs that were not used to update the instance. - """ - to_remove = [] - for key, value in kwargs.items(): - if hasattr(self, key): - setattr(self, key, value) - to_remove.append(key) - - # Remove all the attributes that were updated, without modifying the input dict - unused_kwargs = {key: value for key, value in kwargs.items() if key not in to_remove} - return unused_kwargs diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py new file mode 100644 index 000000000000..d3672ce71b80 --- /dev/null +++ b/src/transformers/distributed/fsdp.py @@ -0,0 +1,487 @@ +# Copyright 2024 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +from __future__ import annotations + +import inspect +import os +from typing import TYPE_CHECKING, Any, Literal + +from ..utils import is_torch_available, is_torch_greater_or_equal, logging, strtobool +from ..utils.quantization_config import QuantizationMethod +from .tensor_parallel import replace_layer_number_by_wildcard + + +if TYPE_CHECKING: + import torch.nn as nn + +if is_torch_available(): + import torch + +if is_torch_available() and is_torch_greater_or_equal("2.5"): + import torch.distributed as dist + from torch.distributed._composable.fsdp import fully_shard + from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy, OffloadPolicy + +logger = logging.get_logger(__name__) + + +def is_fsdp_enabled() -> bool: + """Check if FSDP is active via Accelerate (env var based) — covers FSDP1 only.""" + if not is_torch_available(): + return False + + return ( + torch.distributed.is_available() + and torch.distributed.is_initialized() + and strtobool(os.environ.get("ACCELERATE_USE_FSDP", "False")) == 1 + and strtobool(os.environ.get("FSDP_CPU_RAM_EFFICIENT_LOADING", "False")) == 1 + ) + + +def is_fsdp_managed_module(module: nn.Module) -> bool: + """Check if a module is managed by FSDP (1 or 2).""" + if not is_torch_available(): + return False + if not torch.distributed.is_available(): + return False + + # FSDP2: attribute set by apply_fsdp2() + if getattr(module, "_is_fsdp_managed_module", False): + return True + # FSDP1: wrapped by FullyShardedDataParallel + try: + from torch.distributed.fsdp import FullyShardedDataParallel + except ImportError: + return False + return isinstance(module, FullyShardedDataParallel) + + +def initialize_fsdp( + fsdp_plan: dict[str, Any] | None, + device_mesh=None, + device_map=None, +): + """ + Sets up the device mesh for FSDP2 (Fully Sharded Data Parallel). + This function is called when the model is loaded and fsdp_plan is set. + + Args: + fsdp_plan: Optional FSDP config dict. Manual mode is signaled by the + presence of a ``"modules"`` key; otherwise auto mode is used. + device_mesh: Optional pre-created DeviceMesh for FSDP. + device_map: Optional device map. + + Returns: + Tuple of (device_map, device_mesh, fsdp_size) + """ + if not is_torch_available(): + raise ImportError("PyTorch is required for FSDP support") + + if fsdp_plan is None: + return device_map, device_mesh, None + + if not is_torch_greater_or_equal("2.5"): + raise OSError("FSDP2 is only supported for `torch>=2.5`.") + + if device_mesh is None: + # Detect the accelerator on the machine + device_type = torch._C._get_accelerator().type + current_device = getattr(torch, device_type) + + if not dist.is_initialized(): + try: + rank = int(os.environ["RANK"]) + local_rank = int(os.environ["LOCAL_RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + + backend_map = { + "cuda": "nccl", + "cpu": "gloo", + "xpu": "xccl", + "hpu": "hccl", + "neuron": "neuron", + "tpu": "tpu_dist", + } + backend = backend_map.get(device_type) + if device_type == "cpu" and int(os.environ.get("CCL_WORKER_COUNT", "0")): + backend = "ccl" + if device_type == "xpu" and not is_torch_greater_or_equal("2.8", accept_dev=True): + backend = "ccl" + + dist.init_process_group(backend=backend, rank=rank, world_size=world_size) + if device_type != "cpu": + current_device.set_device(local_rank) + + except Exception as e: + raise OSError( + "We tried to initialize torch.distributed for you, but it failed. Make " + "sure you init torch distributed in your script to use `fsdp_plan`." + ) from e + + if device_type != "cpu": + current_device.set_device(int(os.environ["LOCAL_RANK"])) + index = current_device.current_device() + fsdp_device = torch.device(device_type, index) + device_map = fsdp_device + else: + fsdp_device = torch.device(device_type) + device_map = device_type or {} + + fsdp_size = dist.get_world_size() + device_mesh = torch.distributed.init_device_mesh(fsdp_device.type, (fsdp_size,), mesh_dim_names=("dp_shard",)) + else: + # Use provided device mesh + if device_mesh.ndim > 1: + if "dp_shard" not in device_mesh.mesh_dim_names: + raise ValueError( + "When using `fsdp_plan` with n-d `device_mesh`, it must contain a 'dp_shard' dimension. " + "Please provide a valid `device_mesh`." + ) + device_mesh = device_mesh["dp_shard"] + fsdp_size = device_mesh.size() + device_map = torch.device(f"{device_mesh.device_type}:{int(os.environ['LOCAL_RANK'])}") + + return device_map, device_mesh, fsdp_size + + +def _get_policy_kwargs(fsdp_plan: dict[str, Any]) -> dict[str, Any]: + """Parse `cpu_offload` / `mixed_precision` flags from the user fsdp_plan into fully_shard kwargs.""" + policy_kwargs = {} + if fsdp_plan.get("cpu_offload"): + policy_kwargs["offload_policy"] = CPUOffloadPolicy() + if fsdp_plan.get("mixed_precision"): + policy_kwargs["mp_policy"] = MixedPrecisionPolicy( + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + output_dtype=None, + ) + return policy_kwargs + + +def _parse_manual_plan_entry( + entry: list[str], +) -> tuple[bool, MixedPrecisionPolicy | None, CPUOffloadPolicy | None]: + """ + Returns: + tuple[bool, MixedPrecisionPolicy | None, CPUOffloadPolicy | None]: + - bool: whether to reshard after forward + - MixedPrecisionPolicy | None: mixed precision policy + - CPUOffloadPolicy | None: cpu offload policy + """ + + if not isinstance(entry, list): + raise ValueError( + f"Manual fsdp_plan values must be a list of strings combining strategy/policies, got {type(entry)}" + ) + items = entry + + strategy: Literal["free_full_weight", "keep_full_weight"] | None = None + offload_policy: CPUOffloadPolicy | None = None + mp_policy: MixedPrecisionPolicy | None = None + + for item in items: + if not isinstance(item, str): + raise ValueError( + f"fsdp_plan option must be a string, got {type(item)}. " + "Supported: 'free_full_weight', 'keep_full_weight', 'cpu_offload', 'mixed_precision'." + ) + token = item.lower() + if token in {"free_full_weight", "keep_full_weight"}: + strategy = token + elif token == "cpu_offload": + offload_policy = CPUOffloadPolicy() + elif token == "mixed_precision": + # TODO(3outeille): add support for different dtypes + mp_policy = MixedPrecisionPolicy( + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + output_dtype=torch.bfloat16, + ) + else: + raise ValueError( + "Unknown fsdp_plan option " + f"{item!r}. Supported: 'free_full_weight', 'keep_full_weight', 'cpu_offload', 'mixed_precision'." + ) + + if strategy is None: + strategy = "free_full_weight" + + return strategy != "keep_full_weight", mp_policy, offload_policy + + +def _iter_manual_plan_targets(model, pattern, name_to_module, already_sharded_names): + if pattern in name_to_module: + target = name_to_module[pattern] + if isinstance(target, (torch.nn.ModuleList, torch.nn.ModuleDict, torch.nn.Sequential)): + # (ModuleList, ModuleDict, Sequential) don't have a forward() that gets called - + # the model loops over their children directly. So when a pattern matches a + # container, we shard each child instead. + for child_name, child in target.named_children(): + yield f"{pattern}.{child_name}", child + else: + yield pattern, target + return + + # Prefix match: "model.layers" matches "model.layers.0", etc. + for name, module in model.named_modules(): + if name in already_sharded_names or isinstance( + module, (torch.nn.ModuleList, torch.nn.ModuleDict, torch.nn.Sequential) + ): + continue + if name != pattern and not name.startswith(pattern + "."): + continue + if any( + name.startswith(already_sharded_names_name + ".") for already_sharded_names_name in already_sharded_names + ): + continue + yield name, module + + +def _get_manual_plan_modules(fsdp_plan: dict[str, Any]) -> dict[str, list[str]]: + modules = fsdp_plan.get("modules") + if not isinstance(modules, dict): + raise ValueError("Manual fsdp_plan must define a 'modules' dict.") + return modules + + +def is_tail_pair(entries) -> bool: + """Match the canonical tail pair: one final norm + the output head (or tied embedding).""" + if len(entries) != 2: + return False + names = [name for name, _ in entries] + has_norm = any(n == "norm" or n.endswith(".norm") for n in names) + has_head = any(n in {"lm_head", "embed_tokens"} or n.endswith((".lm_head", ".embed_tokens")) for n in names) + return has_norm and has_head + + +def tied_source_path(model) -> str | None: + """Return the dotted path of the input embedding module (the tied source).""" + input_embed = getattr(model, "get_input_embeddings", lambda: None)() + if input_embed is None: + return None + for name, module in model.named_modules(): + if module is input_embed: + return name + return None + + +def _resolve_plan_key(name_to_module: dict, key: str): + """Resolve a plan key into the matching (name, module) pairs. + + Supports exact module names and tp_plan-style wildcards (via + ``replace_layer_number_by_wildcard``). + """ + if key in name_to_module: + return [(key, name_to_module[key])] + return [(name, mod) for name, mod in name_to_module.items() if replace_layer_number_by_wildcard(name) == key] + + +def _iter_plan_targets(model, plan, is_weights_tied: bool, tied_source: str | None): + """Yield ``(name, module, strategy)`` for every module the plan applies to. + + Expands wildcards via ``_resolve_plan_key`` and pre-applies tying rules: + skips the standalone tied-source entry, and rewrites a keep ``"lm_head"`` + entry to the tied source so the shared parameter is wrapped once. + """ + name_to_module = dict(model.named_modules()) + for key, strategy in plan.items(): + if is_weights_tied and key == tied_source: + continue + if is_weights_tied and key == "lm_head" and strategy == "keep_full_weight" and tied_source is not None: + yield tied_source, name_to_module[tied_source], strategy + continue + for name, module in _resolve_plan_key(name_to_module, key): + yield name, module, strategy + + +def apply_fully_shard_data_parallel( + model, + fsdp_mesh, + fsdp_plan: dict[str, Any] | None, +): + """ + Apply FSDP2 (fully_shard) to a model. + + When ``fsdp_plan`` is ``None`` or doesn't contain a ``"modules"`` key, the + model-declared ``model._fsdp_plan`` drives sharding. Policies (`cpu_offload`, + `mixed_precision`) from ``fsdp_plan`` are applied on top. + + When ``fsdp_plan`` has a ``"modules"`` key, the user fully specifies the + layout (manual mode). + + Examples: + # Plan-driven (uses model._fsdp_plan). + fsdp_plan = None + fsdp_plan = {"cpu_offload": True, "mixed_precision": True} + + # Manual override. + fsdp_plan = { + "modules": { + "model.embed_tokens": ["free_full_weight"], + "model.layers.0.mlp": ["free_full_weight", "cpu_offload", "mixed_precision"], + "model.norm": ["keep_full_weight"], + "lm_head": ["keep_full_weight"], + }, + } + """ + if not is_torch_available(): + raise ImportError("PyTorch is required for FSDP support") + + if not is_torch_greater_or_equal("2.5"): + raise OSError("FSDP2 requires torch>=2.5") + + if fsdp_plan is None: + fsdp_plan = {} + + input_embed = getattr(model, "get_input_embeddings", lambda: None)() + output_embed = getattr(model, "get_output_embeddings", lambda: None)() + is_weights_tied = ( + input_embed is not None + and output_embed is not None + and hasattr(input_embed, "weight") + and hasattr(output_embed, "weight") + and input_embed.weight is output_embed.weight + ) + + if not isinstance(fsdp_plan, dict): + raise ValueError(f"fsdp_plan must be a dict, got {type(fsdp_plan)}") + is_manual = "modules" in fsdp_plan + + if not is_manual: + policy_kwargs = _get_policy_kwargs(fsdp_plan) + + plan = getattr(model, "_fsdp_plan", None) or {} + if not plan: + raise ValueError( + f"{type(model).__name__} has no `_fsdp_plan` declared. Either set " + "`base_model_fsdp_plan` on the config and `_fsdp_plan` on the head class, " + "or pass an explicit `fsdp_plan={'modules': {...}}` manual override." + ) + + tied_source = tied_source_path(model) if is_weights_tied else None + keep_buffer: list[tuple[str, Any]] = [] + + for name, module, strategy in _iter_plan_targets(model, plan, is_weights_tied, tied_source): + if strategy == "keep_full_weight": + keep_buffer.append((name, module)) + continue + fully_shard(module, mesh=fsdp_mesh, reshard_after_forward=True, **policy_kwargs) + 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_tail_pair(keep_buffer): + keep_names = [n for n, _ in keep_buffer] + keep_modules = [m for _, m in keep_buffer] + fully_shard(keep_modules, mesh=fsdp_mesh, reshard_after_forward=False, **policy_kwargs) + logger.debug(f"Grouped tail {keep_names} (reshard=False)") + else: + for name, module in keep_buffer: + fully_shard(module, mesh=fsdp_mesh, reshard_after_forward=False, **policy_kwargs) + logger.debug(f"Applied fully_shard to {name} (reshard=False)") + + # Shard root model + fully_shard(model, mesh=fsdp_mesh, **policy_kwargs) + + logger.info(f"FSDP2 applied to model via _fsdp_plan: {len(plan)} entries") + + else: + # fsdp_plan = { + # "modules": { + # "model.layers.0.self_attn": ["free_full_weight"], # reshard_after_forward=True + # "model.norm": ["keep_full_weight"], # reshard_after_forward=False + # "model.layers.0.mlp": ["free_full_weight", "cpu_offload", "mixed_precision"], + # }, + # } + + name_to_module = dict(model.named_modules()) + already_sharded_names: set[str] = set() + root_mp_policy = MixedPrecisionPolicy() + root_offload_policy = OffloadPolicy() + + for pattern, entry in _get_manual_plan_modules(fsdp_plan).items(): + reshard, mp_policy, offload_policy = _parse_manual_plan_entry(entry) + if mp_policy is not None: + root_mp_policy = mp_policy + if offload_policy is not None: + root_offload_policy = offload_policy + + for name, module in _iter_manual_plan_targets(model, pattern, name_to_module, already_sharded_names): + if name in already_sharded_names: + continue + shard_kwargs = {"mesh": fsdp_mesh, "reshard_after_forward": reshard} + if mp_policy is not None: + shard_kwargs["mp_policy"] = mp_policy + if offload_policy is not None: + shard_kwargs["offload_policy"] = offload_policy + fully_shard(module, **shard_kwargs) + already_sharded_names.add(name) + logger.debug(f"Applied fully_shard to {name}") + + # Shard root model with the same policies as sub-modules. + # MixedPrecisionPolicy.output_dtype casting happens in post_forward + # for every fully_shard-wrapped module, even with no direct parameters. + fully_shard(model, mesh=fsdp_mesh, mp_policy=root_mp_policy, offload_policy=root_offload_policy) + + # Used by generation code to detect FSDP and enable synced_gpus. + model._is_fsdp_managed_module = True + + if is_weights_tied and hasattr(model, "tie_weights"): + # Re-tie weights. + # fully_shard replaces nn.Parameter objects (swapping data for DTensor shards), + # which breaks weight tying (e.g. lm_head.weight is no longer embed_tokens.weight). + # Re-tying makes lm_head._parameters["weight"] point to the new DTensor parameter + # so gradients accumulate correctly into a single buffer. + model.tie_weights() + + return model + + +# ========================= PEFT compatibility ========================= +# TODO(3outeille): make sure new FSDP works with PEFT +def get_fsdp_ckpt_kwargs(): + """ + Returns checkpoint kwargs for FSDP model saving. + + Checks if the `adapter_only` parameter is supported by `save_fsdp_model` from accelerate + and returns the appropriate kwargs. + """ + from accelerate.utils import save_fsdp_model + + if "adapter_only" in list(inspect.signature(save_fsdp_model).parameters): + return {"adapter_only": True} + else: + return {} + + +def update_fsdp_plugin_peft(model, accelerator): + """ + Updates the FSDP plugin for PEFT LoRA/QLoRA compatibility. + + When using FSDP with PEFT LoRA, the auto wrap policy needs to be updated to additionally wrap + LoRA trainable layers separately. When using FSDP with QLoRA, the mixed precision policy needs + to be updated to use the quantization storage data type. + """ + from peft import PeftConfig + from peft.utils.other import fsdp_auto_wrap_policy + + if isinstance(model.active_peft_config, PeftConfig): + accelerator.state.fsdp_plugin.auto_wrap_policy = fsdp_auto_wrap_policy(model) + if ( + getattr(model, "quantization_method", None) == QuantizationMethod.BITS_AND_BYTES + and model.hf_quantizer.quantization_config.bnb_4bit_quant_storage.is_floating_point + ): + accelerator.state.fsdp_plugin.set_mixed_precision( + model.hf_quantizer.quantization_config.bnb_4bit_quant_storage, override=True + ) diff --git a/src/transformers/distributed/sharding_utils.py b/src/transformers/distributed/sharding_utils.py new file mode 100644 index 000000000000..097e30216420 --- /dev/null +++ b/src/transformers/distributed/sharding_utils.py @@ -0,0 +1,491 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +from __future__ import annotations + +import math +from typing import TYPE_CHECKING + +from ..utils import is_torch_available + + +if TYPE_CHECKING: + import torch + from torch.distributed.tensor import DTensor + +if is_torch_available(): + import torch + from torch.distributed._functional_collectives import wait_tensor + from torch.distributed.tensor import DTensor, Replicate + from torch.distributed.tensor._utils import compute_local_shape_and_global_offset + from torch.distributed.tensor.placement_types import Shard, _StridedShard + + +class DtensorShardOperation: + """Shard-on-read: slice a full checkpoint tensor down to this rank's local + DTensor shard, for any combination of placements on a 1-D or n-D mesh. + + Placements primer + ----------------- + Each mesh dim carries one placement describing how it slices the tensor: + + | Placement | Local data on each rank of the mesh dim | + |--------------------------|-------------------------------------------------| + | Replicate | full tensor (no slicing) | + | Shard(d) | contiguous chunk of dim d (rows r*c .. (r+1)*c) | + | _StridedShard(d, sf=N) | one chunk from each of N groups along dim d, | + | | concatenated together (interleaved layout) | + + Different scenarios of different placements + ------------------------------------------ + Placement tuples are ordered outermost-first; for a 2-D (fsdp, tp) mesh + the tuple is (fsdp_placement, tp_placement). + + | Scenario | Placements | + |---------------------------------------------------|---------------------------------------| + | TP-only, non-fused (e.g. q_proj/k_proj/v_proj) | [Shard(d)] | + | TP-only, fused gate/up | [_StridedShard(d, sf=2)] | + | TP + FSDP, same tensor dim | [_StridedShard(d, sf=tp_size), Shard(d)] | + | TP + FSDP, different dims | [Shard(d1), Shard(d2)] | + + Mesh dimensions are listed outermost-first. For a 2-D (fsdp=F, tp=T) mesh, + rank index = fsdp_idx * T + tp_idx. + + Loading (this class) + -------------------- + During `from_pretrained`, each rank reads the full checkpoint tensors, + then calls `shard_tensor(source, tensor_idx=...)` to keep only its local + DTensor shard. The class encapsulates the placements + mesh so the + slicing logic doesn't have to be repeated at every call site. + + Two checkpoint layouts are supported. Running example: an MoE weight + stack of param shape [N_experts, in, out]: + + | Checkpoint layout | tensor_idx | source.shape | Returns | + |------------------------------------|------------|---------------|-----------------------------| + | One stacked tensor | None | [N, in, out] | This rank's slice along | + | | | | every sharded dim. | + | N per-expert tensors (called once | 0..N-1 | [in, out] | Inner-dim slice if this | + | per expert by the caller) | | | rank owns expert `i`, | + | | | | else None (caller drops it).| + """ + + def __init__(self, param: DTensor): + self.device_mesh = param.device_mesh + self.placements = tuple(param.placements) + self.param_ndim = param.ndim + local_shape, offsets = compute_local_shape_and_global_offset(param.shape, self.device_mesh, self.placements) + # Where this rank's slice starts along axis 0, and how many indices + # it covers. Example: param of shape [8, in, out] with Shard(0) on + # 2 ranks gives: + # rank 0 → _axis0_offset=0, _axis0_local_size=4 (owns experts 0..3) + # rank 1 → _axis0_offset=4, _axis0_local_size=4 (owns experts 4..7) + # When the checkpoint stores one tensor per expert, shard_tensor + # checks whether tensor_idx falls in this rank's range to decide + # whether to keep the piece or drop it. + self._axis0_offset = offsets[0] + self._axis0_local_size = local_shape[0] + + def shard_tensor( + self, source: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None + ) -> torch.Tensor | None: + """Slice source down to this rank's shard. + + The checkpoint can store the parameter in two layouts. Take a stack + of N MoE experts of shape [in, out] as a running example — the + param shape is [N, in, out]: + + - Single tensor (tensor_idx is None): the checkpoint holds one + [N, in, out] tensor, so source.shape == param.shape. Every + sharded dim is sliced here, including axis 0. + + - One tensor per piece (tensor_idx given): the checkpoint holds N + separate [in, out] tensors, one per expert. shard_tensor is + called once per expert; on each call source is the [in, out] + tensor for expert number tensor_idx (so 0 <= tensor_idx < N). + Note: source has one fewer dim than the param: the axis-0 + index lives in tensor_idx, not in source.shape. + If this rank does not own tensor_idx along axis 0, return None + and the piece is discarded. Otherwise slice only the inner + dims; the caller (MergeModulelist / Concatenate) collects the + kept pieces and stacks them back along axis 0 to rebuild the + full [N, in, out] param. + """ + source_shape = list(source.shape) if isinstance(source, torch.Tensor) else source.get_shape() + placements = [(md, p) for md, p in enumerate(self.placements) if hasattr(p, "dim")] + + # Dense path + if tensor_idx is None: + if not placements: + return source[...].to(device=device, dtype=dtype) + has_strided = any(not p.is_shard() for _, p in placements) + intervals = [[(0, size)] for size in source_shape] + for mesh_dim, placement in placements: # [i.e: (0, Shard(0)), (1, Shard(-1))] + sub_mesh = self._get_sub_mesh(mesh_dim) + rank, world_size = sub_mesh.get_local_rank(), sub_mesh.size() + source_dim = self._norm_dim(placement.dim) + if not placement.is_shard(): + intervals[source_dim] = self._strided_intervals( + intervals[source_dim], rank, world_size, placement.split_factor + ) + else: + intervals[source_dim] = self._contiguous_intervals(intervals[source_dim], rank, world_size) + # Only _StridedShard can produce multi-interval dims that need cat. + if has_strided: + return self._slice_and_cat(source, intervals, device, dtype) + slices = tuple(slice(*(pieces[0] if pieces else (0, 0))) for pieces in intervals) + return source[slices].to(device=device, dtype=dtype) + + # MoE path: drop the piece if this rank does not own tensor_idx + # along axis 0. Once shard_tensor has been called for all N pieces, + # the caller (MergeModulelist) stacks the kept slices along axis 0 to + # form this rank's local shard of the param. + shards_leading_axis = any(self._norm_dim(p.dim) == 0 for _, p in placements) + owns_index = self._axis0_offset <= tensor_idx < self._axis0_offset + self._axis0_local_size + if shards_leading_axis and not owns_index: + return None + + # Inner dims use only _contiguous_intervals (one piece per dim), so a + # single slice suffices + inner_placements = [(md, p) for md, p in placements if self._norm_dim(p.dim) != 0] + if not inner_placements: + return source[...].to(device=device, dtype=dtype) + slice_per_dim: list[tuple[int, int]] = [(0, size) for size in source_shape] + # placement.dim is indexed in param space (e.g. axis 2 of [N, in, out]). + # source is in source space (e.g. axis 1 of [in, out]), so we translate + # from one to the other by stripping the leading axis. + for mesh_dim, placement in inner_placements: + sub_mesh = self._get_sub_mesh(mesh_dim) + rank, world_size = sub_mesh.get_local_rank(), sub_mesh.size() + param_dim = self._norm_dim(placement.dim) + source_dim = param_dim - 1 + pieces = self._contiguous_intervals([slice_per_dim[source_dim]], rank, world_size) + slice_per_dim[source_dim] = pieces[0] if pieces else (0, 0) + return source[tuple(slice(s, e) for s, e in slice_per_dim)].to(device=device, dtype=dtype) + + def _strided_intervals( + self, intervals: list[tuple[int, int]], rank: int, world_size: int, split_factor: int + ) -> list[tuple[int, int]]: + narrowed = [] + for start, end in intervals: + group_size = math.ceil((end - start) / split_factor) + for group_idx in range(split_factor): + group_start = start + group_idx * group_size + group_end = min(group_start + group_size, end) + if group_end <= group_start: + continue + size, offset = Shard.local_shard_size_and_offset(group_end - group_start, world_size, rank) + if size > 0: + narrowed.append((group_start + offset, group_start + offset + size)) + return narrowed + + def _contiguous_intervals( + self, intervals: list[tuple[int, int]], rank: int, world_size: int + ) -> list[tuple[int, int]]: + total = sum(end - start for start, end in intervals) + my_size, my_offset = Shard.local_shard_size_and_offset(total, world_size, rank) + if my_size == 0: + return [] + + out: list[tuple[int, int]] = [] + flat_pos = 0 + slice_end = my_offset + my_size + for start, end in intervals: + length = end - start + interval_end_flat = flat_pos + length + if interval_end_flat <= my_offset: # entirely before my slice + flat_pos = interval_end_flat + continue + if flat_pos >= slice_end: # entirely after my slice + break + sub_start = max(0, my_offset - flat_pos) + sub_end = min(length, slice_end - flat_pos) + out.append((start + sub_start, start + sub_end)) + flat_pos = interval_end_flat + return out + + def _slice_and_cat(self, source, intervals, device, dtype): + multi_interval_dim: int | None = None + slices: list[slice] = [] + for source_dim, pieces in enumerate(intervals): + if len(pieces) == 1: + start, end = pieces[0] + slices.append(slice(start, end)) + continue + if multi_interval_dim is not None: + raise ValueError("Shard-on-read only supports disjoint ranges on a single checkpoint dimension.") + multi_interval_dim = source_dim + slices.append(slice(None)) # placeholder, filled per-piece below + + # Fast path: every dim is one contiguous interval, read in a single slice. + if multi_interval_dim is None: + return source[tuple(slices)].to(device=device, dtype=dtype) + + # Multi-interval dim: read each piece separately, then concatenate. + pieces_read = [] + for start, end in intervals[multi_interval_dim]: + piece_slices = list(slices) + piece_slices[multi_interval_dim] = slice(start, end) + pieces_read.append(source[tuple(piece_slices)]) + return torch.cat(pieces_read, dim=multi_interval_dim).to(device=device, dtype=dtype) + + def _get_sub_mesh(self, mesh_dim: int): + if self.device_mesh.ndim == 1: + return self.device_mesh + return self.device_mesh[self.device_mesh.mesh_dim_names[mesh_dim]] + + def _norm_dim(self, dim: int) -> int: + # if dim is negative, it should be normalized to the last axis + return dim if dim >= 0 else self.param_ndim + dim + + +def _dtensor_from_local_like(local_tensor: torch.Tensor, ref: DTensor) -> DTensor: + """Wrap `local_tensor` as a DTensor that mirrors `ref`'s mesh, placements, + global shape, and stride.""" + return DTensor.from_local( + local_tensor.contiguous(), + ref.device_mesh, + ref.placements, + run_check=False, + shape=ref.shape, + stride=tuple(ref.stride()), + ) + + +def _replicate_dtensor(tensor: DTensor) -> DTensor: + """All-gather a DTensor to fully Replicate, handling _StridedShard. + + PyTorch's redistribute() does not support _StridedShard as a source: + _StridedShard -> redistribute() -> Replicate ❌ AssertionError + _StridedShard -> redistribute() -> Shard ❌ NotImplementedError + Shard -> redistribute() -> Replicate ✅ works + Replicate -> redistribute() -> Shard ✅ works + Replicate -> redistribute() -> _StridedShard ✅ works + + During reconstruction, we walk placements right-to-left (innermost mesh dim first), + invoke each one's low-level _to_replicate_tensor, and wait for the async collective to finish + at each step. Example — global [256, 1024], mesh (fsdp=2, tp=2), + placements (_StridedShard(0), Shard(0)), local [64, 1024]: + + i=1 (tp, Shard(0)): [64, 1024] -> [128, 1024] + i=0 (fsdp, _StridedShard(0)): [128, 1024] -> [256, 1024] + """ + mesh = tensor.device_mesh + placements = tensor.placements + replicate_all = tuple(Replicate() for _ in range(mesh.ndim)) + + if not any(isinstance(p, _StridedShard) for p in placements): + return tensor.redistribute(placements=replicate_all) + + with torch.no_grad(): + local = tensor._local_tensor + for i in reversed(range(mesh.ndim)): + p = placements[i] + if p.is_replicate(): + continue + logical_shape = list(tensor.shape) + for j, pj in enumerate(placements[:i]): + if not pj.is_replicate(): + size, _ = Shard.local_shard_size_and_offset( + logical_shape[pj.dim], mesh.size(j), mesh.get_local_rank(j) + ) + logical_shape[pj.dim] = size + local = p._to_replicate_tensor(local, mesh, i, logical_shape) + local = wait_tensor(local) + return DTensor.from_local(local, mesh, replicate_all, run_check=False) + + +def _find_strided_shard_placement_from_fused_params(placements): + for i, p in enumerate(placements): + if not isinstance(p, _StridedShard): + continue + # We want to find the first _StridedShard placement that is not composed with another placement on the same dim. + # Because that means it's a fused parameter. Meanwhile if it's acting on the same dim, that means it comes from composing parallelism together + has_partner_on_same_dim = any( + j != i and getattr(other, "dim", None) == p.dim for j, other in enumerate(placements) + ) + if not has_partner_on_same_dim: + return p + return None + + +def _split_fused_dtensor(dt, dim, n_pieces): + local_pieces = list(dt._local_tensor.chunk(n_pieces, dim=dim)) + if len(local_pieces) != n_pieces: + raise RuntimeError( + f"Cannot split DTensor of shape {tuple(dt.shape)} into {n_pieces} pieces along dim {dim}: " + f"got {len(local_pieces)} chunks instead." + ) + new_placements = tuple( + Shard(p.dim) if (isinstance(p, _StridedShard) and p.dim == dim) else p for p in dt.placements + ) + return [ + DTensor.from_local(lp.contiguous(), dt.device_mesh, new_placements, run_check=False) for lp in local_pieces + ] + + +def _merge_unfused_dtensors(pieces, dim, target_placements): + local_cat = torch.cat([p._local_tensor for p in pieces], dim=dim).contiguous() + return DTensor.from_local(local_cat, pieces[0].device_mesh, target_placements, run_check=False) + + +def get_fusion_metadata(optimizer_state_dict): + """Inspect the optimizer state dict and return metadata describing every + fused param that needs to be split for DCP. Does NOT mutate + `optimizer_state_dict`. Pair it with `unfuse_optimizer_state` to apply + the split and `fuse_optimizer_state` to undo it on load. + + Args: + optimizer_state_dict: as returned by `get_optimizer_state_dict(...)`. For + a Mixtral on a (fsdp=2, tp=2) mesh with AdamW and a single MoE layer, + it looks like: + + { + "state": { + "model.layers.0.mlp.experts.gate_up_proj": { + "exp_avg": DTensor(shape=(4, 16, 8), + placements=(Shard(0), _StridedShard(1, sf=2))), + "exp_avg_sq": DTensor(shape=(4, 16, 8), + placements=(Shard(0), _StridedShard(1, sf=2))), + "step": tensor(7.0), + }, + "model.layers.0.mlp.experts.down_proj": { + "exp_avg": DTensor(shape=(4, 8, 16), + placements=(Shard(0), Shard(2))), + "step": tensor(7.0), + }, + }, + "param_groups": [{"lr": 1e-4, ...}], + } + + Returns: + fusion_metadata = { + "model.layers.0.mlp.experts.gate_up_proj": { + "chunk_dim": 1, + "placements": (Shard(0), _StridedShard(1, split_factor=2)), + "unfused_keys": [ + "model.layers.0.mlp.experts.gate_up_proj.0", + "model.layers.0.mlp.experts.gate_up_proj.1", + ], + }, + } + """ + optimizer_state = optimizer_state_dict.get("state", {}) + fusion_metadata: dict[str, dict] = {} + + for param_name, optimizer_fields in optimizer_state.items(): + dtensor = next((v for v in optimizer_fields.values() if isinstance(v, DTensor)), None) + if dtensor is None: + continue + strided_shard_placement = _find_strided_shard_placement_from_fused_params(dtensor.placements) + if strided_shard_placement is None: + continue + fusion_metadata[param_name] = { + "chunk_dim": strided_shard_placement.dim, + "placements": tuple(dtensor.placements), + "unfused_keys": [f"{param_name}.{i}" for i in range(strided_shard_placement.split_factor)], + } + + return fusion_metadata + + +def unfuse_optimizer_state(optimizer_state_dict, fusion_metadata): + """Apply `fusion_metadata` to split each fused param into its `sf` plain-Shard + pieces in place. After the call, `optimizer_state_dict["state"]` has + `.0` .. `.{sf-1}` keys in place of each original fused key; + non-fused params (absent from `fusion_metadata`) are untouched. + + Args: + optimizer_state_dict: same shape as the input to `get_fusion_metadata`. + fusion_metadata: as returned by `get_fusion_metadata(optimizer_state_dict)`. + """ + optimizer_state = optimizer_state_dict.get("state", {}) + + for param_name, metadata in fusion_metadata.items(): + optimizer_fields = optimizer_state[param_name] + + for unfused_key in metadata["unfused_keys"]: + optimizer_state[unfused_key] = {} + + for optim_param_name, value in optimizer_fields.items(): + if isinstance(value, DTensor): + chunks = _split_fused_dtensor(value, metadata["chunk_dim"], len(metadata["unfused_keys"])) + else: + chunks = [value] * len(metadata["unfused_keys"]) # scalar — replicate + for unfused_key, chunk in zip(metadata["unfused_keys"], chunks): + optimizer_state[unfused_key][optim_param_name] = chunk + + del optimizer_state[param_name] + + +def fuse_optimizer_state(optimizer_state_dict, fusion_metadata): + """Fuse the optimizer state dict back in place using the fusion info. + + Inverse of `unfuse_optimizer_state`: concatenates each set of per-piece + sub-dicts along `chunk_dim`, rewraps each DTensor field with the original + `_StridedShard` placement, then replaces the piece keys with the fused key + in `optimizer_state_dict["state"]`. + + Args: + optimizer_state_dict: the unfused state dict (e.g. just filled by + `dcp.load`). For our Mixtral example, before this call it looks like: + + { + "state": { + "model.layers.0.mlp.experts.gate_up_proj.0": { + "exp_avg": DTensor(shape=(4, 8, 8), + placements=(Shard(0), Shard(1))), + "exp_avg_sq": DTensor(shape=(4, 8, 8), + placements=(Shard(0), Shard(1))), + "step": tensor(7.0), + }, + "model.layers.0.mlp.experts.gate_up_proj.1": { + "exp_avg": DTensor(shape=(4, 8, 8), + placements=(Shard(0), Shard(1))), + "exp_avg_sq": DTensor(shape=(4, 8, 8), + placements=(Shard(0), Shard(1))), + "step": tensor(7.0), + }, + "model.layers.0.mlp.experts.down_proj": { + "exp_avg": DTensor(shape=(4, 8, 16), + placements=(Shard(0), Shard(2))), + "step": tensor(7.0), + }, + }, + "param_groups": [{"lr": 1e-4, ...}], + } + + After the call the `.0`/`.1` piece keys are gone and the original + fused key is restored with the `_StridedShard` placement. + + fusion_metadata: as returned by `get_fusion_metadata`. + """ + optimizer_state = optimizer_state_dict.get("state", {}) + + for param_name, metadata in fusion_metadata.items(): + first_piece = optimizer_state[metadata["unfused_keys"][0]] + + merged_fields = {} + for optim_param_name, value in first_piece.items(): + chunks = [optimizer_state[unfused_key][optim_param_name] for unfused_key in metadata["unfused_keys"]] + if isinstance(value, DTensor): + merged_fields[optim_param_name] = _merge_unfused_dtensors( + chunks, metadata["chunk_dim"], metadata["placements"] + ) + else: + merged_fields[optim_param_name] = value # scalar — pick any copy + + optimizer_state[param_name] = merged_fields + + for unfused_key in metadata["unfused_keys"]: + del optimizer_state[unfused_key] diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py new file mode 100644 index 000000000000..126b404594f2 --- /dev/null +++ b/src/transformers/distributed/tensor_parallel.py @@ -0,0 +1,526 @@ +# Copyright 2024 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +from __future__ import annotations + +import contextlib +import re + +from ..utils import logging +from ..utils.generic import GeneralInterface +from ..utils.import_utils import is_torch_available, is_torch_greater_or_equal + + +if is_torch_available(): + import torch + +if is_torch_available() and is_torch_greater_or_equal("2.5"): + import torch.distributed as dist + from torch.distributed.tensor import DTensor, Partial, Replicate, Shard, distribute_tensor + from torch.distributed.tensor.parallel import ( + ColwiseParallel, + PrepareModuleInput, + RowwiseParallel, + SequenceParallel, + ) + from torch.distributed.tensor.parallel.style import ParallelStyle + from torch.distributed.tensor.placement_types import _StridedShard + + # Cache this result has it's a C FFI call which can be pretty time-consuming + _torch_distributed_available = torch.distributed.is_available() + + +logger = logging.get_logger(__name__) + + +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_tp_plan(parameter_name: str, tp_plan: dict[str, str], is_weight=True) -> str | None: + """ + Get the TP style for a parameter from the TP plan. + + The TP plan is a dictionary that maps parameter names to TP 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 tp_plan: + return tp_plan[generic_param_name] + elif is_weight and "." in generic_param_name and (module_name := generic_param_name.rsplit(".", 1)[0]) in tp_plan: + return tp_plan[module_name] + return None + + +# ============================================================================= +# High-Level API Functions +# ============================================================================= + + +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. + + Only weight-sharding rules (colwise, rowwise, vocab, moe_experts) are checked. + Module/activation entries (e.g. PrepareModuleInput, SequenceParallel) set up + communication hooks on modules, not weight sharding, so they are excluded. + """ + + if tp_plan is None: + return + + # Filter out module-level comm hooks — they don't shard weights. + # Plan values are registry names; entries beginning with "activation" or "module" + # configure communication hooks rather than parameter sharding. + weight_plan = { + k: v for k, v in tp_plan.items() if not (v == "activation" or v.startswith(("activation_", "module_"))) + } + + generic_keys = {replace_layer_number_by_wildcard(key) for key in expected_keys} + unsharded_layers = set(generic_keys) + unused_rules = weight_plan.copy() + + for key in generic_keys: + param_name = key.rsplit(".", 1)[0] if "." in key else key + generic_param_name = re.sub(r"\d+", "*", param_name) + + if generic_param_name in weight_plan: + unused_rules.pop(generic_param_name, None) + unsharded_layers.discard(key) + elif "." in generic_param_name and (parent_param_name := generic_param_name.rsplit(".", 1)[0]) in weight_plan: + unused_rules.pop(parent_param_name, None) + unsharded_layers.discard(key) + + if len(unused_rules) > 0: + logger.warning(f"The following TP rules were not applied on any of the layers: {unused_rules}") + if len(unsharded_layers) > 0: + logger.warning(f"The following layers were not sharded: {', '.join(unsharded_layers)}") + + +class TensorParallelStyle(ParallelStyle): + """Base class for transformers TP styles. Installs the pre / around / post + forward hooks. Subclasses that need to shard params override `_apply` to + wrap them as DTensor placeholders before calling `super()._apply(...)`. + + Param wrapping runs on meta (the model is on meta when `apply_tensor_parallel` + is invoked); `distribute_tensor` on meta builds metadata only — no collective. + Real data flows in later, async, via DtensorShardOperation during load. + + Forward-time hooks (override what you need): + - transform_inputs_pre_forward(module, args, kwargs, mesh) → (args, kwargs) + - context_around_forward(module) → context manager wrapping the call + - transform_output_post_forward(module, output, mesh) → output + """ + + def transform_inputs_pre_forward(self, module, args, kwargs, mesh): + return args, kwargs + + def context_around_forward(self, module): + return contextlib.nullcontext() + + def transform_output_post_forward(self, module, output, mesh): + return output + + def _apply(self, module, mesh): + original_forward = module.forward + + def tp_forward(*args, **kwargs): + args, kwargs = self.transform_inputs_pre_forward(module, args, kwargs, mesh) + with self.context_around_forward(module): + output = original_forward(*args, **kwargs) + return self.transform_output_post_forward(module, output, mesh) + + module.forward = tp_forward + return module + + +class PrepareModuleInputOutput(TensorParallelStyle): + """Allgather input (Shard(1) → Replicate) + local split output (Replicate → Shard(1)). + + Used for MoE blocks with SP: the input sequence is gathered before routing, + and the output (after expert allreduce) is split back to match the residual. + Forward output split is a local op (no comm). Backward creates the all-gather. + """ + + def __init__(self, use_local_output=True): + super().__init__() + self.use_local_output = use_local_output + + def transform_inputs_pre_forward(self, module, args, kwargs, mesh): + x = args[0] + if not isinstance(x, DTensor): + x = DTensor.from_local(x, mesh, [Shard(1)], run_check=False) + x = x.redistribute(placements=[Replicate()]).to_local() + return (x,) + args[1:], kwargs + + def transform_output_post_forward(self, module, output, mesh): + if not isinstance(output, DTensor): + output = DTensor.from_local(output, mesh, [Replicate()], run_check=False) + return output.redistribute(placements=[Shard(1)]).to_local() + + +def _accumulate_local_param_grad(original_param: DTensor, local_grad: torch.Tensor) -> torch.Tensor: + """Stitch a local grad back onto the original DTensor parameter. + + During forward we replace the DTensor param with a detached plain-tensor + leaf (see ``_swap_dtensor_params_for_local``) because ``grouped_mm`` / fused + ops don't accept DTensor inputs. That swap breaks the autograd link between + the local leaf's grad and the DTensor param's ``.grad``, so this tensor hook + runs on the leaf and copies/accumulates the grad onto the original DTensor. + + NOTE: An autograd-aware ``param.to_local()`` swap would let backward stitch + the grad automatically, but DTensor's backward path then redistributes the + resulting grad — and that redistribute does not currently support + ``_StridedShard`` placements (used by ``MoEExpertsParallel`` / ``PackedColwiseParallel``). + """ + tensor_meta = original_param._spec.tensor_meta + detached_grad = local_grad.detach() + grad_dtensor = DTensor.from_local( + detached_grad, + original_param.device_mesh, + original_param.placements, + run_check=False, + shape=tensor_meta.shape, + stride=tensor_meta.stride, + ) + with torch.no_grad(): + existing_grad = original_param.grad + if existing_grad is None: + original_param.grad = grad_dtensor + elif isinstance(existing_grad, DTensor): + existing_grad._local_tensor.add_(detached_grad) + else: + existing_grad.add_(detached_grad) + return local_grad + + +@contextlib.contextmanager +def _swap_dtensor_params_for_local(module): + """Temporarily replace DTensor params with local-shard ``Parameter``s for forward. + + ``grouped_mm`` / fused kernels don't accept DTensor inputs, so each DTensor + param is swapped for a detached local ``Parameter``. A tensor hook + (``_accumulate_local_param_grad``) on the local leaf copies the backward + grad back onto the original DTensor. + + The original DTensor params are restored on exit (even on exception) so + save_pretrained / state-dict still see sharded params. + """ + shadows = {} + for name, param in list(module.named_parameters(recurse=False)): + if not isinstance(param, DTensor): + continue + shadows[name] = param + local = torch.nn.Parameter(param._local_tensor.detach(), requires_grad=param.requires_grad) + if param.requires_grad: + local.register_hook(lambda g, p=param: _accumulate_local_param_grad(p, g)) + module._parameters.pop(name) + setattr(module, name, local) + try: + yield + finally: + for name, param in shadows.items(): + if hasattr(module, name): + delattr(module, name) + module.register_parameter(name, param) + + +class PackedColwiseParallel(TensorParallelStyle): + """Column-wise parallel style for fused linear weights packed along the output dimension.""" + + def __init__( + self, + *, + input_layouts=None, + use_local_output: bool = True, + split_factor: int = 2, + ): + super().__init__() + self.input_layouts = (input_layouts or Replicate(),) + self.use_local_output = use_local_output + self.split_factor = split_factor + + def transform_inputs_pre_forward(self, module, args, kwargs, mesh): + input_tensor = args[0] + if not isinstance(input_tensor, DTensor): + input_tensor = DTensor.from_local(input_tensor, mesh, self.input_layouts, run_check=False) + elif input_tensor.placements != self.input_layouts: + input_tensor = input_tensor.redistribute(placements=self.input_layouts) + input_tensor = input_tensor.to_local() + return (input_tensor,) + args[1:], kwargs + + def context_around_forward(self, module): + return _swap_dtensor_params_for_local(module) + + def transform_output_post_forward(self, module, output, mesh): + if output is None or self.use_local_output: + return output + return DTensor.from_local( + output, mesh, (_StridedShard(dim=-1, split_factor=self.split_factor),), run_check=False + ) + + def _apply(self, module, mesh): + if not isinstance(module, torch.nn.Linear): + raise NotImplementedError("PackedColwiseParallel currently only supports nn.Linear!") + # Wrap weight + bias as DTensor placeholders. Runs on meta — + # distribute_tensor builds metadata only, no collective. + placement = _StridedShard(dim=0, split_factor=self.split_factor) + for name in ("weight", "bias"): + meta = module._parameters.get(name) + if meta is None: + continue + module._parameters[name] = torch.nn.Parameter( + distribute_tensor(meta, mesh, [placement], src_data_rank=None), + requires_grad=meta.requires_grad, + ) + return super()._apply(module, mesh) + + def __repr__(self) -> str: + return ( + f"{self.__class__.__name__}(input_layouts={self.input_layouts}, " + f"use_local_output={self.use_local_output}, split_factor={self.split_factor})" + ) + + +if is_torch_available() and is_torch_greater_or_equal("2.5"): + + class _AllReduceBackward(torch.autograd.Function): + """Identity forward, allreduce-sum backward. + + Used for MoE routing weights: the forward value is replicated (same on all + ranks), but the backward gradient is partial (each rank has 1/tp_size from + its expert shard). We need to sum the partial gradients without dividing by + world_size, which is what DTensor's ``Replicate`` backward does incorrectly. + """ + + @staticmethod + def forward(ctx, x, process_group): + ctx.process_group = process_group + return x + + @staticmethod + def backward(ctx, grad): + dist.all_reduce(grad, group=ctx.process_group) + return grad, None + + +class MoEExpertsParallel(TensorParallelStyle): + """Tensor-parallel style for MoE expert modules. + + Shards expert weights as DTensors, then wraps the module's ``forward`` so + that grouped_mm (which needs plain tensors) works transparently. + + Lifecycle phases: + 1. _apply — wrap each expert weight named in shard_plan as a DTensor + placeholder with the declared placement. + 2. transform_inputs_pre_forward — localize hidden_states (Replicate→local, + gives us an all-reduce on the backward gradient for free), then fix + routing-weight gradients (their backward is partial; use allreduce-sum, + not divide-by-world-size). + 3. context_around_forward — swap DTensor params for local leaves so + grouped_mm sees plain tensors; restored on exit so save_pretrained + still sees DTensors. + 4. transform_output_post_forward — under TP-only each rank's output is + partial (only its expert shard contributed); reduce/redistribute to + output_layouts. + """ + + def __init__(self, output_layouts=None, shard_plan=None): + super().__init__() + self.output_layouts = output_layouts or Replicate() + self._moe_shard_plan = shard_plan or {} + + def _apply(self, module, mesh): + # Wrap each expert weight as a DTensor placeholder. Runs on meta — + # distribute_tensor builds metadata only, no collective. + for name, placement in self._moe_shard_plan.items(): + meta = module._parameters.get(name) + if meta is None: + continue + module._parameters[name] = torch.nn.Parameter( + distribute_tensor(meta, mesh, [placement], src_data_rank=None), + requires_grad=meta.requires_grad, + ) + return super()._apply(module, mesh) + + def transform_inputs_pre_forward(self, module, args, kwargs, mesh): + hidden_states, top_k_index, top_k_weights = args + if not isinstance(hidden_states, DTensor): + hidden_states = DTensor.from_local(hidden_states, mesh, [Replicate()], run_check=False) + hidden_states = hidden_states.to_local() + + if isinstance(top_k_weights, DTensor): + top_k_weights = top_k_weights.to_local() + tp_group = mesh.get_group() if mesh.ndim == 1 else mesh.get_group("tp") + top_k_weights = _AllReduceBackward.apply(top_k_weights, tp_group) + + return (hidden_states, top_k_index, top_k_weights), kwargs + + def context_around_forward(self, module): + return _swap_dtensor_params_for_local(module) + + def transform_output_post_forward(self, module, output, mesh): + if output is None: + return None + # Under TP-only each rank has a partial result; under TP+FSDP the + # weights may be fully gathered by FSDP, making the output complete. + has_sharded_params = any( + isinstance(p, DTensor) and any(not pl.is_replicate() for pl in p.placements) for p in module.parameters() + ) + source = Partial() if has_sharded_params else Replicate() + if not isinstance(output, DTensor): + output = DTensor.from_local(output, mesh, [source], run_check=False) + # MoE output is 2D [tokens, hidden]. For SP, Shard(1) means seq dim + # in 3D but token dim (0) in 2D. + target = self.output_layouts + if output.dim() == 2 and isinstance(target, Shard) and target.dim == 1: + target = Shard(0) + if output.placements != (target,): + output = output.redistribute(placements=(target,)) + return output.to_local() + + +class ParallelInterface(GeneralInterface): + """Registry of named TP styles. Configs and modeling files reference these by string name. + + Adding a new entry here is the supported way to introduce a new TP style. + Users can also override or extend at runtime via ``ALL_PARALLEL_STYLES["my_style"] = ...``. + + Naming convention: ``{kind}[_{comm}][_{extra}]``. The ``_{comm}`` suffix is dropped only when + comm is ``"none"`` (no collective). All entries are eager instances; the dict literal lives + behind a torch-availability guard so this module remains importable without torch. + """ + + _global_mapping = ( + { + # Column-parallel + "colwise": ColwiseParallel(input_layouts=Replicate(), output_layouts=Shard(-1)), + "colwise_allgather": ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()), + "colwise_loss_parallel": ColwiseParallel( + input_layouts=Shard(1), output_layouts=Shard(-1), use_local_output=False + ), + "packed_colwise": PackedColwiseParallel(input_layouts=Replicate()), + # Row-parallel + "rowwise_allreduce": RowwiseParallel(input_layouts=Shard(-1), output_layouts=Replicate()), + "rowwise_reduce_scatter": RowwiseParallel(input_layouts=Shard(-1), output_layouts=Shard(1)), + # Vocab / embedding (rowwise sharding on vocab dim) + "vocab_allreduce": RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()), + "vocab_reduce_scatter": RowwiseParallel(input_layouts=Replicate(), output_layouts=Shard(1)), + # Activation / norm (sequence-parallel passthrough) + # use_local_output=True: torch defaults to False here, but downstream modeling + # code expects plain tensors, not DTensors. + "activation": SequenceParallel(use_local_output=True), + "activation_seq_dim_2": SequenceParallel(sequence_dim=2, use_local_output=True), + # Module-level prepare-input. Same use_local_output=True override as above — + # torch's default is False, our modeling code expects plain tensors downstream. + "module_allgather": PrepareModuleInput( + input_layouts=(Shard(1),), desired_input_layouts=(Replicate(),), use_local_output=True + ), + "module_allgather_hidden_states": PrepareModuleInput( + input_kwarg_layouts={"hidden_states": Shard(1)}, + desired_input_kwarg_layouts={"hidden_states": Replicate()}, + use_local_output=True, + ), + "module_allgather_split": PrepareModuleInputOutput(), + # MoE — canonical shard_plan baked in (only variant in use across configs). + # gate_up_proj is packed (gate||up along output dim) so we use _StridedShard + # to interleave; down_proj is plain rowwise on its input dim. + "moe_experts_allreduce": MoEExpertsParallel( + output_layouts=Replicate(), + shard_plan={ + "gate_up_proj": _StridedShard(dim=-2, split_factor=2), + "down_proj": Shard(-1), + }, + ), + } + if is_torch_available() and is_torch_greater_or_equal("2.5") and _torch_distributed_available + else {} + ) + + +ALL_PARALLEL_STYLES: ParallelInterface = ParallelInterface() + + +def apply_tensor_parallel(model, tp_mesh, tp_plan): + """Apply tensor parallelism by calling each style's ``_apply`` on the + matching submodules. + + Walks ``model.named_modules()``, resolves each name against the wildcard + ``tp_plan`` from the model config, and applies the corresponding style + from ``ALL_PARALLEL_STYLES`` (looked up by string name) directly. + """ + distributed_config = getattr(model.config, "distributed_config", None) + sp_requested = getattr(distributed_config, "enable_sequence_parallel", False) + sp_supported = getattr(model.config, "base_model_sp_plan", None) is not None + enable_sp = sp_requested and sp_supported + + if tp_plan is None: + if enable_sp: + tp_plan = dict(model._sp_plan or {}) + else: + tp_plan = dict(model._tp_plan or {}) + + # tie_weights() replaces lm_head.weight with embed_tokens.weight after TP is applied. + # If embed_tokens isn't in the plan, sharding lm_head as a DTensor causes tie to + # clobber it with a plain tensor (and forward then mixes DTensor/Tensor). Skip + # lm_head TP in that case so both ends stay plain and the tie is a real alias. + if getattr(model.config, "tie_word_embeddings", False): + tied_source_in_plan = any(k.endswith("embed_tokens") for k in tp_plan) + if not tied_source_in_plan: + tp_plan.pop("lm_head", None) + + for name, submodule in model.named_modules(): + style_value = _get_parameter_tp_plan(parameter_name=name, tp_plan=tp_plan, is_weight=False) + if style_value is None: + continue + ALL_PARALLEL_STYLES[style_value]._apply(submodule, tp_mesh) + + # Under SP, inputs_embeds is sequence-sharded after embed_tokens, so + # auto-generated position_ids would use the wrong (local) seq_len. + # Inject position_ids from the original input_ids shape before the model forward + if enable_sp: + base_model = getattr(model, model.base_model_prefix, model) + + def _inject_sp_metadata(mod, args, kwargs): + input_ids = kwargs.get("input_ids", args[0] if args else None) + if input_ids is None: + return args, kwargs + if "position_ids" not in kwargs or kwargs["position_ids"] is None: + seq_len = input_ids.shape[1] + kwargs["position_ids"] = torch.arange(seq_len, device=input_ids.device).unsqueeze(0) + return args, kwargs + + base_model.register_forward_pre_hook(_inject_sp_metadata, with_kwargs=True) + + # If the plan uses loss_parallel on lm_head, enable it globally so + # the model's internal loss computation handles DTensor logits correctly. + # loss_parallel patches F.cross_entropy to work with Shard(-1) logits. + # It must be active during both forward and backward, so we enable it + # once rather than as a context manager. + has_loss_parallel = any(v == "colwise_loss_parallel" for v in tp_plan.values()) + if has_loss_parallel: + from torch.distributed.tensor.parallel import loss_parallel + + model._loss_parallel_ctx = loss_parallel() + model._loss_parallel_ctx.__enter__() + + return model diff --git a/src/transformers/distributed/utils.py b/src/transformers/distributed/utils.py new file mode 100644 index 000000000000..50cf2fa4bb28 --- /dev/null +++ b/src/transformers/distributed/utils.py @@ -0,0 +1,296 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +from __future__ import annotations + +import os +from typing import TYPE_CHECKING + +from ..utils import is_torch_available, is_torch_greater_or_equal, logging +from .fsdp import apply_fully_shard_data_parallel +from .sharding_utils import ( + _find_strided_shard_placement_from_fused_params, + _replicate_dtensor, + fuse_optimizer_state, + get_fusion_metadata, + unfuse_optimizer_state, +) +from .tensor_parallel import apply_tensor_parallel + + +logger = logging.get_logger(__name__) + + +if TYPE_CHECKING: + import torch.nn as nn + + from .configuration_utils import DistributedConfig + +if is_torch_available(): + import torch + import torch.distributed.checkpoint as dcp + from torch.distributed.checkpoint.hf_storage import HuggingFaceStorageWriter + from torch.distributed.checkpoint.state_dict import ( + get_model_state_dict, + get_optimizer_state_dict, + set_optimizer_state_dict, + ) + from torch.distributed.tensor import DTensor + + +def _ensure_torch_distributed(device_type: str): + """Initialize torch.distributed if not already initialized.""" + if not torch.distributed.is_initialized(): + try: + rank = int(os.environ["RANK"]) + local_rank = int(os.environ["LOCAL_RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + + backend_map = { + "cuda": "nccl", + "cpu": "gloo", + "xpu": "xccl", + "hpu": "hccl", + "neuron": "neuron", + "tpu": "tpu_dist", + } + backend = backend_map.get(device_type) + + # Bind the accelerator before init so the process group is created with a + # device_id, otherwise collectives like barrier() warn (and may spin up an + # extra NCCL comm) about the missing device binding. + device_id = None + if device_type != "cpu": + getattr(torch, device_type).set_device(local_rank) + device_id = torch.device(device_type, local_rank) + torch.distributed.init_process_group( + backend=backend, rank=rank, world_size=world_size, device_id=device_id + ) + except Exception as e: + raise OSError( + "We tried to initialize torch.distributed for you, but it failed. Make " + "sure you init torch distributed in your script to use distributed training." + ) from e + + +def _distributed_barrier(): + """Barrier bound to the current accelerator device. + + Passing `device_ids` is required when the process group was initialized without a + `device_id`; with it, the call is a no-op compared to plain `barrier()`. Safe to call + when torch.distributed has not been initialized — returns immediately. + """ + if not torch.distributed.is_initialized(): + return + device_type = torch._C._get_accelerator().type + if device_type != "cpu": + torch.distributed.barrier(device_ids=[getattr(torch, device_type).current_device()]) + else: + torch.distributed.barrier() + + +def init_device_mesh(distributed_config: DistributedConfig) -> torch.distributed.device_mesh.DeviceMesh: + if not is_torch_greater_or_equal("2.5"): + raise OSError("Distributed training with DistributedConfig requires `torch>=2.5`.") + + device_type = torch._C._get_accelerator().type + _ensure_torch_distributed(device_type) + + world_size = torch.distributed.get_world_size() + if device_type != "cpu": + getattr(torch, device_type).set_device(int(os.environ.get("LOCAL_RANK", 0))) + + tp_size = distributed_config.tp_size + fsdp_size = distributed_config.fsdp_size + + assert world_size == tp_size * fsdp_size, ( + f"world_size ({world_size}) must be equal to tp_size ({tp_size}) * fsdp_size ({fsdp_size})" + ) + + dims, names = [], [] + if fsdp_size > 1: + dims.append(fsdp_size) + names.append("fsdp") + if tp_size > 1: + dims.append(tp_size) + names.append("tp") + + # Build from a 1D world mesh via _unflatten so that PyTorch can flatten + # sub-dimensions back when needed (e.g. for single all_reduce across + # [fsdp, tp] during grad norm computation instead of 2 sequential ones). + world_mesh = torch.distributed.init_device_mesh(device_type, (world_size,), mesh_dim_names=("world",)) + mesh = world_mesh._unflatten(0, tuple(dims), tuple(names)) + + # Pre-create flattened sub-mesh for multi-dimensional meshes so DTensor + # can use a single collective instead of sequential per-dimension ones. + if len(dims) > 1: + mesh._flatten("_".join(names)) + + return mesh + + +def distribute_model(model, distributed_config: DistributedConfig, device_mesh) -> nn.Module: + """Apply TP and/or FSDP2 to `model` based on the mesh dims in `device_mesh`.""" + model.config.distributed_config = distributed_config + model.device_mesh = device_mesh + mesh_dim_names = device_mesh.mesh_dim_names or () + if "tp" in mesh_dim_names: + tp_mesh = device_mesh["tp"] if device_mesh.ndim > 1 else device_mesh + model = apply_tensor_parallel(model, tp_mesh, distributed_config.tp_plan) + if "fsdp" in mesh_dim_names: + fsdp_mesh = device_mesh["fsdp"] if device_mesh.ndim > 1 else device_mesh + model = apply_fully_shard_data_parallel(model, fsdp_mesh, distributed_config.fsdp_plan) + return model + + +@torch.no_grad() +def clip_grad_norm(parameters, max_norm: float, norm_type: float = 2.0): + """Grad-norm clip that works when params live on different DTensor meshes. + + ``torch.nn.utils.get_total_norm`` stacks per-grad norms; that fails when grads + live on different meshes (e.g. TP-wrapped params on the (fsdp, tp) mesh and + FSDP-only params on the (fsdp,) sub-mesh). We sidestep it by replicating each + DTensor grad to a plain local tensor, computing the norm over those, and + scaling the original DTensor grads in place — the placement of the original + grads doesn't matter for the per-element clip. + """ + grads = [p.grad for p in parameters if p.grad is not None] + local_grads = [_replicate_dtensor(g).to_local() if isinstance(g, DTensor) else g for g in grads] + total_norm = torch.nn.utils.get_total_norm(local_grads, norm_type=norm_type) + torch.nn.utils.clip_grads_with_norm_(grads, max_norm=max_norm, total_norm=total_norm) + return total_norm + + +def gather_full_state_dict(model) -> dict[str, torch.Tensor]: + """Gather all sharded params to full plain tensors for saving. + + Handles FSDP unshard and TP DTensor gather. + Streams one parameter at a time to avoid holding all full tensors on GPU. + Only rank 0 accumulates the result; other ranks return ``{}``. + """ + is_rank0 = torch.distributed.get_rank() == 0 + state_dict = get_model_state_dict(model) + + result = {} + for key, tensor in state_dict.items(): + if not isinstance(tensor, DTensor): + if is_rank0: + result[key] = tensor.detach().to(device="cpu", copy=True).contiguous() + continue + + with torch.no_grad(): + full = _replicate_dtensor(tensor).to_local() + if is_rank0: + result[key] = full.detach().to(device="cpu", copy=True).contiguous() + del full + return result + + +def save_model_checkpoint_distributed(model, checkpoint_dir: str) -> None: + """Save model parameters as standard HF-format sharded safetensors using + DCP + HuggingFaceStorageWriter with consolidation enabled. + + Every rank first writes its own shard in parallel under + `/sharded/`, then a consolidation pass reads those shards + and emits HF-compatible `model-*-of-N.safetensors` (+ index) at + `/`. The result is a directory `from_pretrained` reads + through its normal path — no special flag needed at load time. + + DTensors carrying an uncomposed `_StridedShard` placement (e.g. fused + gate||up MoE weights) are replicated to a full tensor on every rank + before the save, otherwise DCP cannot encode that placement. + """ + state_dict = get_model_state_dict(model) + for key, value in list(state_dict.items()): + if ( + isinstance(value, DTensor) + and _find_strided_shard_placement_from_fused_params(value.placements) is not None + ): + state_dict[key] = _replicate_dtensor(value) + + dcp.save( + state_dict, + storage_writer=HuggingFaceStorageWriter( + path=checkpoint_dir, + save_distributed=True, + enable_consolidation=True, + ), + ) + # Wait for rank 0 to finish writing the HF safetensors so other + # ranks don't return (and hit `from_pretrained`) before the files exist. + _distributed_barrier() + + +def has_mixed_tensor_and_dtensor_params(params) -> bool: + has_dtensor = False + has_tensor = False + for param in params: + if isinstance(param, DTensor): + has_dtensor = True + elif isinstance(param, torch.Tensor): + has_tensor = True + + if has_dtensor and has_tensor: + return True + return False + + +def maybe_disable_foreach_and_fused_for_mixed_dtensor_groups(optimizer) -> None: + """ + When get_optimizer_state_dict() or set_optimizer_state_dict() runs on an optimizer with no state yet, + PyTorch first materializes that state by doing a no-op step() with zero gradients. If an optimizer + group mixes regular tensors and DTensors, the batched foreach/fused optimizer kernels cannot process + that mixed group, so we turn those kernels off for such groups before distributed optimizer save/ + load. + """ + for i, param_group in enumerate(optimizer.param_groups): + if has_mixed_tensor_and_dtensor_params(param_group.get("params", ())): + logger.warning_once( + f"Param group {i} mixes regular tensors and DTensors; disabling foreach/fused " + "optimizer kernels for that group so distributed optimizer save/load can materialize state." + ) + param_group["foreach"] = False + if "fused" in param_group: + param_group["fused"] = False + + +def save_optimizer_distributed(model, optimizer, checkpoint_dir: str) -> None: + """Save optimizer state via DCP. + + Params whose DTensors carry a lonely `_StridedShard` placement (e.g. Mixtral + `gate_up_proj`) are locally split into plain-`Shard` pieces at the boundary + so DCP only ever sees DTensors it can encode as one contiguous chunk per + rank. + """ + maybe_disable_foreach_and_fused_for_mixed_dtensor_groups(optimizer) + optimizer_state_dict = get_optimizer_state_dict(model, optimizer) + fusion_metadata = get_fusion_metadata(optimizer_state_dict) + unfuse_optimizer_state(optimizer_state_dict, fusion_metadata) + dcp.save({"optimizer": optimizer_state_dict}, checkpoint_id=checkpoint_dir) + + +def load_optimizer_distributed(model, optimizer, checkpoint_dir: str) -> None: + """Load optimizer state via DCP. + + Symmetric to `save_optimizer_distributed`: build the unfused template, let DCP fill + it, then merge fused params back to their original `_StridedShard` form + before handing the state_dict back to the optimizer. + """ + maybe_disable_foreach_and_fused_for_mixed_dtensor_groups(optimizer) + optimizer_state_dict = get_optimizer_state_dict(model, optimizer) + fusion_metadata = get_fusion_metadata(optimizer_state_dict) + unfuse_optimizer_state(optimizer_state_dict, fusion_metadata) + dcp.load({"optimizer": optimizer_state_dict}, checkpoint_id=checkpoint_dir) + fuse_optimizer_state(optimizer_state_dict, fusion_metadata) + set_optimizer_state_dict(model, optimizer, optimizer_state_dict) + maybe_disable_foreach_and_fused_for_mixed_dtensor_groups(optimizer) diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py index dcfbab1edf23..af20dd8d2242 100644 --- a/src/transformers/generation/utils.py +++ b/src/transformers/generation/utils.py @@ -33,6 +33,7 @@ QuantizedCache, StaticCache, ) +from ..distributed.fsdp import is_fsdp_managed_module from ..dynamic_module_utils import ( check_python_requirements, get_cached_module_file, @@ -40,7 +41,6 @@ resolve_trust_remote_code, ) from ..integrations.deepspeed import is_deepspeed_zero3_enabled -from ..integrations.fsdp import is_fsdp_managed_module from ..masking_utils import create_masks_for_generate from ..tokenization_python import ExtensionsTrie from ..utils import ( @@ -2152,7 +2152,7 @@ def _extract_generation_mode_kwargs( "assistant_model": assistant_model, "streamer": streamer, } - world_size = dist.get_world_size() if dist.is_available() and dist.is_initialized() else 1 # type: ignore + world_size = dist.get_world_size() if dist.is_available() and dist.is_initialized() else 1 generation_mode_kwargs["synced_gpus"] = ( (is_deepspeed_zero3_enabled() or is_fsdp_managed_module(self)) and world_size > 1 if synced_gpus is None @@ -2599,7 +2599,7 @@ def _has_unfinished_sequences(self, this_peer_finished: bool, synced_gpus: bool, # The following logic allows an early break if all peers finished generating their sequence this_peer_finished_flag = torch.tensor(0.0 if this_peer_finished else 1.0, device=device) # send 0.0 if we finished, 1.0 otherwise - dist.all_reduce(this_peer_finished_flag, op=dist.ReduceOp.SUM) # type: ignore + dist.all_reduce(this_peer_finished_flag, op=dist.ReduceOp.SUM) # did all peers finish? the reduced sum will be 0.0 then if this_peer_finished_flag.item() == 0.0: return False diff --git a/src/transformers/integrations/__init__.py b/src/transformers/integrations/__init__.py index 336db3773f76..f6282cf4609e 100755 --- a/src/transformers/integrations/__init__.py +++ b/src/transformers/integrations/__init__.py @@ -50,7 +50,6 @@ "eetq": ["replace_with_eetq_linear"], "fbgemm_fp8": ["FbgemmFp8Linear", "FbgemmFp8Llama4TextExperts", "replace_with_fbgemm_fp8_linear"], "finegrained_fp8": ["FP8Linear", "replace_with_fp8_linear"], - "fsdp": ["is_fsdp_enabled", "is_fsdp_managed_module"], "ggml": [ "GGUF_CONFIG_DEFAULTS_MAPPING", "GGUF_CONFIG_MAPPING", @@ -160,11 +159,6 @@ "convert_and_export_with_cache", ] -_import_structure["tensor_parallel"] = [ - "shard_and_distribute_module", - "ALL_PARALLEL_STYLES", - "translate_to_torch_parallel_style", -] try: if not is_torch_greater_or_equal("2.5"): raise OptionalDependencyNotAvailable() @@ -209,7 +203,6 @@ from .eetq import replace_with_eetq_linear from .fbgemm_fp8 import FbgemmFp8Linear, FbgemmFp8Llama4TextExperts, replace_with_fbgemm_fp8_linear from .finegrained_fp8 import FP8Linear, replace_with_fp8_linear - from .fsdp import is_fsdp_enabled, is_fsdp_managed_module from .ggml import ( GGUF_CONFIG_DEFAULTS_MAPPING, GGUF_CONFIG_MAPPING, @@ -305,12 +298,6 @@ else: from .executorch import TorchExportableModuleWithStaticCache, convert_and_export_with_cache - from .tensor_parallel import ( - ALL_PARALLEL_STYLES, - shard_and_distribute_module, - translate_to_torch_parallel_style, - ) - try: if not is_torch_greater_or_equal("2.5"): raise OptionalDependencyNotAvailable() diff --git a/src/transformers/integrations/accelerate.py b/src/transformers/integrations/accelerate.py index f29a3194ce02..13aeb2410bfb 100644 --- a/src/transformers/integrations/accelerate.py +++ b/src/transformers/integrations/accelerate.py @@ -26,6 +26,7 @@ from safetensors import safe_open from safetensors.torch import save_file +from ..distributed.fsdp import is_fsdp_enabled from ..utils import ( is_accelerate_available, is_torch_available, @@ -34,7 +35,6 @@ ) from ..utils.quantization_config import QuantizationMethod from .deepspeed import is_deepspeed_zero3_enabled -from .fsdp import is_fsdp_enabled if is_torch_available(): diff --git a/src/transformers/integrations/fsdp.py b/src/transformers/integrations/fsdp.py deleted file mode 100644 index 7cda7ad55acc..000000000000 --- a/src/transformers/integrations/fsdp.py +++ /dev/null @@ -1,92 +0,0 @@ -# Copyright 2024 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# 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. -from __future__ import annotations - -import inspect -import os -from typing import TYPE_CHECKING - -from ..utils import is_torch_available, strtobool -from ..utils.quantization_config import QuantizationMethod - - -if TYPE_CHECKING: - from torch import nn - - -def is_fsdp_managed_module(module: nn.Module) -> bool: - if not is_torch_available(): - return False - - import torch - - if not torch.distributed.is_available(): - return False - - import torch.distributed.fsdp - - return isinstance(module, torch.distributed.fsdp.FullyShardedDataParallel) or getattr( - module, "_is_fsdp_managed_module", False - ) - - -def is_fsdp_enabled(): - if is_torch_available(): - import torch - - return ( - torch.distributed.is_available() - and torch.distributed.is_initialized() - and strtobool(os.environ.get("ACCELERATE_USE_FSDP", "False")) == 1 - and strtobool(os.environ.get("FSDP_CPU_RAM_EFFICIENT_LOADING", "False")) == 1 - ) - - return False - - -def get_fsdp_ckpt_kwargs(): - """ - Returns checkpoint kwargs for FSDP model saving. - - Checks if the `adapter_only` parameter is supported by `save_fsdp_model` from accelerate - and returns the appropriate kwargs. - """ - from accelerate.utils import save_fsdp_model - - if "adapter_only" in list(inspect.signature(save_fsdp_model).parameters): - return {"adapter_only": True} - else: - return {} - - -def update_fsdp_plugin_peft(model, accelerator): - """ - Updates the FSDP plugin for PEFT LoRA/QLoRA compatibility. - - When using FSDP with PEFT LoRA, the auto wrap policy needs to be updated to additionally wrap - LoRA trainable layers separately. When using FSDP with QLoRA, the mixed precision policy needs - to be updated to use the quantization storage data type. - """ - from peft import PeftConfig - from peft.utils.other import fsdp_auto_wrap_policy - - if isinstance(model.active_peft_config, PeftConfig): - accelerator.state.fsdp_plugin.auto_wrap_policy = fsdp_auto_wrap_policy(model) - if ( - getattr(model, "quantization_method", None) == QuantizationMethod.BITS_AND_BYTES - and model.hf_quantizer.quantization_config.bnb_4bit_quant_storage.is_floating_point - ): - accelerator.state.fsdp_plugin.set_mixed_precision( - model.hf_quantizer.quantization_config.bnb_4bit_quant_storage, override=True - ) diff --git a/src/transformers/integrations/moe.py b/src/transformers/integrations/moe.py index 4f107f03e6fb..02369999c11f 100644 --- a/src/transformers/integrations/moe.py +++ b/src/transformers/integrations/moe.py @@ -366,7 +366,7 @@ def _grouped_linear( out = _grouped_mm(input, weight, offs=offs) else: # (S, input_dim) @ grouped (num_experts, output_dim, input_dim).T -> (S, output_dim) - out = _grouped_mm(input, weight.transpose(-2, -1), offs=offs) + out = _grouped_mm(input, weight.transpose(-2, -1).contiguous(), offs=offs) if bias is not None: # We should be able to pass bias to the grouped_mm call, but it's not yet supported. diff --git a/src/transformers/integrations/mxfp4.py b/src/transformers/integrations/mxfp4.py index 67d9420659af..482ef891771f 100644 --- a/src/transformers/integrations/mxfp4.py +++ b/src/transformers/integrations/mxfp4.py @@ -511,28 +511,15 @@ def mlp_forward(self, hidden_states): def dequantize(module, param_name, param_value, target_device, dq_param_name, **kwargs): - from ..integrations.tensor_parallel import shard_and_distribute_module - - model = kwargs.get("model") - empty_param = kwargs.get("empty_param") - casting_dtype = kwargs.get("casting_dtype") - to_contiguous = kwargs.get("to_contiguous") - rank = kwargs.get("rank") device_mesh = kwargs.get("device_mesh") + if device_mesh is not None: + raise NotImplementedError( + "MXFP4 quantization is not yet compatible with DTensor-based tensor parallelism. " + "Please disable TP or use a non-quantized model." + ) for proj in ["gate_up_proj", "down_proj"]: if proj in param_name: - if device_mesh is not None: - param_value = shard_and_distribute_module( - model, - param_value, - empty_param, - dq_param_name, - casting_dtype, - to_contiguous, - rank, - device_mesh, - ) blocks_attr = f"{proj}_blocks" scales_attr = f"{proj}_scales" setattr(module, param_name.rsplit(".", 1)[1], param_value) @@ -557,24 +544,19 @@ def load_and_swizzle_mxfp4(module, param_name, param_value, target_device, trito triton_kernels_hub.matmul_ogs.FlexCtx, triton_kernels_hub.matmul_ogs.InFlexData, ) - from ..integrations.tensor_parallel import shard_and_distribute_module - model = kwargs.get("model") - empty_param = kwargs.get("empty_param") - casting_dtype = kwargs.get("casting_dtype") - to_contiguous = kwargs.get("to_contiguous") - rank = kwargs.get("rank") device_mesh = kwargs.get("device_mesh") + if device_mesh is not None: + raise NotImplementedError( + "MXFP4 quantization is not yet compatible with DTensor-based tensor parallelism. " + "Please disable TP or use a non-quantized model." + ) + if "blocks" in param_name: proj = param_name.split(".")[-1].split("_blocks")[0] if "scales" in param_name: proj = param_name.split(".")[-1].split("_scales")[0] - if device_mesh is not None: - shard_and_distribute_module( - model, param_value, empty_param, param_name, casting_dtype, to_contiguous, rank, device_mesh - ) - else: - setattr(module, param_name.rsplit(".", 1)[1], torch.nn.Parameter(param_value, requires_grad=False)) + setattr(module, param_name.rsplit(".", 1)[1], torch.nn.Parameter(param_value, requires_grad=False)) blocks_attr = f"{proj}_blocks" scales_attr = f"{proj}_scales" blocks = getattr(module, blocks_attr) # at this point values were loaded from ckpt diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py deleted file mode 100644 index 76335a2fa00d..000000000000 --- a/src/transformers/integrations/tensor_parallel.py +++ /dev/null @@ -1,1560 +0,0 @@ -# Copyright 2024 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# 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. -from __future__ import annotations - -import math -import operator -import os -import re -from functools import 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 - - -if is_torch_available(): - import torch - import torch.distributed as dist - from torch import nn - - # Cache this result has it's a C FFI call which can be pretty time-consuming - _torch_distributed_available = torch.distributed.is_available() - - -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 -): - r""" - Sets up the device mesh and initialized the backend for tensor parallelism. - This function is called when the model is loaded and the TP plan is set to 'auto'. - """ - if tp_size is not None and tp_plan is None: - raise ValueError("tp_plan has to be set when tp_size is passed.") - if tp_plan is not None and device_map is not None: - raise ValueError("`tp_plan` and `device_map` are mutually exclusive. Choose either one for parallelization.") - if device_mesh is None: - if not is_torch_greater_or_equal("2.5"): - raise OSError("Tensor parallel is only supported for `torch>=2.5`.") - - # Detect the accelerator on the machine. If no accelerator is available, it returns CPU. - device_type = torch._C._get_accelerator().type - if device_type == "mps": - raise RuntimeError("Tensor parallelism is not supported on MPS devices.") - current_device = getattr(torch, device_type) - if not torch.distributed.is_initialized(): - try: - rank = int(os.environ["RANK"]) - local_rank = int(os.environ["LOCAL_RANK"]) - world_size = int(os.environ["WORLD_SIZE"]) - - backend_map = { - "cuda": "nccl", - "cpu": "gloo", - "xpu": "xccl", - "hpu": "hccl", - "neuron": "neuron", - "tpu": "tpu_dist", - } - backend = backend_map.get(device_type) - - torch.distributed.init_process_group(backend=backend, rank=rank, world_size=world_size) - current_device = getattr(torch, device_type) - if device_type != "cpu": - current_device.set_device(local_rank) - - except Exception as e: - raise OSError( - "We tried to initialize torch.distributed for you, but it failed. Make " - "sure you init torch distributed in your script to use `tp_plan`." - ) from e - - if device_type != "cpu": - current_device.set_device(int(os.environ["LOCAL_RANK"])) - index = current_device.current_device() - tp_device = torch.device(device_type, index) - device_map = tp_device - else: - tp_device = torch.device(device_type) - device_map = device_type or {} - - tp_size = tp_size if tp_size is not None else torch.distributed.get_world_size() - device_mesh = torch.distributed.init_device_mesh(tp_device.type, (tp_size,)) - else: - if device_mesh.ndim > 1: - if "tp" not in device_mesh.mesh_dim_names: - raise ValueError( - "When using `tp_plan` and n-d `device_mesh`, it must contain a 'tp' dimension. " - "Please provide a valid `device_mesh`." - ) - device_mesh = device_mesh["tp"] - tp_size = device_mesh.size() - device_map = torch.device(f"{device_mesh.device_type}:{int(os.environ['LOCAL_RANK'])}") - - return device_map, device_mesh, tp_size - - -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_tp_plan(parameter_name: str, tp_plan: dict[str, str], is_weight=True) -> str | None: - """ - Get the TP style for a parameter from the TP plan. - - The TP plan is a dictionary that maps parameter names to TP 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 tp_plan: - return tp_plan[generic_param_name] - elif is_weight and "." in generic_param_name and (module_name := generic_param_name.rsplit(".", 1)[0]) in tp_plan: - return tp_plan[module_name] - return None - - -# ============================================================================= -# Tensor Sharding Utilities -# ============================================================================= - - -if is_torch_available(): - str_to_dtype = { - "BOOL": torch.bool, - "U8": torch.uint8, - "I8": torch.int8, - "I16": torch.int16, - "F16": torch.float16, - "BF16": torch.bfloat16, - "I32": torch.int32, - "F32": torch.float32, - "F64": torch.float64, - "I64": torch.int64, - "F8_E4M3": torch.float8_e4m3fn, - } - - -def _blocks_to_block_sizes(total_size: int, blocks: int | list[int]) -> list[int]: - """ - Convert block count or proportions to block sizes. - - This function accepts - - - The number of blocks (int), in which case the block size is - total_size//blocks; or - - A list of block sizes (list[int]). - - In the second case, if sum(blocks) < total_size, the ratios between - the block sizes will be preserved. For instance, if blocks is - [2, 1, 1] and total_size is 1024, the returned block sizes are - [512, 256, 256]. - """ - if isinstance(blocks, list): - total_blocks = sum(blocks) - assert total_size % total_blocks == 0, f"Cannot split {total_size} in proportional blocks: {blocks}" - part_size = total_size // total_blocks - return [part_size * block for block in blocks] - else: - assert total_size % blocks == 0, f"Prepacked is not divisible by {blocks}" - single_size = total_size // blocks - return [single_size] * blocks - - -def get_packed_weights(param, empty_param, device_mesh, rank, dim): - """ - When weights are packed (gate_up_proj), we need to make sure each shard gets its correct share. - So if you have: gate_proj ( 16, 5120, 8190) - and up_proj ( 16, 5120, 8190) - packed as gate_up_proj ( 16, 5120, 2 * 8190) - And you shard along the last dimension, you need to interleave the gate and up values: - - Now, if we shard along the last dimension across TP_size (Tensor Parallelism size), we must interleave the values from gate and up projections correctly. - - Let's take TP_size = 4 for an example: - - Packed tensor `gate_up_proj` - --------------------------------------------------------------- - [ G0 G1 G2 G3 | G4 G5 G6 G7 | ... | U0 U1 U2 U3 | U4 U5 U6 U7 | ... ] - ↑─────────────↑ ↑─────────────↑ ↑─────────────↑ ↑─────────────↑ - Gate Slice 0 Gate Slice 1 Up Slice 0 Up Slice 1 - - Explanation: - - The first half of the tensor (left of the center) holds the gate_proj values. - - The second half (right of the center) holds the up_proj values. - - For TP=4, we divide each half into 4 slices. In this example, we show two slices for brevity. - - Each shard receives one slice from the gate part and the corresponding slice from the up part. - - For instance: - • Shard 0 gets: [ Gate Slice 0, Up Slice 0 ] = [ G0, G1, G2, G3, U0, U1, U2, U3 ] - • Shard 1 gets: [ Gate Slice 1, Up Slice 1 ] = [ G4, G5, G6, G7, U4, U5, U6, U7 ] - • … and so on. - - This ensures that each shard receives an equal portion of both gate and up projections, maintaining consistency across tensor parallelism. - """ - slice_ = param - total_size = empty_param.shape[dim] - world_size = device_mesh.size() - block_sizes = _blocks_to_block_sizes(total_size=total_size, blocks=2) - - tensors_slices = [] - block_offset = 0 - for block_size in block_sizes: - shard_block_size = block_size // world_size - start = rank * shard_block_size - stop = (rank + 1) * shard_block_size - tensors_slices += range(block_offset + start, block_offset + stop) - block_offset += block_size - - slice_dtype = slice_.get_dtype() - # Handle F8_E4M3 dtype by converting to float16 before slicing - # Without upcasting, the slicing causes : RuntimeError: "index_cpu" not implemented for 'Float8_e4m3fn' - casted = False - if slice_dtype == "F8_E4M3" or slice_dtype == "F8_E5M2": - slice_ = slice_[...].to(torch.float16) - casted = True - - if dim == 0: - tensor = slice_[tensors_slices, ...] - elif dim == 1 or dim == -2: - tensor = slice_[:, tensors_slices, ...] - elif dim == 2 or dim == -1: - tensor = slice_[..., tensors_slices] - else: - raise ValueError(f"Unsupported dim {dim}, only dim 0, 1 or 2 are supported") - - if casted: - return tensor - else: - return tensor.to(str_to_dtype[slice_dtype]) - - -def repack_weights( - packed_parameter: torch.Tensor, - sharded_dim: int, # The dimension index in the global tensor that was sharded - world_size: int, - num_blocks: int = 2, -) -> torch.Tensor: - """ - Reorders a tensor that was reconstructed from sharded packed weights into its canonical packed format. - - For example, if a weight was packed (e.g., gate_proj and up_proj) and then sharded, - DTensor.full_tensor() might produce an interleaved layout like [G0, U0, G1, U1, ...] - along the sharded dimension. This function reorders it to [G0, G1, ..., U0, U1, ...]. - This is an inverse operation to get_packed_weights. - - Args: - reconstructed_tensor: The tensor reconstructed from DTensor (e.g., via .full_tensor().contiguous()). - sharded_dim: The dimension index in the reconstructed_tensor that was originally sharded. - world_size: The tensor parallel world size. - num_packed_projs: The number of projections that were packed together (e.g., 2 for gate_up_proj). - - Returns: - The reordered tensor in canonical packed format. - """ - - if num_blocks != 2: - raise ValueError( - "Num blocks different from 2 is not supported yet. This is most likely a bug in your implementation as we only pack gate and up projections together." - ) - - actual_sharded_dim = sharded_dim if sharded_dim >= 0 else sharded_dim + packed_parameter.ndim - total_size_on_sharded_dim = packed_parameter.shape[actual_sharded_dim] - original_block_size_on_dim = total_size_on_sharded_dim // num_blocks - shard_chunk_size = original_block_size_on_dim // world_size - - prefix_shape = packed_parameter.shape[:actual_sharded_dim] - suffix_shape = packed_parameter.shape[actual_sharded_dim + 1 :] - - tensor_view = packed_parameter.view( - *prefix_shape, - world_size, - num_blocks, - shard_chunk_size, - *suffix_shape, - ) - - # Permute to bring num_packed_projs first, then world_size, then shard_chunk_size - # This groups all chunks of G together, then all chunks of U together. - # Target order of these middle dimensions: (num_packed_projs, world_size, shard_chunk_size) - # Current order of view's middle dimensions: (world_size, num_packed_projs, shard_chunk_size) - # Absolute indices of the dimensions to be permuted (world_size, num_packed_projs) - axis_ws_abs = len(prefix_shape) - axis_npp_abs = len(prefix_shape) + 1 - - permute_order = list(range(tensor_view.ndim)) - permute_order[axis_ws_abs], permute_order[axis_npp_abs] = permute_order[axis_npp_abs], permute_order[axis_ws_abs] - - tensor_permuted = tensor_view.permute(*permute_order) - - # Reshape back to the original tensor's ndim, with the sharded dimension now correctly ordered as [G_all, U_all]. - # The final shape should be the same as reconstructed_tensor. - final_ordered_tensor = tensor_permuted.reshape_as(packed_parameter) - - return final_ordered_tensor - - -def get_tensor_shard(param, empty_param, device_mesh, rank, dim, tensor_idx: int | None = None): - """ - Generalized tensor sharding across a multi-dimensional device mesh. - Extract only the fraction of the parameter owned by the given `rank` when the parameter would have gone sharding at provided `dim`. - Extraction follows the pytorch `Shard` placement so that sharding and materializing back to full tensor follows `Shard` semantics. - `Shard` follows torch.chunk style sharding of the tensor. We demonstrate some cases below on how sharding happens including some edge cases - such as some ranks having an empty tensor as shard. Below implementation is robut to all these cases. - - Case (1) - empty_param (16, 5120, 8190) - dim 0 - device_mesh.size() 4 - rank 0 gets (4, 5120, 8190) (0 ... 4, 5120, 8190) - rank 1 gets (4, 5120, 8190) (4 ... 8, 5120, 8190) - rank 2 gets (4, 5120, 8190) (8 ... 12, 5120, 8190) - rank 3 gets (4, 5120, 8190) (12 ... 16, 5120, 8190) - - Case (2) - empty_param (16, 5120, 8190) - dim 0 - device_mesh.size() 14 - rank 0 gets (2, 5120, 8190) (0 ... 2, 5120, 8190) - rank 1 gets (2, 5120, 8190) (2 ... 4, 5120, 8190) - rank 2 gets (2, 5120, 8190) (4 ... 6, 5120, 8190) - rank 3 gets (2, 5120, 8190) (6 ... 8, 5120, 8190) - rank 4 gets (2, 5120, 8190) (8 ... 10, 5120, 8190) - rank 5 gets (2, 5120, 8190) (10 ... 12, 5120, 8190) - rank 6 gets (2, 5120, 8190) (12 ... 14, 5120, 8190) - rank 7 gets (2, 5120, 8190) (14 ... 16, 5120, 8190) - rank 8 gets (0, 5120, 8190) - rank 9 gets (0, 5120, 8190) - rank 10 gets (0, 5120, 8190) - rank 11 gets (0, 5120, 8190) - rank 12 gets (0, 5120, 8190) - rank 13 gets (0, 5120, 8190) - - Case (3) - empty_param (16, 5120, 8190) - dim 0 - device_mesh.size() 3 - rank 0 gets (6, 5120, 8190) (0 ... 6, 5120, 8190) - rank 1 gets (6, 5120, 8190) (6 ... 12, 5120, 8190) - rank 2 gets (4, 5120, 8190) (12 ... 16, 5120, 8190) - - In case (2), empty shards are returned with appropriate dimension to allow for operations to work smoothly. - Args: - param (torch.Tensor): The tensor to shard. - empty_param (torch.Tensor): A tensor used for shape reference. - device_mesh (torch.Tensor): Shape [d_0, ..., d_n] representing the mesh. - rank (int): Global rank of the current process/device. - dim (int): Dimension along which to shard the tensor. - """ - param_dim = empty_param.ndim - mesh_shape = device_mesh.shape - world_size = reduce(operator.mul, mesh_shape) - # Get param shape: works for both torch.Tensor and safetensors TensorInfo - param_shape = list(param.shape) if isinstance(param, torch.Tensor) else param.get_shape() - if dim < 0: - dim = param_dim + dim - if empty_param.dim() == 3 and dim == 1 and len(param_shape) == 2: - dim = 0 - elif empty_param.dim() == 3 and dim == 2 and len(param_shape) == 2: - dim = 1 - - shard_size = math.ceil(param_shape[dim] / world_size) - start = rank * shard_size - end = min(start + shard_size, param_shape[dim]) - - if dim >= param_dim: - raise ValueError(f"dim {dim} is out of bounds for tensor of dimension {param_dim}") - - if rank >= world_size: - raise ValueError(f"Rank {rank} is out of bounds for mesh size {world_size}") - - # we have the full tensor not 1 part of it. - # in that case, we just assume that the weight was properly saved - # and thus because we TP if the layer is colwise it should not use this. Layer should be packed_colwise - # to inform that it needs to read form a packed tensor. It will also take care of the module list thingy. - # here we take care of potential chunking / layer split / layer chunking. - # The only "hard" case is? if we collect q,k,v -> merge it into qkv. In that case - # actually we still shard dim=0 does not change - # so only case is if the dim of the empty param is 3 and the shard dim is 0 -> we put the - # tensor on a certain device (with the input tensor_index) - if tensor_idx is not None and empty_param.dim() == 3 and dim == 0 and len(param_shape) == 2: - # special case we don't "shard" just send this entire tensor to the correct rank. - if start <= tensor_idx < end: - # this tensor does need to be materialized on this device: - return param[:] - else: - return torch.empty([], dtype=torch.int64, device=rank) - - slice_indices = [slice(None)] * len(param_shape) - - if start < param_shape[dim]: - slice_indices[dim] = slice(start, end) - param = param[tuple(slice_indices)] - if isinstance(param, list): # TODO handle the modulelist case! - param = [p[:] for p in param] - return param - - param_shape[dim] = 0 - return torch.empty(tuple(param_shape), dtype=torch.int64) # empty allocates memory.... - - -def _split_along_last_dim(x, world_size): - """Split tensor along last dimension into world_size chunks.""" - return torch.chunk(x, world_size, dim=-1) - - -# ============================================================================= -# Distributed Communication Primitives -# ============================================================================= -# -# Naming convention: -# - Functions describe their FORWARD behavior -# - Backward behavior is the "conjugate" operation for gradient flow -# -# Available operations: -# ┌────────────────────┬─────────────────────┬─────────────────────┐ -# │ Function │ Forward │ Backward │ -# ├────────────────────┼─────────────────────┼─────────────────────┤ -# │ all_reduce │ all-reduce (sum) │ identity │ -# │ all_reduce_backward│ identity │ all-reduce (sum) │ -# │ all_gather │ all-gather │ split (local chunk) │ -# │ split │ split (local chunk) │ all-gather │ -# │ reduce_scatter │ reduce-scatter │ all-gather │ -# └────────────────────┴─────────────────────┴─────────────────────┘ -# =================== - - -class _AllReduceBackward(torch.autograd.Function): - """Identity forward, all-reduce backward. Used before colwise layers (f in Megatron).""" - - @staticmethod - def forward(ctx, x, device_mesh): - ctx.device_mesh = device_mesh - return x - - @staticmethod - def backward(ctx, grad_output): - device_mesh = ctx.device_mesh - if device_mesh.size() == 1: - return grad_output, None - grad_output = grad_output.contiguous() - dist.all_reduce(grad_output, op=dist.ReduceOp.SUM, group=device_mesh.get_group()) - return grad_output, None - - -class _AllReduceForward(torch.autograd.Function): - """All-reduce forward, identity backward. Used after rowwise layers (g in Megatron).""" - - @staticmethod - def forward(ctx, x, device_mesh): - if device_mesh.size() == 1: - return x - dist.all_reduce(x, op=dist.ReduceOp.SUM, group=device_mesh.get_group()) - return x - - @staticmethod - def backward(ctx, grad_output): - return grad_output, None - - -class _AllGather(torch.autograd.Function): - """All-gather forward, split backward. Gathers sharded outputs.""" - - @staticmethod - def forward(ctx, x, device_mesh): - ctx.device_mesh = device_mesh - world_size = device_mesh.size() - - if world_size == 1: - return x - - last_dim = x.dim() - 1 - rank = device_mesh.get_local_rank() - group = device_mesh.get_group() - - x = x.contiguous() - tensor_list = [torch.empty_like(x) for _ in range(world_size)] - tensor_list[rank] = x - dist.all_gather(tensor_list, x, group=group) - return torch.cat(tensor_list, dim=last_dim).contiguous() - - @staticmethod - def backward(ctx, grad_output): - device_mesh = ctx.device_mesh - world_size = device_mesh.size() - - if world_size == 1: - return grad_output, None - - rank = device_mesh.get_local_rank() - chunks = _split_along_last_dim(grad_output, world_size) - return chunks[rank].contiguous(), None - - -class _Split(torch.autograd.Function): - """Split forward, all-gather backward. Scatters replicated input.""" - - @staticmethod - def forward(ctx, x, device_mesh): - ctx.device_mesh = device_mesh - world_size = device_mesh.size() - - if world_size == 1: - return x - - rank = device_mesh.get_local_rank() - chunks = _split_along_last_dim(x, world_size) - return chunks[rank].contiguous() - - @staticmethod - def backward(ctx, grad_output): - device_mesh = ctx.device_mesh - world_size = device_mesh.size() - - if world_size == 1: - return grad_output, None - - last_dim = grad_output.dim() - 1 - rank = device_mesh.get_local_rank() - group = device_mesh.get_group() - - grad_output = grad_output.contiguous() - tensor_list = [torch.empty_like(grad_output) for _ in range(world_size)] - tensor_list[rank] = grad_output - dist.all_gather(tensor_list, grad_output, group=group) - return torch.cat(tensor_list, dim=last_dim).contiguous(), None - - -class _ReduceScatter(torch.autograd.Function): - """Reduce-scatter forward, all-gather backward. For sequence parallel.""" - - @staticmethod - def forward(ctx, x, device_mesh): - ctx.device_mesh = device_mesh - world_size = device_mesh.size() - - if world_size == 1: - return x - - last_dim = x.dim() - 1 - group = device_mesh.get_group() - - input_chunks = list(x.chunk(world_size, dim=last_dim)) - output_shape = list(x.shape) - output_shape[last_dim] //= world_size - output = torch.empty(output_shape, dtype=x.dtype, device=x.device) - - dist.reduce_scatter(output, input_chunks, op=dist.ReduceOp.SUM, group=group) - return output - - @staticmethod - def backward(ctx, grad_output): - device_mesh = ctx.device_mesh - world_size = device_mesh.size() - - if world_size == 1: - return grad_output, None - - last_dim = grad_output.dim() - 1 - rank = device_mesh.get_local_rank() - group = device_mesh.get_group() - - grad_output = grad_output.contiguous() - tensor_list = [torch.empty_like(grad_output) for _ in range(world_size)] - tensor_list[rank] = grad_output - dist.all_gather(tensor_list, grad_output, group=group) - return torch.cat(tensor_list, dim=last_dim).contiguous(), None - - -# ============================================================================= -# Convenience wrappers -# ============================================================================= - - -def all_reduce_backward(x, device_mesh): - """Identity forward, all-reduce backward. Use before colwise layers.""" - return _AllReduceBackward.apply(x, device_mesh) - - -def all_reduce_forward(x, device_mesh): - """All-reduce forward, identity backward. Use after rowwise layers.""" - return _AllReduceForward.apply(x, device_mesh) - - -def all_gather(x, device_mesh): - """All-gather forward, split backward.""" - return _AllGather.apply(x, device_mesh) - - -def split(x, device_mesh): - """Split forward, all-gather backward.""" - return _Split.apply(x, device_mesh) - - -def reduce_scatter(x, device_mesh): - """Reduce-scatter forward, all-gather backward.""" - return _ReduceScatter.apply(x, device_mesh) - - -def distribute_module( - module: nn.Module, - device_mesh=None, - input_fn=None, - output_fn=None, -) -> nn.Module: - """ - Copy pasted from torch's function but we remove the communications (partitioning) - as well as buffer registering that is similarly not efficient. - """ - if input_fn is not None: - module.register_forward_pre_hook(lambda mod, inputs: input_fn(mod, inputs, device_mesh)) - if output_fn is not None: - module.register_forward_hook(lambda mod, inputs, outputs: output_fn(mod, outputs, device_mesh)) - return module - - -class TensorParallelLayer: - """General tensor parallel layer for transformers""" - - device_mesh = None - rank = None - empty_param = None - - def __init__(self, device_mesh=None, rank=None, empty_param=None): - self.rank = rank - self.device_mesh = device_mesh - self.empty_param = empty_param - - def _prepare_input_fn(self, mod, inputs, device_mesh): - raise NotImplementedError - - def _prepare_output_fn(self, mod, outputs, device_mesh): - raise NotImplementedError - - def shard_tensor( - self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None - ) -> torch.Tensor: - raise NotImplementedError - - def prepare_module_tp(self, module: nn.Module, device_mesh, **kwargs) -> nn.Module: - distribute_module( - module, - device_mesh, - self._prepare_input_fn, - self._prepare_output_fn, - ) - - def get_expected_sharded_shape(self, full_shape: tuple[int, ...] | torch.Size) -> tuple[int, ...]: - """ - Compute the expected shape after TP sharding for a given full shape. - - Args: - full_shape: The full (unsharded) parameter shape - - Returns: - The expected sharded shape for this rank - """ - # Default: no sharding, return full shape - return tuple(full_shape) - - def update_module_attributes(self, module: nn.Module): - """ - Update module attributes (e.g. in_features, out_features) to reflect sharded dimensions. - - Args: - module: The module to update - - Returns: - None, update the module in-place - """ - pass - - -class ColwiseParallel(TensorParallelLayer): - """ - Column-wise parallel: weight is sharded on dim -2 (output features). - Forward: input replicated -> output sharded on last dim. - If gather_output=True, output is all-gathered to produce full tensor. - """ - - def __init__(self, gather_output: bool = False, **kwargs): - super().__init__(**kwargs) - self.gather_output = gather_output - - def _prepare_input_fn(self, mod, inputs, device_mesh): - input_tensor = inputs[0] if inputs else inputs - return all_reduce_backward(input_tensor, device_mesh) - - def _prepare_output_fn(self, mod, outputs, device_mesh): - if self.gather_output: - return all_gather(outputs, device_mesh) - return outputs - - def shard_tensor( - 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) - else: - parameter = get_tensor_shard(param, self.empty_param, self.device_mesh, self.rank, -2) - return parameter.to(device=device, dtype=dtype) - - def get_expected_sharded_shape(self, full_shape: tuple[int, ...] | torch.Size) -> tuple[int, ...]: - world_size = self.device_mesh.size() - shape = list(full_shape) - # Colwise shards dim -2, but 1D tensors (bias) shard on dim -1 - dim = -1 if len(shape) == 1 else -2 - dim = len(shape) + dim if dim < 0 else dim - shard_size = math.ceil(shape[dim] / world_size) - start = self.rank * shard_size - end = min(start + shard_size, shape[dim]) - shape[dim] = end - start - return tuple(shape) - - def update_module_attributes(self, module: nn.Module): - # If we gather the output, the output dimension of the module is not sharded, so no need to update out_features. - # Otherwise, we need to update out_features to reflect the sharded dimension. - if not self.gather_output and hasattr(module, "out_features"): - module.out_features = self.get_expected_sharded_shape((module.out_features,))[0] - - -class ReplicatedWithGradAllReduce(TensorParallelLayer): - """ - Replicated parameter with gradient all-reduce. - - For parameters like q_norm/k_norm that sit between colwise and rowwise - layers. The parameter is replicated (not sharded), but its gradient - accumulates from local heads only in TP mode. This class registers a - backward hook to all-reduce the parameter gradient. - """ - - def _prepare_input_fn(self, mod, inputs, device_mesh): - return inputs - - def _prepare_output_fn(self, mod, outputs, device_mesh): - return outputs - - def shard_tensor(self, param, tensor_idx=None, device=None, dtype=None): - return param[...].to(device=device, dtype=dtype) - - def prepare_module_tp(self, module, device_mesh, **kwargs): - # Use a module-level backward hook (not param.register_hook) because parameters are replaced during weight loading after this method runs. - # Module hooks survive parameter replacement. - def _backward_hook(mod, grad_input, grad_output, mesh=device_mesh): - for param in mod.parameters(): - if param.grad is not None: - all_reduce_forward(param.grad, mesh) - - module.register_full_backward_hook(_backward_hook) - - -class MlaKvAProjParallel(TensorParallelLayer): - """ - For MLA attention used in DeepSeek-V2 style models (deepseek_v2, longcat_flash, glm_moe_dsa, glm4_moe_lite): - kv_a_proj_with_mqa output is [kv_lora_rank + qk_rope_head_dim] (can have different naming but important thing - to understand is that it is split) - Example below (from modeling_longcat_flash.py): - - kv_a_proj_with_mqa - | - split - / \ - k_pass k_rot <-- "bypasses kv_b_proj" - | | (goes straight to attention, - kv_a_layernorm | never touches kv_b_proj) - | | - kv_b_proj | - (colwise) | - | | - k_pass k_rot - \\ / - cat - | - key_states - - k_pass is passed to kv_b_proj (colwise) which has built-in all_reduce_backward so we don't have a partial gradient for it. - However, k_rot goes straight to attention, never touches kv_b_proj. So we need to average gradient across all ranks otherwise we only get gradient for one rank (partial gradient). - """ - - def _prepare_output_fn(self, mod, output, device_mesh): - if not hasattr(mod.config, "qk_rope_head_dim"): - raise AttributeError( - f"Config for {type(mod).__name__} does not have `qk_rope_head_dim`. " - "MlaKvAProjParallel requires `qk_rope_head_dim` to be defined in the model config. " - "Please add it to the model's config or update the TP plan mapping." - ) - rope_dim = mod.config.qk_rope_head_dim - pass_output, rope_output = output.split([output.shape[-1] - rope_dim, rope_dim], dim=-1) - rope_output = all_reduce_backward(rope_output, device_mesh) - return torch.cat([pass_output, rope_output], dim=-1) - - def shard_tensor(self, param, tensor_idx=None, device=None, dtype=None): - return param[...].to(device=device, dtype=dtype) - - def prepare_module_tp(self, module, device_mesh, config=None, **kwargs): - module.config = config - distribute_module(module, device_mesh, output_fn=self._prepare_output_fn) - - -class RowwiseParallel(TensorParallelLayer): - """ - Row-wise parallel: weight is sharded on dim -1 (input features). - Forward: input (optionally split) -> output partial -> all-reduce to replicate. - - Args: - split_input: If True, splits replicated input before matmul. Use when input - comes from a non-parallelizable operation (chunk/slice). - Default False (expects pre-sharded input from colwise layer). - """ - - def __init__(self, split_input: bool = False, **kwargs): - super().__init__(**kwargs) - self.split_input = split_input - - def _prepare_input_fn(self, mod, inputs, device_mesh): - if hasattr(mod, "bias") and mod.bias is not None: - mod._bias = mod.bias - mod.bias = None - - input_tensor = inputs[0] if inputs else inputs - - if self.split_input: - # Input is replicated, split it to match sharded weight - return split(input_tensor, device_mesh) - return input_tensor - - def _prepare_output_fn(self, mod, outputs, device_mesh): - outputs = all_reduce_forward(outputs, device_mesh) - if hasattr(mod, "_bias") and mod._bias is not None: - outputs = outputs + mod._bias - return outputs - - def shard_tensor( - 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: - parameter = param[...] - else: - parameter = get_tensor_shard(param, self.empty_param, self.device_mesh, self.rank, -1) - return parameter.to(device=device, dtype=dtype) - - def get_expected_sharded_shape(self, full_shape: tuple[int, ...] | torch.Size) -> tuple[int, ...]: - # 1D tensors (bias) are NOT sharded in rowwise - if len(full_shape) == 1: - return tuple(full_shape) - world_size = self.device_mesh.size() - shape = list(full_shape) - dim = -1 - dim = len(shape) + dim if dim < 0 else dim - shard_size = math.ceil(shape[dim] / world_size) - start = self.rank * shard_size - end = min(start + shard_size, shape[dim]) - shape[dim] = end - start - return tuple(shape) - - def update_module_attributes(self, module: nn.Module): - if hasattr(module, "in_features"): - # To fall in the 2D case in get_expected_sharded_shape, - # otherwise it will be treated as 1D and not sharded - shape = (1, module.in_features) - module.in_features = self.get_expected_sharded_shape(shape)[1] - - -class PackedColwiseParallel(ColwiseParallel): - """Packed column-wise parallel for fused weights like gate_up_proj.""" - - def shard_tensor( - 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) - else: - expected_shape = self.get_expected_sharded_shape(self.empty_param.shape) - if dim < len(expected_shape): - # Input is unpacked (e.g., gate_proj that will be concatenated to gate_up_proj) - # Use regular tensor shard - concatenation will happen after - parameter = get_tensor_shard(param, self.empty_param, self.device_mesh, self.rank, -2) - else: - # Input is already packed, use packed sharding - parameter = get_packed_weights(param, self.empty_param, self.device_mesh, self.rank, -2) - return parameter.to(device=device, dtype=dtype) - - -class PackedRowwiseParallel(RowwiseParallel): - """Packed row-wise parallel for fused weights like gate_up_proj.""" - - def shard_tensor( - 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: - parameter = param[...] - else: - # Check if input tensor is unpacked (shape mismatch with expected packed size) - # This happens when using MergeModulelist + Concatenate for fused weights like gate_up_proj - param_shape = param.shape if isinstance(param, torch.Tensor) else param.get_shape() - expected_packed_dim = self.empty_param.shape[-1] if self.empty_param.dim() >= 1 else 0 - actual_dim = param_shape[-1] if len(param_shape) >= 1 else 0 - - if actual_dim < expected_packed_dim: - # Input is unpacked, use regular tensor shard - parameter = get_tensor_shard(param, self.empty_param, self.device_mesh, self.rank, -1) - else: - # Input is already packed, use packed sharding - parameter = get_packed_weights(param, self.empty_param, self.device_mesh, self.rank, -1) - return parameter.to(device=device, dtype=dtype) - - -class EmbeddingParallel(TensorParallelLayer): - """EmbeddingParallel: shards embedding table, handles masked lookups for vocab parallelism.""" - - def __init__(self, *, embedding_dim_sharding: int = 0, **kwargs): - super().__init__(**kwargs) - self.embedding_dim_sharding = embedding_dim_sharding - - def _prepare_input_fn(self, mod, inputs, device_mesh): - input_tensor = inputs[0] if inputs else inputs - - # For vocab-parallel (dim 0), we need to handle masking and offsetting - if self.embedding_dim_sharding == 0: - rank = device_mesh.get_local_rank() - - # Get vocab range for this rank - # Use weight.shape[0] to get the actual local (sharded) size, not num_embeddings - # which may not be updated after sharding - per_partition_size = mod.weight.shape[0] - vocab_start_index = rank * per_partition_size - vocab_end_index = vocab_start_index + per_partition_size - - # Build mask for out-of-vocabulary tokens - input_mask = (input_tensor < vocab_start_index) | (input_tensor >= vocab_end_index) - mod._input_mask = input_mask - - # Offset input to local indices and mask invalid ones - masked_input = input_tensor.clone() - vocab_start_index - masked_input[input_mask] = 0 # Set to valid local index - - return masked_input - - return input_tensor - - def _prepare_output_fn(self, mod, outputs, device_mesh): - # For vocab-parallel (dim 0), zero out embeddings for out-of-range tokens before all-reduce - if self.embedding_dim_sharding == 0 and hasattr(mod, "_input_mask"): - input_mask = mod._input_mask - # Use multiplication instead of in-place assignment to preserve gradients - mask_expanded = input_mask.unsqueeze(-1).expand_as(outputs) - outputs = outputs * (~mask_expanded).to(outputs.dtype) - del mod._input_mask - - return all_reduce_forward(outputs, device_mesh) - - def shard_tensor( - 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) - else: - parameter = get_tensor_shard( - param, - self.empty_param, - self.device_mesh, - self.rank, - self.embedding_dim_sharding, - ) - return parameter.to(device=device, dtype=dtype) - - def get_expected_sharded_shape(self, full_shape: tuple[int, ...] | torch.Size) -> tuple[int, ...]: - world_size = self.device_mesh.size() - shape = list(full_shape) - # EmbeddingParallel shards on self.embedding_dim_sharding (default 0) - # 1D tensors (bias) shard on dim -1 - dim = -1 if len(shape) == 1 else self.embedding_dim_sharding - dim = len(shape) + dim if dim < 0 else dim - shard_size = math.ceil(shape[dim] / world_size) - start = self.rank * shard_size - end = min(start + shard_size, shape[dim]) - shape[dim] = end - start - return tuple(shape) - - def update_module_attributes(self, module: nn.Module): - if hasattr(module, "num_embeddings") and self.embedding_dim_sharding == 0: - module.num_embeddings = self.get_expected_sharded_shape((module.num_embeddings,))[0] - if hasattr(module, "embedding_dim") and self.embedding_dim_sharding == 1: - module.embedding_dim = self.get_expected_sharded_shape((module.embedding_dim,))[0] - - -class SequenceParallel(TensorParallelLayer): - """ - Sequence Parallel: input/output sharded on sequence dimension. - Weights are replicated. - """ - - def __init__(self, sequence_dim: int = 1, use_local_output: bool = False, use_dtensor=False, **kwargs): - super().__init__(**kwargs) - self.sequence_dim = sequence_dim - - def _prepare_input_fn(self, mod, inputs, device_mesh): - input_tensor = inputs[0] if inputs else inputs - # For sequence parallel, input is sharded on sequence dim - # All-gather for the layer, then reduce-scatter after - return all_gather(input_tensor, device_mesh) - - def _prepare_output_fn(self, mod, outputs, device_mesh): - return reduce_scatter(outputs, device_mesh) - - def shard_tensor( - self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None - ) -> torch.Tensor: - return param[...].to(device=device, dtype=dtype) - - -class GroupedGemmParallel(TensorParallelLayer): - """ - Applies Expert Parallelism to MoE experts by loading the correct experts on each device. - """ - - def __init__(self, **kwargs): - super().__init__(**kwargs) - - def shard_tensor( - 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} % {self.device_mesh.size()} != 0" - ) - local_num_experts = global_num_experts // self.device_mesh.size() - shard_size = local_num_experts - start = self.rank * shard_size - end = (self.rank + 1) * shard_size - # special case we don't "shard" just send this entire tensor to the correct rank. - shape = param.get_shape() if not isinstance(param, torch.Tensor) else param.shape - if tensor_idx is not None and start <= tensor_idx < end: - # this tensor does need to be materialized on this device: - return param[:].to(device=device) - elif tensor_idx is None: # a bias or a weight, but already merged - return param[start:end].to(device=device, dtype=dtype) - elif len(shape) >= 1 and tensor_idx is not None: - return None - else: # bias case - return param[:].to(device=device, dtype=dtype) - - def get_expected_sharded_shape(self, full_shape: tuple[int, ...] | torch.Size) -> tuple[int, ...]: - # GroupedGemm shards on dim 0 (experts dimension) - world_size = self.device_mesh.size() - shape = list(full_shape) - local_num_experts = shape[0] // world_size - shape[0] = local_num_experts - return tuple(shape) - - 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] - - -class RouterParallel(TensorParallelLayer): - """ - Allows to reshape the router scores to support running expert parallel. - """ - - def __init__(self, **kwargs): - super().__init__(**kwargs) - - def _prepare_input_fn(self, mod, inputs, device_mesh): - return inputs[0] if inputs else inputs - - def _prepare_output_fn(self, mod, outputs, device_mesh): - """ - Remap global expert indices to local and zero out non-local scores. - - Example: 4 tokens, top_k=4, 128 experts, EP=8. num_local_experts = 128/8 = 16. - - Router produces (all ranks see the same values): - router_scores: (4, 4) — top-k routing weights - router_indices: (4, 4) — global expert IDs - [ 52, 42, 119, 67], - [102, 89, 61, 40], - [ 82, 103, 4, 34], - [ 93, 23, 109, 11], - - Each index maps to a rank: index // 16 gives the owning rank. - [3, 2, 7, 4], - [6, 5, 3, 2], - [5, 6, 0, 2], - [5, 1, 6, 0], - - For rank 0 (owns experts 0-15), we remap local indices with fmod and - fill non-local with sentinel=16 (used for one_hot masking): - router_indices (rank 0): - [ 16, 16, 16, 16], - [ 16, 16, 16, 16], - [ 16, 16, 4, 16], - [ 16, 16, 16, 11], - - Scores for non-local experts are zeroed out via masked_fill: - router_scores (rank 0): - [0.0, 0.0, 0.0, 0.0], - [0.0, 0.0, 0.0, 0.0], - [0.0, 0.0, 0.3, 0.0], ← only expert 4 (local) keeps its score - [0.0, 0.0, 0.0, 0.1], ← only expert 11 (local) keeps its score - - both router_scores and router_indices stay (seq, top_k) shape. - They are paired element-wise: scores[i] is the weight for indices[i]. - All expert forward implementations (grouped_mm, batched_mm, eager) flatten - both with reshape(-1) and rely on this pairing. Changing the shape of one - without the other breaks routing! - - Each rank believes it is alone and computes only its part of the hidden states. - The sentinel index (num_local_experts) is skipped by one_hot encoding or clamped - + masked in grouped_mm/batched_mm. After the expert forward, an all_reduce sums - partial outputs across EP ranks to produce the full result. - """ - ep_rank, ep_size = device_mesh.get_local_rank(), device_mesh.size() - num_experts = getattr(mod, "num_experts", None) - if num_experts is None: - num_experts = getattr(getattr(mod, "config", None), "num_experts", None) - if num_experts is None: - raise AttributeError(f"Router module {type(mod).__name__} is missing num_experts and config.num_experts") - - if num_experts % ep_size != 0: - raise ValueError( - f"The number of experts must be divisible by number of ep_size: {num_experts} % {ep_size} != 0" - ) - num_local_experts = num_experts // ep_size - router_logits, router_scores, router_indices = outputs - non_local_mask = (router_indices // num_local_experts) != ep_rank - router_scores = router_scores.masked_fill(non_local_mask, 0.0) - router_indices = router_indices.masked_fill(non_local_mask, -1) - # As -1 % 1 is 0, we can only use mask fill when num_local_experts is 1 - if num_local_experts > 1: - router_indices = torch.fmod(router_indices, num_local_experts) - else: - router_indices = router_indices.masked_fill(router_indices > 0, 0).masked_fill(router_indices < 0, -1) - router_indices = router_indices.masked_fill(router_indices == -1, num_local_experts) - return router_logits, router_scores, router_indices - - def shard_tensor( - self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None - ) -> torch.Tensor: - return param[...].to(device=device, dtype=dtype) - - -class MoeTensorParalellExperts(TensorParallelLayer): - """ - Note: For tensor parallel, the MoEExpertsParallel TP layer handles gradient sync: - - all_reduce_backward on hidden_states (for colwise gate_up_proj gradient) - - all_reduce_backward on top_k_weights (for router gradient) - - all_reduce_forward on output (for partial expert outputs) - """ - - def __init__(self, **kwargs): - super().__init__(**kwargs) - - def _prepare_input_fn(self, mod, inputs, device_mesh): - # inputs = (hidden_states, top_k_index, top_k_weights) - hidden_states = inputs[0] - top_k_index = inputs[1] - top_k_weights = inputs[2] - - # all_reduce_backward on hidden_states for correct colwise (gate_up_proj) gradient - hidden_states = all_reduce_backward(hidden_states, device_mesh) - - # all_reduce_backward on routing weights for correct router gradient - # This is needed because ∂L/∂routing_weights = ∂L/∂output * partial_expert_output - # and partial_expert_output is different on each GPU before all-reduce - top_k_weights = all_reduce_backward(top_k_weights, device_mesh) - - return (hidden_states, top_k_index, top_k_weights) - - def _prepare_output_fn(self, mod, outputs, device_mesh): - # all_reduce_forward to sum partial expert outputs across GPUs - return all_reduce_forward(outputs, device_mesh) - - def shard_tensor( - self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None - ) -> torch.Tensor: - # This class doesn't shard tensors - sharding is handled by packed_colwise/rowwise - # on the individual weight tensors (gate_up_proj/down_proj) - return param[...].to(device=device, dtype=dtype) - - -class MoeIdentityExpertParallel(TensorParallelLayer): - """ - TP class for zero/identity experts in MoE layers. - - Under TP, the parent MoeTensorParalellExperts does all_reduce_forward (sum) - on the expert module output. Identity experts produce the same output on - every rank, so the sum gives world_size * output. This class divides the - input by world_size to compensate. - """ - - def _prepare_input_fn(self, mod, inputs, device_mesh): - input_tensor = inputs[0] if inputs else inputs - # TODO(fmom): when 2D-device mesh, need to select a //-ism axis to divide the input tensor by. - return input_tensor / device_mesh.size() - - def shard_tensor(self, param, tensor_idx=None, device=None, dtype=None): - return param[...].to(device=device, dtype=dtype) - - def prepare_module_tp(self, module, device_mesh, **kwargs): - distribute_module(module, device_mesh, input_fn=self._prepare_input_fn) - - -class ParallelInterface(GeneralInterface): - # Class instance object, so that a call to `register` can be reflected into all other files correctly, even if - # a new instance is created (in order to locally override a given entry) - _global_mapping = ( - { - "embedding_rowwise": EmbeddingParallel(embedding_dim_sharding=0), - "embedding_colwise": EmbeddingParallel(embedding_dim_sharding=1), - "colwise_gather_output": ColwiseParallel(gather_output=True), - "colwise": ColwiseParallel(), - "rowwise": RowwiseParallel(), - "rowwise_split_input": RowwiseParallel(split_input=True), - "packed_colwise": PackedColwiseParallel(), - "packed_rowwise": PackedRowwiseParallel(), - "sequence_parallel": SequenceParallel(), - "grouped_gemm": GroupedGemmParallel(), - "ep_router": RouterParallel(), - "moe_tp_experts": MoeTensorParalellExperts(), - "moe_identity_expert": MoeIdentityExpertParallel(), - "replicated_with_grad_allreduce": ReplicatedWithGradAllReduce(), - "mla_kv_a_proj": MlaKvAProjParallel(), - } - if is_torch_available() and _torch_distributed_available - else {} - ) - - # Map plan names to sharding dimensions for weights - # For weights: colwise shards dim -2, rowwise shards dim -1 - # For embedding: rowwise shards dim 0 (vocab), colwise shards dim -2 (hidden) - plan_to_weight_dim: dict[str, int | None] = { - "colwise": -2, - "colwise_gather_output": -2, - "packed_colwise": -2, - "rowwise": -1, - "rowwise_split_input": -1, - "packed_rowwise": -1, - "embedding_rowwise": 0, - "embedding_colwise": 1, - "sequence_parallel": None, - "replicated_with_grad_allreduce": None, - "mla_kv_a_proj": None, - } - - # Bias sharding: colwise shards bias, rowwise doesn't (bias is replicated and all-reduced) - plan_to_bias_dim: dict[str, int | None] = { - "colwise": -1, - "colwise_gather_output": -1, - "packed_colwise": -1, - "rowwise": None, - "rowwise_split_input": None, - "packed_rowwise": None, - "embedding_rowwise": None, - "embedding_colwise": None, - "sequence_parallel": None, - "replicated_with_grad_allreduce": None, - "mla_kv_a_proj": None, - } - - @classmethod - def register_plan_to_weight_dim(cls, key: str, value: int | None): - cls.plan_to_weight_dim[key] = value - - @classmethod - def register_plan_to_bias_dim(cls, key: str, value: int | None): - cls.plan_to_bias_dim[key] = value - - -ALL_PARALLEL_STYLES: ParallelInterface = ParallelInterface() - - -# ============================================================================= -# High-Level API Functions -# ============================================================================= - - -def gather_full_tensor( - local_tensor: torch.Tensor, shard_dim: int, device_mesh: dist.device_mesh.DeviceMesh -) -> torch.Tensor: - """ - All-gather a sharded tensor along the specified dimension to reconstruct the full tensor. - - Args: - local_tensor: The local shard of the tensor on this rank - shard_dim: The dimension along which the tensor was sharded - device_mesh: The device mesh for distributed communication - - Returns: - The full reconstructed tensor (same on all ranks) - """ - world_size = device_mesh.size() - # In case of TP+DP configuration, the TP group should be used for gathering, not the full DP group - process_group = device_mesh.get_group("tp") if "tp" in (device_mesh.mesh_dim_names or {}) else None - - # Normalize negative dimension - if shard_dim < 0: - shard_dim = local_tensor.ndim + shard_dim - - # Gather all shards - gathered_tensors = [torch.empty_like(local_tensor) for _ in range(world_size)] - dist.all_gather(gathered_tensors, local_tensor.contiguous(), group=process_group) - - # Concatenate along the shard dimension - return torch.cat(gathered_tensors, dim=shard_dim) - - -def gather_state_dict_for_save( - state_dict: dict[str, torch.Tensor], - tp_plan: dict[str, str], - device_mesh, - tp_size: int, -) -> dict[str, torch.Tensor]: - """ - Gather sharded tensors to reconstruct full tensors for saving. - - This function all-gathers each sharded tensor along its shard dimension - to reconstruct the full unsharded tensor for checkpoint saving. - - Args: - state_dict: The model state dict with local sharded tensors - tp_plan: The tensor parallel plan mapping layer patterns to shard styles - device_mesh: The device mesh for distributed communication - tp_size: The tensor parallel world size - - Returns: - State dict with full (gathered) tensors - """ - # Use the global mappings from ParallelInterface (can be extended by users) - plan_to_weight_dim = ALL_PARALLEL_STYLES.plan_to_weight_dim - plan_to_bias_dim = ALL_PARALLEL_STYLES.plan_to_bias_dim - - result = {} - for key, tensor in state_dict.items(): - # Find the matching TP plan for this parameter - param_name = key.rsplit(".", 1)[0] if "." in key else key - param_type = key.rsplit(".", 1)[1] if "." in key else None - generic_param_name = re.sub(r"\d+", "*", param_name) - # Also check the full key for nn.Parameter (e.g., MoE experts without .weight suffix) - generic_full_key = re.sub(r"\d+", "*", key) - - # Check if this parameter has a TP plan - current_plan = None - if generic_full_key in tp_plan: - # Full key match (e.g., "model.layers.*.mlp.experts.gate_up_proj" for MoE experts) - current_plan = tp_plan[generic_full_key] - elif generic_param_name in tp_plan: - current_plan = tp_plan[generic_param_name] - elif "." in generic_param_name: - parent_param_name = generic_param_name.rsplit(".", 1)[0] - if parent_param_name in tp_plan: - current_plan = tp_plan[parent_param_name] - - if current_plan is None or current_plan not in plan_to_weight_dim: - # Not sharded, keep as-is - result[key] = tensor - continue - - # Determine sharding dimension based on param type - if param_type == "bias": - shard_dim = plan_to_bias_dim.get(current_plan) - else: - shard_dim = plan_to_weight_dim.get(current_plan) - - if shard_dim is None: - # Replicated, keep as-is - result[key] = tensor - continue - - # Gather full tensor and handle packed weights repacking - full_tensor = gather_full_tensor(tensor, shard_dim, device_mesh) - if current_plan in ("packed_colwise", "packed_rowwise"): - full_tensor = repack_weights(full_tensor, shard_dim, tp_size, 2) - result[key] = full_tensor.contiguous() - - return result - - -def add_tensor_parallel_hooks_to_module( - model, - module, - current_module_plan, - layer_name, - device_mesh, -): - r""" - This function is called in `PretrainedModel.post_init()`. It is responsible of adding hooks - to the modules of the `model`, based on the `PretrainedModel._tp_plan`. - - This is the place where we add the `pre_forward` and `post_forwards` hooks. These are defined - for each `TensorParallelLayer` as `_prepare_input_fn` and `_prepare_output_fn`. - - Args: - model (`PretrainedModel`): The model containing the modules. - module (`nn.Module`): The current module to which we want to add the hooks. - current_module_plan (`str` or `None`): The tensor parallel plan for the current module, if any. - layer_name (`str`): The qualified name of the current module. - device_mesh (`dist.device_mesh.DeviceMesh`): The device mesh for distributed communication. - - """ - if current_module_plan is not None: - tp_layer = ALL_PARALLEL_STYLES[current_module_plan] - try: - tp_layer.prepare_module_tp(module, device_mesh, config=model.config) - except NotImplementedError as e: - logger.warning( - f"Trying to prepare {layer_name}, but it's not supported. Corresponding module: {module} Fix it's TP " - f"plan: {e}" - ) - - module._hf_tp_plan = current_module_plan - module._hf_device_mesh = device_mesh - module.__repr__ = lambda: f"{module.__repr__()}\nTP Plan: {current_module_plan}" - - -def shard_and_distribute_module( - model, param, empty_param, parameter_name, param_casting_dtype, is_contiguous, rank, device_mesh -): - r""" - This function is called in `from_pretrained` when loading a model's checkpoints. - It receives the pointer to the parameter (or the parameter itself) and takes care of "sharding". - All process run this function, so they just load the partition of the tensor that they require. - - Main uses cases: - - column / rowise parallelism, you just shard all the weights of the layer (weight and bias) - - packed layers: you slice the weights, then shard like above - - custom operation: - - you want to add an all-gather at the end of a local layer. - - you want to have a layer that is isolated from the rest of the world (because torch.DTensor does not work well with `.view` for instance) - - """ - param_name, param_type = parameter_name.rsplit(".", 1) if "." in parameter_name else parameter_name - tp_plan = model.tp_plan or {} - module_to_tp = model.get_submodule(param_name) - rank = int(rank) - current_shard_plan = _get_parameter_tp_plan(parameter_name, tp_plan) - - if dist.get_rank() == 0: - if current_shard_plan is None: - logger.info(f"Tensor sharding plan for {param_name} not found, using default 'replicate' plan.") - else: - logger.info(f"Tensor sharding plan for {param_name}: {current_shard_plan}") - - tp_layer = None - 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.shard_tensor(param, tensor_idx=None, dtype=param_casting_dtype, device=rank) - if is_contiguous: - param = param.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}" - ) - else: - param = param[:].to(param_casting_dtype) - - # SUPER IMPORTANT we have to use setattr - # otherwise loading is crazy slow - if not isinstance(param, torch.nn.Parameter): - param = torch.nn.Parameter(param, requires_grad=empty_param.is_floating_point()) - setattr(module_to_tp, param_type, param) - if tp_layer is not None: - tp_layer.update_module_attributes(module_to_tp) - return param - - -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. - """ - - if tp_plan is None: - return - - generic_keys = {replace_layer_number_by_wildcard(key) for key in expected_keys} - unsharded_layers = set(generic_keys) - unused_rules = tp_plan.copy() - - for key in generic_keys: - param_name = key.rsplit(".", 1)[0] if "." in key else key - generic_param_name = re.sub(r"\d+", "*", param_name) - - if generic_param_name in tp_plan: - unused_rules.pop(generic_param_name, None) - unsharded_layers.discard(key) - elif "." in generic_param_name and (parent_param_name := generic_param_name.rsplit(".", 1)[0]) in tp_plan: - unused_rules.pop(parent_param_name, None) - unsharded_layers.discard(key) - - if len(unused_rules) > 0: - logger.warning(f"The following TP rules were not applied on any of the layers: {unused_rules}") - if len(unsharded_layers) > 0: - logger.warning(f"The following layers were not sharded: {', '.join(unsharded_layers)}") - - -def distribute_model(model, tp_plan, distributed_config, device_mesh, tp_size): - """Distribute a model according to the TP plan.""" - model._tp_size = tp_size - model._device_mesh = device_mesh - if distributed_config is not None: - if isinstance(distributed_config, dict): - distributed_config = DistributedConfig.from_dict(distributed_config) - model.config.distributed_config = distributed_config - # Set the new requested tp_plan on the model - if isinstance(tp_plan, dict): - model.tp_plan = tp_plan - model_plan = model.tp_plan - if model_plan is not None and _torch_distributed_available: - for v in model_plan.values(): - if v not in ALL_PARALLEL_STYLES: - raise ValueError(f"Unsupported tensor parallel style {v}. Supported styles are {ALL_PARALLEL_STYLES}") - for name, module in model.named_modules(): - if not getattr(module, "_is_hooked", False): - plan = _get_parameter_tp_plan(parameter_name=name, tp_plan=model_plan, is_weight=False) - add_tensor_parallel_hooks_to_module( - model, - module, - plan, - name, - device_mesh, - ) - module._is_hooked = True - return model diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 7e27fc5a20ab..b17250bd6f45 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -39,6 +39,7 @@ 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 from torch.distributions import constraints from torch.utils.checkpoint import checkpoint @@ -52,9 +53,22 @@ revert_weight_conversion, ) from .distributed import DistributedConfig +from .distributed.fsdp import is_fsdp_enabled +from .distributed.sharding_utils import _dtensor_from_local_like +from .distributed.tensor_parallel import ( + _get_parameter_tp_plan, + verify_tp_plan, +) +from .distributed.utils import ( + _distributed_barrier, + distribute_model, + gather_full_state_dict, + init_device_mesh, + save_model_checkpoint_distributed, +) from .dynamic_module_utils import custom_object_save from .generation import CompileConfig, GenerationConfig -from .integrations import PeftAdapterMixin, deepspeed_config, hub_kernels, is_deepspeed_zero3_enabled, is_fsdp_enabled +from .integrations import PeftAdapterMixin, deepspeed_config, hub_kernels, is_deepspeed_zero3_enabled from .integrations.accelerate import ( _get_device_map, accelerate_disk_offload, @@ -75,15 +89,6 @@ from .integrations.peft import maybe_load_adapters from .integrations.sdpa_attention import sdpa_attention_forward from .integrations.sdpa_paged import sdpa_attention_paged_forward -from .integrations.tensor_parallel import ( - ALL_PARALLEL_STYLES, - _get_parameter_tp_plan, - distribute_model, - gather_state_dict_for_save, - initialize_tensor_parallelism, - shard_and_distribute_module, - verify_tp_plan, -) from .loss.loss_utils import LOSS_MAPPING from .modeling_flash_attention_utils import ( FLASH_ATTENTION_COMPATIBILITY_MATRIX, @@ -93,7 +98,7 @@ ) from .modeling_rope_utils import ROPE_INIT_FUNCTIONS from .monkey_patching import apply_patches, patch_output_recorders -from .pytorch_utils import id_tensor_storage +from .pytorch_utils import _torch_distributed_available, id_tensor_storage from .quantizers import HfQuantizer from .quantizers.auto import get_hf_quantizer from .quantizers.quantizers_utils import get_module_from_name @@ -143,8 +148,6 @@ from ._typing import DeviceMeshLike -_torch_distributed_available = torch.distributed.is_available() - if is_sagemaker_mp_enabled(): import smdistributed.modelparallel.torch as smp from smdistributed.modelparallel import __version__ as SMP_VERSION @@ -182,6 +185,7 @@ class LoadStateDictConfig: dtype_plan: dict = field(default_factory=dict) hf_quantizer: HfQuantizer | None = None device_mesh: "DeviceMeshLike | None" = None + tp_plan: dict[str, str] | None = None weights_only: bool = True weight_mapping: list[WeightConverter | WeightRenaming] | None = None disable_mmap: bool | None = None @@ -1244,6 +1248,10 @@ class PreTrainedModel(nn.Module, EmbeddingAccessMixin, ModuleUtilsMixin, PushToH # For top-level models, this attribute is currently defined in respective model code. For base models, this attribute comes # from `config.base_model_tp_plan` during `post_init`. _tp_plan: dict[str, str] = None + # Sequence-parallel plan used when `distributed_config.enable_sequence_parallel` is set. For top-level models, this + # attribute is defined on the head class (e.g. `*ForCausalLM`). For base models, it comes from + # `config.base_model_sp_plan` during `post_init`. + _sp_plan: dict[str, str] = None # Tensor parallel degree to which model is sharded to _tp_size = None # A pipeline parallel plan specifying the layers which may not be present on all ranks when PP is enabled. For top-level @@ -1383,13 +1391,26 @@ def post_init(self): correctly in the case of composite models (that is, the top level model should know about those properties from its children). """ # Attach the different parallel plans and tied weight keys to the top-most model, so that everything is - # easily available - self._tp_plan, self._ep_plan, self._pp_plan = {}, {}, {} - # If current model is a base model, attach `base_model_tp_plan` and `base_model_pp_plan` from config + # easily available. Seed with the class-level plans (e.g. `*ForCausalLM._tp_plan = {"lm_head": ...}`) + # before the instance attribute shadows them — task-head plans live on the head class. + cls_tp_plan = getattr(self, "_tp_plan", None) or {} + cls_sp_plan = getattr(self, "_sp_plan", None) or {} + cls_fsdp_plan = getattr(self, "_fsdp_plan", None) or {} + self._tp_plan = dict(cls_tp_plan) + self._sp_plan = dict(cls_sp_plan) + self._ep_plan = {} + self._pp_plan = {} + self._fsdp_plan = dict(cls_fsdp_plan) + # If current model is a base model, attach `base_model_*_plan` from config if self.base_model is self: self._pp_plan = self.config.base_model_pp_plan.copy() if self.config.base_model_pp_plan is not None else {} self._tp_plan = self.config.base_model_tp_plan.copy() if self.config.base_model_tp_plan is not None else {} + self._sp_plan = self.config.base_model_sp_plan.copy() if self.config.base_model_sp_plan is not None else {} self._ep_plan = self.config.base_model_ep_plan.copy() if self.config.base_model_ep_plan is not None else {} + self._fsdp_plan = ( + self.config.base_model_fsdp_plan.copy() if self.config.base_model_fsdp_plan is not None else {} + ) + # Current submodel should register its tied weights self.all_tied_weights_keys = self.get_expanded_tied_weights_keys(all_submodels=False) # Current submodel should register its `_keep_in_fp32_modules` @@ -1411,8 +1432,12 @@ def post_init(self): self._ep_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()}) if plan := getattr(module, "_tp_plan", None): self._tp_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()}) + if plan := getattr(module, "_sp_plan", None): + self._sp_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()}) if plan := getattr(module, "_pp_plan", None): self._pp_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()}) + if plan := getattr(module, "_fsdp_plan", None): + self._fsdp_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()}) # Always attach the keys of the children (if the children's config says to NOT tie, then it's empty) if tied_keys := getattr(module, "all_tied_weights_keys", None): self.all_tied_weights_keys.update({f"{name}.{k}": f"{name}.{v}" for k, v in tied_keys.copy().items()}) @@ -1461,14 +1486,6 @@ def tp_plan(self, plan: dict[str, str] | None): if not isinstance(plan, dict): raise ValueError("Can only set a dictionary as `tp_plan`") - # Ensure the styles are all valid - for layer_pattern, parallel_style in plan.items(): - if parallel_style not in ALL_PARALLEL_STYLES: - raise ValueError( - f"Unsupported tensor parallel style '{parallel_style}' for layer '{layer_pattern}'. " - f"Supported styles are {list(ALL_PARALLEL_STYLES.keys())}" - ) - # Validate that the layer patterns match existing model structure. We check this by getting all parameter # names and seeing if any match the patterns model_param_names = [name for name, _ in self.named_parameters()] @@ -2065,6 +2082,8 @@ def _can_set_attn_implementation(cls) -> bool: """Detect whether the class supports setting its attention implementation dynamically. It is an ugly check based on opening the file, but avoids maintaining yet another property flag. """ + # Skip dynamic wrappers like FSDP2's FSDP, whose __module__ is inside torch.* + cls = next((k for k in cls.__mro__ if not k.__module__.startswith("torch.")), cls) class_module = sys.modules.get(cls.__module__) # Missing module entry (e.g. cleared by a test) or custom model in a jupyter notebook / repl -> do not allow to set it if class_module is None or not hasattr(class_module, "__file__"): @@ -2084,6 +2103,8 @@ def _can_set_experts_implementation(cls) -> bool: """Detect whether the class supports setting its experts implementation dynamically. It is an ugly check based on opening the file, but avoids maintaining yet another property flag. """ + # Skip dynamic wrappers like FSDP2's FSDP, whose __module__ is inside torch.* + cls = next((k for k in cls.__mro__ if not k.__module__.startswith("torch.")), cls) class_module = sys.modules.get(cls.__module__) # Missing module entry (e.g. cleared by a test) or custom model in a jupyter notebook / repl -> do not allow to set it if class_module is None or not hasattr(class_module, "__file__"): @@ -3304,6 +3325,7 @@ def save_pretrained( token: str | bool | None = None, save_peft_format: bool = True, save_original_format: bool = True, + distributed_checkpoint: bool = False, **kwargs, ): """ @@ -3369,7 +3391,7 @@ def save_pretrained( ) # we need to check against tp_size, not tp_plan, as tp_plan is substituted to the class one - if self._tp_size is not None and not is_huggingface_hub_greater_or_equal("0.31.4"): + if self.tp_size is not None and not is_huggingface_hub_greater_or_equal("0.31.4"): raise ImportError( "Saving a model with tensor parallelism requires `huggingface_hub` version 0.31.4 or higher." ) @@ -3379,9 +3401,13 @@ def save_pretrained( return os.makedirs(save_directory, exist_ok=True) + save_on_this_rank = is_main_process + if torch.distributed.is_initialized() and getattr(self, "device_mesh", None) is not None: + # DTensor gather materializes a full plain state dict only on global rank 0. + save_on_this_rank = save_on_this_rank and torch.distributed.get_rank() == 0 save_directory_path = os.fspath(save_directory) - if push_to_hub: + if push_to_hub and save_on_this_rank: commit_message = kwargs.pop("commit_message", None) repo_id = kwargs.pop("repo_id", save_directory_path.split(os.path.sep)[-1]) create_pr = kwargs.pop("create_pr", False) @@ -3406,46 +3432,72 @@ def save_pretrained( # If we have a custom model, we copy the file defining it in the folder and set the attributes so it can be # loaded from the Hub. - if self._auto_class is not None: + if self._auto_class is not None and save_on_this_rank: custom_object_save(self, save_directory, config=self.config) - # Save the config - if is_main_process: - if not _hf_peft_config_loaded: - model_to_save.config.save_pretrained(save_directory) - if self.can_generate(): - model_to_save.generation_config.save_pretrained(save_directory) + # Don't persist distributed_config in saved config — it's runtime-only + # (otherwise AutoConfig absorbs it on reload, preventing from_pretrained from seeing it as a kwarg). + # Keep a runtime copy around because TP/FSDP save helpers still rely on it after config serialization. + distributed_config = getattr(model_to_save.config, "distributed_config", None) + if distributed_config is not None: + del model_to_save.config.distributed_config - if _hf_peft_config_loaded: - logger.info( - "Detected adapters on the model, saving the model in the PEFT format, only adapter weights will be saved." - ) - state_dict = model_to_save.get_adapter_state_dict(state_dict=state_dict) + # Save the config + try: + if save_on_this_rank: + if not _hf_peft_config_loaded: + model_to_save.config.save_pretrained(save_directory) + if self.can_generate(): + model_to_save.generation_config.save_pretrained(save_directory) - if save_peft_format: + if _hf_peft_config_loaded: logger.info( - "To match the expected format of the PEFT library, all keys of the state dict of adapters will be prepended with `base_model.model`." + "Detected adapters on the model, saving the model in the PEFT format, only adapter weights will be saved." ) - peft_state_dict = {} - for key, value in state_dict.items(): - peft_state_dict[f"base_model.model.{key}"] = value - state_dict = peft_state_dict + state_dict = model_to_save.get_adapter_state_dict(state_dict=state_dict) + + if save_peft_format: + logger.info( + "To match the expected format of the PEFT library, all keys of the state dict of adapters will be prepended with `base_model.model`." + ) + peft_state_dict = {} + for key, value in state_dict.items(): + peft_state_dict[f"base_model.model.{key}"] = value + state_dict = peft_state_dict - active_adapter = self.active_adapters() + active_adapter = self.active_adapters() - if len(active_adapter) > 1: - raise ValueError( - "Multiple active adapters detected, saving multiple active adapters is not supported yet. You can save adapters separately one by one " - "by iteratively calling `model.set_adapter(adapter_name)` then `model.save_pretrained(...)`" - ) - active_adapter = active_adapter[0] + if len(active_adapter) > 1: + raise ValueError( + "Multiple active adapters detected, saving multiple active adapters is not supported yet. You can save adapters separately one by one " + "by iteratively calling `model.set_adapter(adapter_name)` then `model.save_pretrained(...)`" + ) + active_adapter = active_adapter[0] - current_peft_config = self.peft_config[active_adapter] - current_peft_config.save_pretrained(save_directory) + current_peft_config = self.peft_config[active_adapter] + current_peft_config.save_pretrained(save_directory) + finally: + if distributed_config is not None: + model_to_save.config.distributed_config = distributed_config - # Get the model state_dict + if distributed_checkpoint: + if torch.distributed.is_initialized() and getattr(self, "device_mesh", None) is None: + raise ValueError( + "save_pretrained(distributed_checkpoint=True) requires the model to have been " + "initialized with a distributed_config (device_mesh is None)." + ) + save_model_checkpoint_distributed(self, save_directory) + return + + # Get the model state_dict (handles FSDP unshard + TP gather in one call) + used_distributed_gather = False if state_dict is None: - state_dict = model_to_save.state_dict() + if getattr(self, "device_mesh", None) is not None: + # Pass self (not model_to_save) so device_mesh/tp_size/tp_plan are available + state_dict = gather_full_state_dict(self) + used_distributed_gather = True + else: + state_dict = model_to_save.state_dict() # if any model parameters are offloaded, we need to know it for later is_offloaded = False @@ -3471,10 +3523,6 @@ def save_pretrained( if ignore_key in state_dict: del state_dict[ignore_key] - # If model was sharded with TP, gather full tensors for saving - if self._tp_size is not None: - state_dict = gather_state_dict_for_save(state_dict, self._tp_plan, self._device_mesh, self._tp_size) - # Remove tied weights as safetensors do not handle them state_dict = remove_tied_weights_from_state_dict(state_dict, model_to_save) @@ -3516,54 +3564,55 @@ def save_pretrained( filename.startswith(weights_no_suffix) and os.path.isfile(full_filename) and filename not in state_dict_split.filename_to_tensors - and is_main_process + and save_on_this_rank and reg.fullmatch(filename_no_suffix) is not None ): os.remove(full_filename) # Save the model - for shard_file, tensor_names in logging.tqdm( - state_dict_split.filename_to_tensors.items(), desc="Writing model shards" - ): - filename = os.path.join(save_directory, shard_file) - shard_state_dict = {} - for tensor_name in tensor_names: - # Get the tensor, and remove it from state_dict to avoid keeping the ref - tensor = state_dict.pop(tensor_name) - - # If the param was offloaded, we need to load it back from disk to resave it. It's a strange pattern, - # but it would otherwise not be contained in the saved shard if we were to simply move the file - # or something - if is_offloaded and tensor.device.type == "meta": - tensor = load_offloaded_parameter(model_to_save, tensor_name) - - # only do contiguous after it's permuted correctly in case of TP - shard_state_dict[tensor_name] = tensor.contiguous() - - # TODO: it would be very nice to do the writing concurrently, but safetensors never releases the GIL, - # so it's not possible for now.... - # Write the shard to disk - safe_save_file(shard_state_dict, filename, metadata=metadata) - # Cleanup the data before next loop (important with offloading, so we don't blowup cpu RAM) - del shard_state_dict - - if index is None: - path_to_weights = os.path.join(save_directory, weights_name) - logger.info(f"Model weights saved in {path_to_weights}") - else: - save_index_file = SAFE_WEIGHTS_INDEX_NAME - save_index_file = os.path.join(save_directory, _add_variant(save_index_file, variant)) - # Save the index as well - with open(save_index_file, "w", encoding="utf-8") as f: - content = json.dumps(index, indent=2, sort_keys=True) + "\n" - f.write(content) - logger.info( - f"The model is bigger than the maximum size per checkpoint ({max_shard_size}) and is going to be " - f"split in {len(state_dict_split.filename_to_tensors)} checkpoint shards. You can find where each parameters has been saved in the " - f"index located at {save_index_file}." - ) + if save_on_this_rank: + for shard_file, tensor_names in logging.tqdm( + state_dict_split.filename_to_tensors.items(), desc="Writing model shards" + ): + filename = os.path.join(save_directory, shard_file) + shard_state_dict = {} + for tensor_name in tensor_names: + # Get the tensor, and remove it from state_dict to avoid keeping the ref + tensor = state_dict.pop(tensor_name) + + # If the param was offloaded, we need to load it back from disk to resave it. It's a strange pattern, + # but it would otherwise not be contained in the saved shard if we were to simply move the file + # or something + if is_offloaded and tensor.device.type == "meta": + tensor = load_offloaded_parameter(model_to_save, tensor_name) + + # only do contiguous after it's permuted correctly in case of TP + shard_state_dict[tensor_name] = tensor.contiguous() + + # TODO: it would be very nice to do the writing concurrently, but safetensors never releases the GIL, + # so it's not possible for now.... + # Write the shard to disk + safe_save_file(shard_state_dict, filename, metadata=metadata) + # Cleanup the data before next loop (important with offloading, so we don't blowup cpu RAM) + del shard_state_dict + + if index is None: + path_to_weights = os.path.join(save_directory, weights_name) + logger.info(f"Model weights saved in {path_to_weights}") + else: + save_index_file = SAFE_WEIGHTS_INDEX_NAME + save_index_file = os.path.join(save_directory, _add_variant(save_index_file, variant)) + # Save the index as well + with open(save_index_file, "w", encoding="utf-8") as f: + content = json.dumps(index, indent=2, sort_keys=True) + "\n" + f.write(content) + logger.info( + f"The model is bigger than the maximum size per checkpoint ({max_shard_size}) and is going to be " + f"split in {len(state_dict_split.filename_to_tensors)} checkpoint shards. You can find where each parameters has been saved in the " + f"index located at {save_index_file}." + ) - if push_to_hub: + if push_to_hub and save_on_this_rank: # Eventually create an empty model card model_card = create_and_tag_model_card(repo_id, self.model_tags, token=token) @@ -3579,6 +3628,13 @@ def save_pretrained( create_pr=create_pr, ) + # `gather_full_state_dict` concentrates the full state on rank 0 only; + # other ranks then loop over an empty shard list and would race ahead + # of rank 0's safetensors writes. Barrier so any subsequent + # `from_pretrained` on this path sees the consolidated files. + if used_distributed_gather: + _distributed_barrier() + @wraps(PushToHubMixin.push_to_hub) def push_to_hub(self, *args, **kwargs): tags = self.model_tags if self.model_tags is not None else [] @@ -3972,13 +4028,22 @@ def from_pretrained( max_memory (`Dict`, *optional*): A dictionary device identifier to maximum memory if using `device_map`. Will default to the maximum memory available for each GPU and the available CPU RAM if unset. - tp_plan (`Optional[Union[dict, str]]`, *optional*): - A torch tensor parallel plan, see [here](https://pytorch.org/tutorials/intermediate/TP_tutorial.html). Use `tp_plan="auto"` to - use the predefined plan based on the model. If it's a dict, then it should match between module names and desired layout. - Note that if you use it, you should launch your script accordingly with `torchrun [args] script.py`. This will be much - faster than using a `device_map`, but has limitations. - tp_size (`str`, *optional*): - A torch tensor parallel degree. If not provided would default to world size. + distributed_config ([`DistributedConfig`], *optional*): + Configuration for native distributed training (FSDP2 + TP) via `torch.distributed`. Mutually + exclusive with `quantization_config` (for now) and `device_map`. When set, accelerate is not used for + device placement or dispatch. Launch with `torchrun --nproc_per_node=N script.py`. + + Accepts `tp_size`, `tp_plan`, `fsdp_size`, `fsdp_plan`. When a size is specified without a + plan, the plan defaults to the model's predefined behavior: TP uses the model's predefined + tensor parallel sharding plan, FSDP wraps each transformer layer individually with FSDP2 + (`fully_shard`). Both plans also accept a `dict` for manual control: `tp_plan` maps + parameter names to parallel styles (e.g. `{"model.layers.*.self_attn.q_proj": "colwise"}`), + `fsdp_plan` maps module names to wrap (e.g. `{"model.layers.0": {}, "model.layers.1": {}}`). + + Examples: + - TP-only: `DistributedConfig(tp_size=4)` + - FSDP-only: `DistributedConfig(fsdp_size=4)` + - 2D parallel: `DistributedConfig(tp_size=2, fsdp_size=2)` on 4 GPUs device_mesh (`torch.distributed.DeviceMesh`, *optional*): A torch device mesh. If not provided would default to world size. Used only for tensor parallel for now. If provided, it has to contain dimension named `"tp"` in case it's > 1 dimensional, this dimension will be used for tensor parallelism @@ -4072,8 +4137,6 @@ def from_pretrained( adapter_name = kwargs.pop("adapter_name", "default") generation_config = kwargs.pop("generation_config", None) gguf_file = kwargs.pop("gguf_file", None) - tp_plan = kwargs.pop("tp_plan", None) - tp_size = kwargs.pop("tp_size", None) distributed_config: DistributedConfig = kwargs.pop("distributed_config", None) device_mesh = kwargs.pop("device_mesh", None) trust_remote_code = kwargs.pop("trust_remote_code", None) @@ -4082,8 +4145,17 @@ def from_pretrained( kernel_config = kwargs.pop("kernel_config", None) key_mapping = kwargs.pop("key_mapping", None) - if distributed_config is not None and tp_plan is None: - tp_plan = "auto" + if distributed_config is not None: + if device_map is not None: + raise ValueError( + "`distributed_config` and `device_map` are mutually exclusive. " + "`distributed_config` handles device placement natively via torch.distributed." + ) + # NOTE(3outeille): support quantization (fp4/fp8) with distributed training later + if quantization_config is not None: + raise ValueError( + "Quantization is not currently supported with distributed training. Please disable quantization or distributed_config." + ) # Not used anymore -- remove them from the kwargs for name in ["mirror", "_fast_init", "low_cpu_mem_usage", "from_tf", "from_flax", "offload_state_dict"]: @@ -4114,17 +4186,18 @@ def from_pretrained( "`state_dict` cannot be passed together with a model name or a `gguf_file`. Use one of the two loading strategies." ) - if device_map == "auto" and int(os.environ.get("WORLD_SIZE", "0")): - logger.info( - "You've set device_map=`auto` while triggering a distributed run with torchrun. This might lead to unexpected behavior. " - "If your plan is to load the model on each device, you should set device_map={" - ": PartialState().process_index} where PartialState comes from accelerate library" - ) + if distributed_config is not None: + device_mesh = init_device_mesh(distributed_config) + else: + # Accelerate path + if device_map == "auto" and int(os.environ.get("WORLD_SIZE", "0")): + logger.info( + "You've set device_map=`auto` while triggering a distributed run with torchrun. This might lead to unexpected behavior. " + "If your plan is to load the model on each device, you should set device_map={" + ": PartialState().process_index} where PartialState comes from accelerate library" + ) - if tp_plan is not None or tp_size is not None: # TP warnings, and setup - device_map, device_mesh, tp_size = initialize_tensor_parallelism( - tp_plan, tp_size=tp_size, device_mesh=device_mesh, device_map=device_map - ) + device_map = check_and_set_device_map(device_map) # validate & normalize (requires accelerate) if gguf_file is not None and not is_accelerate_available(): raise ValueError("accelerate is required when loading a GGUF file `pip install accelerate`.") @@ -4137,7 +4210,6 @@ def from_pretrained( download_kwargs_with_commit, **adapter_kwargs, ) - device_map = check_and_set_device_map(device_map) # warn, error and fix the device map user_agent = {"file_type": "model", "framework": "pytorch", "from_auto_class": from_auto_class} if from_pipeline is not None: @@ -4265,17 +4337,20 @@ def from_pretrained( # instantiated model, as the flags can be modified by instances sometimes) dtype_plan = model._get_dtype_plan(dtype) - # Obtain the weight conversion mapping for this model if any are registered and apply to all submodels recursively + # Obtain the weight conversion mapping for this model if any are registered weight_conversions = get_model_conversion_mapping(model, key_mapping, hf_quantizer) - if _torch_distributed_available and device_mesh is not None: # add hooks to nn.Modules: no weights - model = distribute_model(model, tp_plan, distributed_config, device_mesh, tp_size) - - # Prepare the full device map - if device_map is not None: - device_map = _get_device_map(model, device_map, max_memory, hf_quantizer) + if distributed_config is not None: + model = distribute_model(model, distributed_config, device_mesh) + else: + # Accelerate path: auto device mapping + if device_map is not None: + device_map = _get_device_map(model, device_map, max_memory, hf_quantizer) # Finalize model weight initialization + active_tp_plan = ( + getattr(model, "_tp_plan", None) if getattr(distributed_config, "tp_size", None) is not None else None + ) load_config = LoadStateDictConfig( pretrained_model_name_or_path=pretrained_model_name_or_path, ignore_mismatched_sizes=ignore_mismatched_sizes, @@ -4287,6 +4362,7 @@ def from_pretrained( dtype_plan=dtype_plan, hf_quantizer=hf_quantizer, device_mesh=device_mesh, + tp_plan=active_tp_plan, weights_only=weights_only, weight_mapping=weight_conversions, use_safetensors=use_safetensors, @@ -4355,7 +4431,7 @@ def _load_pretrained_model( expected_keys = list(model.state_dict().keys()) if expected_keys is None else expected_keys if logger.level >= logging.WARNING: - verify_tp_plan(expected_keys, getattr(model, "_tp_plan", None)) + verify_tp_plan(expected_keys, load_config.tp_plan) # This offload index if for params explicitly on the "disk" in the device_map disk_offload_index = None @@ -4427,7 +4503,7 @@ def _load_pretrained_model( model=model, state_dict=merged_state_dict, load_config=load_config, - tp_plan=model.tp_plan, + tp_plan=load_config.tp_plan, disk_offload_index=disk_offload_index, ) @@ -4579,8 +4655,8 @@ def tp_size(self): """ Returns the model's tensor parallelism degree. """ - # if None, the model didn't undergo tensor parallel sharding - return self._tp_size + dc = getattr(self.config, "distributed_config", None) + return dc.tp_size if dc is not None else None @property def supports_pp_plan(self): @@ -4735,15 +4811,20 @@ def _move_missing_keys_from_meta_to_device( # will be re-initialized for nothing (which can be quite long) for key in missing_keys - self.all_tied_weights_keys.keys(): param = self.get_parameter_or_buffer(key) - param_device = get_device(device_map, key, valid_torch_device=True) - value = torch.empty_like(param, device=param_device) - # For TP, we may need to shard the param - if device_mesh is not None: - shard_and_distribute_module( - self, value, param, key, None, False, device_mesh.get_local_rank(), device_mesh + if isinstance(param, DTensor): + # DTensor from parallelize_module on meta — materialize on actual device + local_value = torch.empty( + param._local_tensor.shape, + dtype=param.dtype, + device=torch.device(param.device_mesh.device_type, torch.cuda.current_device()), ) - # Otherwise, just move it to device + new_dtensor = _dtensor_from_local_like(local_value, param) + with torch.no_grad(): + new_param = torch.nn.Parameter(new_dtensor, requires_grad=param.requires_grad) + torch.utils.swap_tensors(param, new_param) else: + param_device = get_device(device_map, key, valid_torch_device=True) + value = torch.empty_like(param, device=param_device) _load_parameter_into_model(self, key, value) # We need to move back non-persistent buffers as well, as they are not part of loaded weights anyway for key, buffer in self.named_non_persistent_buffers(): @@ -4826,7 +4907,7 @@ def mark_tied_weights_as_initialized(self, loading_info): later as they will be tied (overwritten) anyway. This is very important as most embeddings are tied, and they are huge params (vocabularies are often 256k), so running inits on them is very costly.""" - for tied_param in getattr(self, "all_tied_weights_keys", {}).keys(): + for tied_param in self.all_tied_weights_keys.keys(): param = self.get_parameter(tied_param) setattr(param, "_is_hf_initialized", True) diff --git a/src/transformers/models/afmoe/modeling_afmoe.py b/src/transformers/models/afmoe/modeling_afmoe.py index 421119b33deb..1fb30e564e1c 100644 --- a/src/transformers/models/afmoe/modeling_afmoe.py +++ b/src/transformers/models/afmoe/modeling_afmoe.py @@ -615,8 +615,10 @@ def forward( @auto_docstring class AfmoeForCausalLM(AfmoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/afmoe/modular_afmoe.py b/src/transformers/models/afmoe/modular_afmoe.py index f3ff9f15b103..0875127a3215 100644 --- a/src/transformers/models/afmoe/modular_afmoe.py +++ b/src/transformers/models/afmoe/modular_afmoe.py @@ -395,7 +395,9 @@ def forward( class AfmoeForCausalLM(LlamaForCausalLM, AfmoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/apertus/configuration_apertus.py b/src/transformers/models/apertus/configuration_apertus.py index 1e0122160b30..c7fcc4159c06 100644 --- a/src/transformers/models/apertus/configuration_apertus.py +++ b/src/transformers/models/apertus/configuration_apertus.py @@ -49,11 +49,26 @@ class ApertusConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -61,6 +76,12 @@ class ApertusConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 131072 hidden_size: int = 4096 intermediate_size: int = 14336 diff --git a/src/transformers/models/apertus/modeling_apertus.py b/src/transformers/models/apertus/modeling_apertus.py index 7d14dd3d14c8..8c5757377b03 100644 --- a/src/transformers/models/apertus/modeling_apertus.py +++ b/src/transformers/models/apertus/modeling_apertus.py @@ -422,8 +422,10 @@ def forward( @auto_docstring class ApertusForCausalLM(ApertusPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/apertus/modular_apertus.py b/src/transformers/models/apertus/modular_apertus.py index 3c9eb6d8b6ea..075a6de4776b 100644 --- a/src/transformers/models/apertus/modular_apertus.py +++ b/src/transformers/models/apertus/modular_apertus.py @@ -67,11 +67,26 @@ class ApertusConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -79,6 +94,12 @@ class ApertusConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 131072 hidden_size: int = 4096 intermediate_size: int = 14336 diff --git a/src/transformers/models/arcee/configuration_arcee.py b/src/transformers/models/arcee/configuration_arcee.py index ee711d608204..e1d957bcac43 100644 --- a/src/transformers/models/arcee/configuration_arcee.py +++ b/src/transformers/models/arcee/configuration_arcee.py @@ -49,9 +49,23 @@ class ArceeConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -59,6 +73,12 @@ class ArceeConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 32000 hidden_size: int = 2560 intermediate_size: int = 18432 diff --git a/src/transformers/models/arcee/modeling_arcee.py b/src/transformers/models/arcee/modeling_arcee.py index 8d2d05bf2952..05ebab40a879 100644 --- a/src/transformers/models/arcee/modeling_arcee.py +++ b/src/transformers/models/arcee/modeling_arcee.py @@ -424,8 +424,10 @@ def forward( @auto_docstring(checkpoint="arcee-ai/AFM-4.5B") class ArceeForCausalLM(ArceePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/arcee/modular_arcee.py b/src/transformers/models/arcee/modular_arcee.py index 91cd8e13f1ed..a382b2e7191b 100644 --- a/src/transformers/models/arcee/modular_arcee.py +++ b/src/transformers/models/arcee/modular_arcee.py @@ -53,9 +53,23 @@ class ArceeConfig(LlamaConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } vocab_size: int = 32000 diff --git a/src/transformers/models/aria/configuration_aria.py b/src/transformers/models/aria/configuration_aria.py index 7694656905dd..dfc1b8e74258 100644 --- a/src/transformers/models/aria/configuration_aria.py +++ b/src/transformers/models/aria/configuration_aria.py @@ -44,10 +44,25 @@ class AriaTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -55,6 +70,12 @@ class AriaTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 32000 hidden_size: int = 4096 diff --git a/src/transformers/models/aria/modeling_aria.py b/src/transformers/models/aria/modeling_aria.py index e4fab906e561..91fcf5d106da 100644 --- a/src/transformers/models/aria/modeling_aria.py +++ b/src/transformers/models/aria/modeling_aria.py @@ -758,8 +758,10 @@ def forward( @auto_docstring class AriaTextForCausalLM(AriaTextPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config: AriaTextConfig): super().__init__(config) diff --git a/src/transformers/models/aria/modular_aria.py b/src/transformers/models/aria/modular_aria.py index 7e50fe3812fa..121e14898537 100644 --- a/src/transformers/models/aria/modular_aria.py +++ b/src/transformers/models/aria/modular_aria.py @@ -113,10 +113,10 @@ class AriaTextConfig(LlamaConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", } intermediate_size: int = 4096 diff --git a/src/transformers/models/audioflamingo3/modeling_audioflamingo3.py b/src/transformers/models/audioflamingo3/modeling_audioflamingo3.py index 0abf85865db8..ca77ef4a6cb1 100644 --- a/src/transformers/models/audioflamingo3/modeling_audioflamingo3.py +++ b/src/transformers/models/audioflamingo3/modeling_audioflamingo3.py @@ -410,6 +410,7 @@ class AudioFlamingo3ForConditionalGeneration(AudioFlamingo3PreTrainedModel, Gene _keep_in_fp32_modules_strict = None _supports_attention_backend = True _tp_plan = None + _sp_plan = None _pp_plan = None def __init__(self, config): diff --git a/src/transformers/models/audioflamingo3/modular_audioflamingo3.py b/src/transformers/models/audioflamingo3/modular_audioflamingo3.py index dfb2c1f54d35..a170d8ee4124 100644 --- a/src/transformers/models/audioflamingo3/modular_audioflamingo3.py +++ b/src/transformers/models/audioflamingo3/modular_audioflamingo3.py @@ -144,6 +144,7 @@ def __init__(self, config: AudioFlamingo3Config): class AudioFlamingo3ForConditionalGeneration(VoxtralForConditionalGeneration): _supports_attention_backend = True _tp_plan = None + _sp_plan = None _pp_plan = None _keep_in_fp32_modules_strict = None diff --git a/src/transformers/models/bamba/modeling_bamba.py b/src/transformers/models/bamba/modeling_bamba.py index 2b42a603d841..c3b3a8a24bf1 100644 --- a/src/transformers/models/bamba/modeling_bamba.py +++ b/src/transformers/models/bamba/modeling_bamba.py @@ -1071,8 +1071,10 @@ def _update_mamba_mask(self, attention_mask, past_key_values): @auto_docstring class BambaForCausalLM(BambaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/bitnet/modeling_bitnet.py b/src/transformers/models/bitnet/modeling_bitnet.py index 14c1581b250f..4994407d74c1 100644 --- a/src/transformers/models/bitnet/modeling_bitnet.py +++ b/src/transformers/models/bitnet/modeling_bitnet.py @@ -423,7 +423,9 @@ def forward( class BitNetForCausalLM(BitNetPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = None + _sp_plan = None _pp_plan = None + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/bitnet/modular_bitnet.py b/src/transformers/models/bitnet/modular_bitnet.py index 7f3a248766b2..c8376f131859 100644 --- a/src/transformers/models/bitnet/modular_bitnet.py +++ b/src/transformers/models/bitnet/modular_bitnet.py @@ -110,6 +110,7 @@ class BitNetModel(LlamaModel): class BitNetForCausalLM(LlamaForCausalLM): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = None + _sp_plan = None _pp_plan = None def forward( diff --git a/src/transformers/models/cohere/configuration_cohere.py b/src/transformers/models/cohere/configuration_cohere.py index d52a3e008427..4c05dce75e47 100644 --- a/src/transformers/models/cohere/configuration_cohere.py +++ b/src/transformers/models/cohere/configuration_cohere.py @@ -53,16 +53,39 @@ class CohereConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } vocab_size: int = 256000 hidden_size: int = 8192 intermediate_size: int = 22528 diff --git a/src/transformers/models/cohere/modeling_cohere.py b/src/transformers/models/cohere/modeling_cohere.py index b8bf50af9bf4..309ddc142dc3 100644 --- a/src/transformers/models/cohere/modeling_cohere.py +++ b/src/transformers/models/cohere/modeling_cohere.py @@ -454,8 +454,10 @@ def forward( @auto_docstring class CohereForCausalLM(CoherePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/cohere2/configuration_cohere2.py b/src/transformers/models/cohere2/configuration_cohere2.py index 48c2df360354..19197e26204d 100644 --- a/src/transformers/models/cohere2/configuration_cohere2.py +++ b/src/transformers/models/cohere2/configuration_cohere2.py @@ -52,10 +52,25 @@ class Cohere2Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -63,6 +78,12 @@ class Cohere2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 256000 hidden_size: int = 8192 intermediate_size: int = 22528 diff --git a/src/transformers/models/cohere2/modeling_cohere2.py b/src/transformers/models/cohere2/modeling_cohere2.py index f43b2a0ef412..6a584c687c79 100644 --- a/src/transformers/models/cohere2/modeling_cohere2.py +++ b/src/transformers/models/cohere2/modeling_cohere2.py @@ -433,8 +433,10 @@ def forward( @auto_docstring class Cohere2ForCausalLM(Cohere2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/cohere2/modular_cohere2.py b/src/transformers/models/cohere2/modular_cohere2.py index d19055a1b787..2720840a931e 100644 --- a/src/transformers/models/cohere2/modular_cohere2.py +++ b/src/transformers/models/cohere2/modular_cohere2.py @@ -73,10 +73,25 @@ class Cohere2Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -84,6 +99,12 @@ class Cohere2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 256000 hidden_size: int = 8192 intermediate_size: int = 22528 diff --git a/src/transformers/models/csm/modeling_csm.py b/src/transformers/models/csm/modeling_csm.py index dd0a1f0aa749..99f643e93b44 100644 --- a/src/transformers/models/csm/modeling_csm.py +++ b/src/transformers/models/csm/modeling_csm.py @@ -549,7 +549,9 @@ def forward(self, hidden_states, codebook_indices=None): class CsmDepthDecoderForCausalLM(CsmPreTrainedModel, GenerationMixin): _tied_weights_keys = None _tp_plan = None + _sp_plan = None _pp_plan = None + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/csm/modular_csm.py b/src/transformers/models/csm/modular_csm.py index c82636613ee1..f1e95384ab52 100644 --- a/src/transformers/models/csm/modular_csm.py +++ b/src/transformers/models/csm/modular_csm.py @@ -274,6 +274,7 @@ def forward(self, hidden_states, codebook_indices=None): class CsmDepthDecoderForCausalLM(LlamaForCausalLM, GenerationMixin): _tied_weights_keys = None _tp_plan = None + _sp_plan = None _pp_plan = None def __init__(self, config): diff --git a/src/transformers/models/cwm/configuration_cwm.py b/src/transformers/models/cwm/configuration_cwm.py index ecc3743da19d..6eb7175ad020 100644 --- a/src/transformers/models/cwm/configuration_cwm.py +++ b/src/transformers/models/cwm/configuration_cwm.py @@ -49,10 +49,25 @@ class CwmConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -60,6 +75,12 @@ class CwmConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 128256 hidden_size: int = 6144 intermediate_size: int = 21504 diff --git a/src/transformers/models/cwm/modeling_cwm.py b/src/transformers/models/cwm/modeling_cwm.py index 3e0eb0504be0..b20fd9b7436f 100644 --- a/src/transformers/models/cwm/modeling_cwm.py +++ b/src/transformers/models/cwm/modeling_cwm.py @@ -426,8 +426,10 @@ def forward( @auto_docstring class CwmForCausalLM(CwmPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/data2vec/modeling_data2vec_audio.py b/src/transformers/models/data2vec/modeling_data2vec_audio.py index bb3af4ecf25a..048747656b9e 100755 --- a/src/transformers/models/data2vec/modeling_data2vec_audio.py +++ b/src/transformers/models/data2vec/modeling_data2vec_audio.py @@ -29,8 +29,8 @@ from ... import initialization as init from ...activations import ACT2FN +from ...distributed.fsdp import is_fsdp_managed_module from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...masking_utils import create_bidirectional_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer diff --git a/src/transformers/models/dbrx/modeling_dbrx.py b/src/transformers/models/dbrx/modeling_dbrx.py index 58735fb55c0b..a984b4585a55 100644 --- a/src/transformers/models/dbrx/modeling_dbrx.py +++ b/src/transformers/models/dbrx/modeling_dbrx.py @@ -642,7 +642,9 @@ def load_balancing_loss_func( class DbrxForCausalLM(DbrxPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "transformer.wte.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: DbrxConfig): diff --git a/src/transformers/models/dbrx/modular_dbrx.py b/src/transformers/models/dbrx/modular_dbrx.py index d737d59e1a8b..548af86ef4f6 100644 --- a/src/transformers/models/dbrx/modular_dbrx.py +++ b/src/transformers/models/dbrx/modular_dbrx.py @@ -430,7 +430,9 @@ def forward( class DbrxForCausalLM(DbrxPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "transformer.wte.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: DbrxConfig): diff --git a/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py b/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py index 1b8005f8efef..49b3cb9c6387 100644 --- a/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py +++ b/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py @@ -54,18 +54,39 @@ class DeepseekV2Config(PreTrainedConfig): base_model_tp_plan = { "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.q_b_proj": "colwise", - "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj", "layers.*.self_attn.kv_b_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + # MLA: don't shard q_a_proj / kv_a_proj_with_mqa — their outputs feed + # layernorms whose weights are full-size. Only b-projections (after + # the norm) and o_proj are sharded. + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_b_proj": "colwise", + "layers.*.self_attn.kv_b_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts": "moe_experts_allreduce", + # Shared experts output must stay Replicate to match experts output + # (they're summed inside the MoE block, before the outer allgather_split + # handles the SP boundary). + "layers.*.mlp.shared_experts.gate_proj": "colwise", + "layers.*.mlp.shared_experts.up_proj": "colwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -73,6 +94,12 @@ class DeepseekV2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 32000 hidden_size: int = 4096 intermediate_size: int = 11008 diff --git a/src/transformers/models/deepseek_v2/modeling_deepseek_v2.py b/src/transformers/models/deepseek_v2/modeling_deepseek_v2.py index 3ef8266218f7..1f9ba33ec6a5 100644 --- a/src/transformers/models/deepseek_v2/modeling_deepseek_v2.py +++ b/src/transformers/models/deepseek_v2/modeling_deepseek_v2.py @@ -541,8 +541,10 @@ def forward( @auto_docstring class DeepseekV2ForCausalLM(DeepseekV2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/deepseek_v2/modular_deepseek_v2.py b/src/transformers/models/deepseek_v2/modular_deepseek_v2.py index 5644c7dc2990..a6c6494c0415 100644 --- a/src/transformers/models/deepseek_v2/modular_deepseek_v2.py +++ b/src/transformers/models/deepseek_v2/modular_deepseek_v2.py @@ -69,18 +69,39 @@ class DeepseekV2Config(LlamaConfig): base_model_tp_plan = { "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.q_b_proj": "colwise", - "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj", "layers.*.self_attn.kv_b_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + # MLA: don't shard q_a_proj / kv_a_proj_with_mqa — their outputs feed + # layernorms whose weights are full-size. Only b-projections (after + # the norm) and o_proj are sharded. + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_b_proj": "colwise", + "layers.*.self_attn.kv_b_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts": "moe_experts_allreduce", + # Shared experts output must stay Replicate to match experts output + # (they're summed inside the MoE block, before the outer allgather_split + # handles the SP boundary). + "layers.*.mlp.shared_experts.gate_proj": "colwise", + "layers.*.mlp.shared_experts.up_proj": "colwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } model_type = "deepseek_v2" diff --git a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py index 4178547a5ff2..32e67ab4d90a 100644 --- a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py @@ -49,21 +49,25 @@ class DeepseekV3Config(PreTrainedConfig): model_type = "deepseek_v3" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } attribute_map = { "num_local_experts": "n_routed_experts", } @@ -111,5 +115,20 @@ def __post_init__(self, **kwargs): self.head_dim = self.qk_rope_head_dim super().__post_init__(**kwargs) + def convert_rope_params_to_dict(self, **kwargs): + rope_scaling = kwargs.pop("rope_scaling", None) + self.rope_parameters = rope_scaling or self.rope_parameters + self.rope_parameters = self.rope_parameters if self.rope_parameters is not None else {} + + # Standardize and validate the correctness of rotary position embeddings parameters + self.rope_parameters.setdefault("rope_theta", kwargs.pop("rope_theta", self.default_theta)) + self.standardize_rope_params() + + # Convert to float because RoPE fn expect a float. Models on the hub were saved as int + for key in ["beta_fast", "beta_slow", "factor"]: + if key in self.rope_parameters: + self.rope_parameters[key] = float(self.rope_parameters[key]) + return kwargs + __all__ = ["DeepseekV3Config"] diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index fe3acd9aeddd..8d828320b19b 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -634,8 +634,10 @@ def forward( @auto_docstring class DeepseekV3ForCausalLM(DeepseekV3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/deepseek_v4/configuration_deepseek_v4.py b/src/transformers/models/deepseek_v4/configuration_deepseek_v4.py index 8f8d818d8ced..9475e8d6f1d0 100644 --- a/src/transformers/models/deepseek_v4/configuration_deepseek_v4.py +++ b/src/transformers/models/deepseek_v4/configuration_deepseek_v4.py @@ -125,7 +125,13 @@ class DeepseekV4Config(PreTrainedConfig): "layers.*.mlp.gate": "ep_router", "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", "layers.*.mlp.experts.down_proj": "grouped_gemm", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", } vocab_size: int = 129280 diff --git a/src/transformers/models/deepseek_v4/modeling_deepseek_v4.py b/src/transformers/models/deepseek_v4/modeling_deepseek_v4.py index 163a0cee77a3..664a9bcbc1cd 100644 --- a/src/transformers/models/deepseek_v4/modeling_deepseek_v4.py +++ b/src/transformers/models/deepseek_v4/modeling_deepseek_v4.py @@ -1395,8 +1395,10 @@ def load_balancing_loss_func( @auto_docstring class DeepseekV4ForCausalLM(DeepseekV4PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/dia/generation_dia.py b/src/transformers/models/dia/generation_dia.py index d22b2fff0d8d..76506e66f00c 100644 --- a/src/transformers/models/dia/generation_dia.py +++ b/src/transformers/models/dia/generation_dia.py @@ -18,6 +18,7 @@ import torch import torch.distributed as dist +from ...distributed.fsdp import is_fsdp_managed_module from ...generation.logits_process import ( DiaClassifierFreeGuidanceLogitsProcessor, DiaEOSChannelFilterLogitsProcessor, @@ -29,7 +30,6 @@ from ...generation.streamers import BaseStreamer from ...generation.utils import GenerateOutput, GenerationConfig, GenerationMixin, GenerationMode from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...modeling_utils import PreTrainedModel from ...utils import logging diff --git a/src/transformers/models/diffllama/modeling_diffllama.py b/src/transformers/models/diffllama/modeling_diffllama.py index d80ccd572dc3..b822da50e2af 100644 --- a/src/transformers/models/diffllama/modeling_diffllama.py +++ b/src/transformers/models/diffllama/modeling_diffllama.py @@ -660,8 +660,10 @@ def forward( @auto_docstring class DiffLlamaForCausalLM(DiffLlamaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/doge/configuration_doge.py b/src/transformers/models/doge/configuration_doge.py index 8518d9021458..feac227b5a97 100644 --- a/src/transformers/models/doge/configuration_doge.py +++ b/src/transformers/models/doge/configuration_doge.py @@ -55,14 +55,14 @@ class DogeConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.dt_proj": "rowwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.dt_proj": "rowwise_allreduce", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", - "layers.*.mlp.router_gate": "colwise_gather_output", - "layers.*.mlp.down_embed": "rowwise_split_input", - "layers.*.mlp.up_embed": "rowwise_split_input", + "layers.*.mlp.down_proj": "rowwise_allreduce", + "layers.*.mlp.router_gate": "colwise_allgather", + "layers.*.mlp.down_embed": "vocab_allreduce", + "layers.*.mlp.up_embed": "vocab_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -70,6 +70,12 @@ class DogeConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 32768 hidden_size: int = 1024 intermediate_size: int = 2048 diff --git a/src/transformers/models/doge/modeling_doge.py b/src/transformers/models/doge/modeling_doge.py index 4aad59b52a9a..1f112465369d 100644 --- a/src/transformers/models/doge/modeling_doge.py +++ b/src/transformers/models/doge/modeling_doge.py @@ -716,8 +716,10 @@ def load_balancing_loss_func( @auto_docstring class DogeForCausalLM(DogePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/doge/modular_doge.py b/src/transformers/models/doge/modular_doge.py index 8b78126c0a00..bc7e22089ea3 100644 --- a/src/transformers/models/doge/modular_doge.py +++ b/src/transformers/models/doge/modular_doge.py @@ -84,14 +84,14 @@ class DogeConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.dt_proj": "rowwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.dt_proj": "rowwise_allreduce", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", - "layers.*.mlp.router_gate": "colwise_gather_output", - "layers.*.mlp.down_embed": "rowwise_split_input", - "layers.*.mlp.up_embed": "rowwise_split_input", + "layers.*.mlp.down_proj": "rowwise_allreduce", + "layers.*.mlp.router_gate": "colwise_allgather", + "layers.*.mlp.down_embed": "vocab_allreduce", + "layers.*.mlp.up_embed": "vocab_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -99,6 +99,12 @@ class DogeConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 32768 hidden_size: int = 1024 intermediate_size: int = 2048 diff --git a/src/transformers/models/dots1/configuration_dots1.py b/src/transformers/models/dots1/configuration_dots1.py index 4d568bf4a565..9846f961250f 100644 --- a/src/transformers/models/dots1/configuration_dots1.py +++ b/src/transformers/models/dots1/configuration_dots1.py @@ -51,18 +51,14 @@ class Dots1Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { @@ -70,6 +66,13 @@ class Dots1Config(PreTrainedConfig): "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + attribute_map = { "num_local_experts": "n_routed_experts", } diff --git a/src/transformers/models/dots1/modeling_dots1.py b/src/transformers/models/dots1/modeling_dots1.py index 95b21258ffd5..05b1a6d20b4c 100644 --- a/src/transformers/models/dots1/modeling_dots1.py +++ b/src/transformers/models/dots1/modeling_dots1.py @@ -569,8 +569,10 @@ def forward( @auto_docstring class Dots1ForCausalLM(Dots1PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/dots1/modular_dots1.py b/src/transformers/models/dots1/modular_dots1.py index 06402d63e28c..94bd074871c0 100644 --- a/src/transformers/models/dots1/modular_dots1.py +++ b/src/transformers/models/dots1/modular_dots1.py @@ -65,18 +65,14 @@ class Dots1Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { @@ -84,6 +80,13 @@ class Dots1Config(PreTrainedConfig): "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + attribute_map = { "num_local_experts": "n_routed_experts", } diff --git a/src/transformers/models/emu3/modeling_emu3.py b/src/transformers/models/emu3/modeling_emu3.py index eab9ee97da93..bd3f3011a539 100644 --- a/src/transformers/models/emu3/modeling_emu3.py +++ b/src/transformers/models/emu3/modeling_emu3.py @@ -1275,8 +1275,10 @@ def forward( @auto_docstring class Emu3ForCausalLM(Emu3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} config: Emu3TextConfig def __init__(self, config): diff --git a/src/transformers/models/ernie4_5/configuration_ernie4_5.py b/src/transformers/models/ernie4_5/configuration_ernie4_5.py index 896fb0b99402..5d805e1114e6 100644 --- a/src/transformers/models/ernie4_5/configuration_ernie4_5.py +++ b/src/transformers/models/ernie4_5/configuration_ernie4_5.py @@ -50,10 +50,25 @@ class Ernie4_5Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -61,6 +76,12 @@ class Ernie4_5Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 103424 hidden_size: int = 1024 intermediate_size: int = 3072 diff --git a/src/transformers/models/ernie4_5/modeling_ernie4_5.py b/src/transformers/models/ernie4_5/modeling_ernie4_5.py index ae533c7c6ef8..73f15f713be4 100644 --- a/src/transformers/models/ernie4_5/modeling_ernie4_5.py +++ b/src/transformers/models/ernie4_5/modeling_ernie4_5.py @@ -422,8 +422,10 @@ def forward( @auto_docstring class Ernie4_5ForCausalLM(Ernie4_5PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/ernie4_5_moe/configuration_ernie4_5_moe.py b/src/transformers/models/ernie4_5_moe/configuration_ernie4_5_moe.py index 0c0c0edbb760..035daabd966d 100644 --- a/src/transformers/models/ernie4_5_moe/configuration_ernie4_5_moe.py +++ b/src/transformers/models/ernie4_5_moe/configuration_ernie4_5_moe.py @@ -66,16 +66,14 @@ class Ernie4_5_MoeConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -83,6 +81,12 @@ class Ernie4_5_MoeConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 103424 pad_token_id: int | None = 0 bos_token_id: int | None = 1 diff --git a/src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py b/src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py index 9c106a90010d..939201f4ff19 100644 --- a/src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py +++ b/src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py @@ -654,8 +654,10 @@ def load_balancing_loss_func( @auto_docstring class Ernie4_5_MoeForCausalLM(Ernie4_5_MoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/ernie4_5_vl_moe/configuration_ernie4_5_vl_moe.py b/src/transformers/models/ernie4_5_vl_moe/configuration_ernie4_5_vl_moe.py index e4eea836f107..34caaa02dbeb 100644 --- a/src/transformers/models/ernie4_5_vl_moe/configuration_ernie4_5_vl_moe.py +++ b/src/transformers/models/ernie4_5_vl_moe/configuration_ernie4_5_vl_moe.py @@ -50,9 +50,9 @@ class Ernie4_5_VLMoeVisionConfig(PreTrainedConfig): base_model_tp_plan = { "blocks.*.attn.qkv": "colwise", - "blocks.*.attn.proj": "rowwise", + "blocks.*.attn.proj": "rowwise_allreduce", "blocks.*.mlp.fc1": "colwise", - "blocks.*.mlp.fc2": "rowwise", + "blocks.*.mlp.fc2": "rowwise_allreduce", } intermediate_size: int = 4 * 1280 temporal_merge_size: int = 2 @@ -86,13 +86,13 @@ class Ernie4_5_VLMoeTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -100,6 +100,12 @@ class Ernie4_5_VLMoeTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 103424 pad_token_id: int | None = None bos_token_id: int | None = None diff --git a/src/transformers/models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py b/src/transformers/models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py index 3736acd51103..0ec36d704726 100644 --- a/src/transformers/models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py +++ b/src/transformers/models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py @@ -94,9 +94,9 @@ class Ernie4_5_VLMoeVisionConfig(Qwen2VLVisionConfig): base_model_tp_plan = { "blocks.*.attn.qkv": "colwise", - "blocks.*.attn.proj": "rowwise", + "blocks.*.attn.proj": "rowwise_allreduce", "blocks.*.mlp.fc1": "colwise", - "blocks.*.mlp.fc2": "rowwise", + "blocks.*.mlp.fc2": "rowwise_allreduce", } hidden_size: int = 1280 @@ -134,13 +134,13 @@ class Ernie4_5_VLMoeTextConfig(Ernie4_5_MoeConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } ignore_keys_at_rope_validation = {"mrope_section"} diff --git a/src/transformers/models/eurobert/configuration_eurobert.py b/src/transformers/models/eurobert/configuration_eurobert.py index f64c4f7e5a11..ab2c33a28828 100644 --- a/src/transformers/models/eurobert/configuration_eurobert.py +++ b/src/transformers/models/eurobert/configuration_eurobert.py @@ -54,10 +54,25 @@ class EuroBertConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -65,6 +80,12 @@ class EuroBertConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 128256 hidden_size: int = 768 intermediate_size: int = 3072 diff --git a/src/transformers/models/eurobert/modeling_eurobert.py b/src/transformers/models/eurobert/modeling_eurobert.py index b93dd0649f14..655d10ab123d 100644 --- a/src/transformers/models/eurobert/modeling_eurobert.py +++ b/src/transformers/models/eurobert/modeling_eurobert.py @@ -408,7 +408,9 @@ def forward( @auto_docstring class EuroBertForMaskedLM(EuroBertPreTrainedModel): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: EuroBertConfig): diff --git a/src/transformers/models/eurobert/modular_eurobert.py b/src/transformers/models/eurobert/modular_eurobert.py index 19f965c84f73..b40c4d043707 100644 --- a/src/transformers/models/eurobert/modular_eurobert.py +++ b/src/transformers/models/eurobert/modular_eurobert.py @@ -141,7 +141,9 @@ def forward( @auto_docstring class EuroBertForMaskedLM(EuroBertPreTrainedModel): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: EuroBertConfig): diff --git a/src/transformers/models/exaone4/configuration_exaone4.py b/src/transformers/models/exaone4/configuration_exaone4.py index f29cab8dd8ea..fd3bade9f662 100644 --- a/src/transformers/models/exaone4/configuration_exaone4.py +++ b/src/transformers/models/exaone4/configuration_exaone4.py @@ -63,12 +63,27 @@ class Exaone4Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -76,6 +91,12 @@ class Exaone4Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 102400 hidden_size: int = 4096 intermediate_size: int = 16384 diff --git a/src/transformers/models/exaone4/modeling_exaone4.py b/src/transformers/models/exaone4/modeling_exaone4.py index fab10b9b6937..d6aa12baedad 100644 --- a/src/transformers/models/exaone4/modeling_exaone4.py +++ b/src/transformers/models/exaone4/modeling_exaone4.py @@ -440,8 +440,10 @@ def forward( @auto_docstring class Exaone4ForCausalLM(Exaone4PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/exaone4/modular_exaone4.py b/src/transformers/models/exaone4/modular_exaone4.py index c6d9202170a0..7a88b77629f1 100644 --- a/src/transformers/models/exaone4/modular_exaone4.py +++ b/src/transformers/models/exaone4/modular_exaone4.py @@ -92,12 +92,27 @@ class Exaone4Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -105,6 +120,12 @@ class Exaone4Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 102400 hidden_size: int = 4096 intermediate_size: int = 16384 diff --git a/src/transformers/models/exaone_moe/configuration_exaone_moe.py b/src/transformers/models/exaone_moe/configuration_exaone_moe.py index 81ded9366cdb..83a873314700 100644 --- a/src/transformers/models/exaone_moe/configuration_exaone_moe.py +++ b/src/transformers/models/exaone_moe/configuration_exaone_moe.py @@ -69,19 +69,25 @@ class ExaoneMoeConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } + + base_model_sp_plan = None base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 102400 hidden_size: int = 4096 intermediate_size: int = 16384 diff --git a/src/transformers/models/exaone_moe/modeling_exaone_moe.py b/src/transformers/models/exaone_moe/modeling_exaone_moe.py index a7f80fc979c4..d2d8543f1417 100644 --- a/src/transformers/models/exaone_moe/modeling_exaone_moe.py +++ b/src/transformers/models/exaone_moe/modeling_exaone_moe.py @@ -563,8 +563,10 @@ def forward( @auto_docstring class ExaoneMoeForCausalLM(ExaoneMoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/exaone_moe/modular_exaone_moe.py b/src/transformers/models/exaone_moe/modular_exaone_moe.py index 75ec2b0bfd27..8d50c3db2295 100644 --- a/src/transformers/models/exaone_moe/modular_exaone_moe.py +++ b/src/transformers/models/exaone_moe/modular_exaone_moe.py @@ -80,6 +80,14 @@ class ExaoneMoeConfig(Exaone4Config): >>> configuration = model.config ```""" + base_model_sp_plan = None + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 102400 hidden_size: int = 4096 intermediate_size: int = 16384 diff --git a/src/transformers/models/falcon_h1/modeling_falcon_h1.py b/src/transformers/models/falcon_h1/modeling_falcon_h1.py index 3be9f882008e..2ea3c884437e 100644 --- a/src/transformers/models/falcon_h1/modeling_falcon_h1.py +++ b/src/transformers/models/falcon_h1/modeling_falcon_h1.py @@ -1166,8 +1166,10 @@ def _update_mamba_mask(self, attention_mask, past_key_values): @auto_docstring class FalconH1ForCausalLM(FalconH1PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/flex_olmo/configuration_flex_olmo.py b/src/transformers/models/flex_olmo/configuration_flex_olmo.py index 7b08a79b801b..a89d9f0452e8 100644 --- a/src/transformers/models/flex_olmo/configuration_flex_olmo.py +++ b/src/transformers/models/flex_olmo/configuration_flex_olmo.py @@ -50,13 +50,10 @@ class FlexOlmoConfig(PreTrainedConfig): attribute_map = {"num_local_experts": "num_experts"} default_theta = 500000.0 base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.k_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.v_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.o_proj": "rowwise_split_input", # input is replicated due to the added norm on q and k - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.q_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.k_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.v_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.o_proj": "vocab_allreduce", # input is replicated due to the added norm on q and k } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -64,6 +61,12 @@ class FlexOlmoConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 100352 hidden_size: int = 4096 intermediate_size: int = 11008 diff --git a/src/transformers/models/flex_olmo/modeling_flex_olmo.py b/src/transformers/models/flex_olmo/modeling_flex_olmo.py index 100e6fa35554..a9403a4110b0 100644 --- a/src/transformers/models/flex_olmo/modeling_flex_olmo.py +++ b/src/transformers/models/flex_olmo/modeling_flex_olmo.py @@ -597,8 +597,10 @@ def load_balancing_loss_func( @auto_docstring class FlexOlmoForCausalLM(FlexOlmoPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/flex_olmo/modular_flex_olmo.py b/src/transformers/models/flex_olmo/modular_flex_olmo.py index 01f32227f31f..c7b73dea8f15 100644 --- a/src/transformers/models/flex_olmo/modular_flex_olmo.py +++ b/src/transformers/models/flex_olmo/modular_flex_olmo.py @@ -60,13 +60,10 @@ class FlexOlmoConfig(PreTrainedConfig): attribute_map = {"num_local_experts": "num_experts"} default_theta = 500000.0 base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.k_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.v_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.o_proj": "rowwise_split_input", # input is replicated due to the added norm on q and k - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.q_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.k_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.v_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.o_proj": "vocab_allreduce", # input is replicated due to the added norm on q and k } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -74,6 +71,12 @@ class FlexOlmoConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 100352 hidden_size: int = 4096 intermediate_size: int = 11008 diff --git a/src/transformers/models/gemma/configuration_gemma.py b/src/transformers/models/gemma/configuration_gemma.py index 0ef2e8b31b8d..7b80a3fda63e 100644 --- a/src/transformers/models/gemma/configuration_gemma.py +++ b/src/transformers/models/gemma/configuration_gemma.py @@ -50,10 +50,25 @@ class GemmaConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -61,6 +76,12 @@ class GemmaConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 256000 hidden_size: int = 3072 intermediate_size: int = 24576 diff --git a/src/transformers/models/gemma/modeling_gemma.py b/src/transformers/models/gemma/modeling_gemma.py index c6c5a55b8790..f64702b55f19 100644 --- a/src/transformers/models/gemma/modeling_gemma.py +++ b/src/transformers/models/gemma/modeling_gemma.py @@ -450,8 +450,10 @@ def forward( @auto_docstring class GemmaForCausalLM(GemmaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/gemma/modular_gemma.py b/src/transformers/models/gemma/modular_gemma.py index 25f436473fbe..8a24552eaa36 100644 --- a/src/transformers/models/gemma/modular_gemma.py +++ b/src/transformers/models/gemma/modular_gemma.py @@ -69,10 +69,25 @@ class GemmaConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -80,6 +95,12 @@ class GemmaConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 256000 hidden_size: int = 3072 intermediate_size: int = 24576 diff --git a/src/transformers/models/gemma2/configuration_gemma2.py b/src/transformers/models/gemma2/configuration_gemma2.py index 11d7b012099f..8ff4fc88629a 100644 --- a/src/transformers/models/gemma2/configuration_gemma2.py +++ b/src/transformers/models/gemma2/configuration_gemma2.py @@ -54,10 +54,25 @@ class Gemma2Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -65,6 +80,12 @@ class Gemma2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 256000 hidden_size: int = 2304 intermediate_size: int = 9216 diff --git a/src/transformers/models/gemma2/modeling_gemma2.py b/src/transformers/models/gemma2/modeling_gemma2.py index 20673571b2d2..46eb89681619 100644 --- a/src/transformers/models/gemma2/modeling_gemma2.py +++ b/src/transformers/models/gemma2/modeling_gemma2.py @@ -476,8 +476,10 @@ def forward( @auto_docstring class Gemma2ForCausalLM(Gemma2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/gemma2/modular_gemma2.py b/src/transformers/models/gemma2/modular_gemma2.py index 2edd9ef5f101..d400331af615 100644 --- a/src/transformers/models/gemma2/modular_gemma2.py +++ b/src/transformers/models/gemma2/modular_gemma2.py @@ -81,10 +81,25 @@ class Gemma2Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -92,6 +107,12 @@ class Gemma2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 256000 hidden_size: int = 2304 intermediate_size: int = 9216 diff --git a/src/transformers/models/gemma3/configuration_gemma3.py b/src/transformers/models/gemma3/configuration_gemma3.py index f25c9cef21fd..adf71ed56bba 100644 --- a/src/transformers/models/gemma3/configuration_gemma3.py +++ b/src/transformers/models/gemma3/configuration_gemma3.py @@ -61,12 +61,27 @@ class Gemma3TextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -74,6 +89,12 @@ class Gemma3TextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 262_208 hidden_size: int = 2304 intermediate_size: int = 9216 @@ -99,6 +120,7 @@ class Gemma3TextConfig(PreTrainedConfig): final_logit_softcapping: float | None = None attn_logit_softcapping: float | None = None use_bidirectional_attention: bool | None = False + default_theta = {"global": 1_000_000.0, "local": 10_000.0} def __post_init__(self, **kwargs): diff --git a/src/transformers/models/gemma3/modeling_gemma3.py b/src/transformers/models/gemma3/modeling_gemma3.py index 1a7b79afcac1..b175dc6b67a4 100644 --- a/src/transformers/models/gemma3/modeling_gemma3.py +++ b/src/transformers/models/gemma3/modeling_gemma3.py @@ -585,8 +585,10 @@ def forward( @auto_docstring class Gemma3ForCausalLM(Gemma3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} config: Gemma3TextConfig def __init__(self, config: Gemma3TextConfig): diff --git a/src/transformers/models/gemma3/modular_gemma3.py b/src/transformers/models/gemma3/modular_gemma3.py index 82afbc884c65..e83243f21c4a 100644 --- a/src/transformers/models/gemma3/modular_gemma3.py +++ b/src/transformers/models/gemma3/modular_gemma3.py @@ -87,13 +87,35 @@ class Gemma3TextConfig(Gemma2Config, PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", + } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + default_theta = {"global": 1_000_000.0, "local": 10_000.0} vocab_size: int = 262_208 diff --git a/src/transformers/models/gemma3n/configuration_gemma3n.py b/src/transformers/models/gemma3n/configuration_gemma3n.py index e61c5f0038e7..973ff753c28b 100644 --- a/src/transformers/models/gemma3n/configuration_gemma3n.py +++ b/src/transformers/models/gemma3n/configuration_gemma3n.py @@ -78,23 +78,27 @@ class Gemma3nTextConfig(PreTrainedConfig): model_type = "gemma3n_text" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise", - "layers.*.self_attn.k_proj": "colwise", - "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.v_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.q_proj": "colwise_allgather", + "layers.*.self_attn.k_proj": "colwise_allgather", + "layers.*.self_attn.v_proj": "colwise_allgather", + "layers.*.self_attn.o_proj": "vocab_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } + base_model_sp_plan = None base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 262_400 hidden_size: int = 2048 intermediate_size: int | list[int] = 16_384 @@ -117,6 +121,7 @@ class Gemma3nTextConfig(PreTrainedConfig): sliding_window: int = 512 layer_types: list[str] | None = None final_logit_softcapping: float = 30.0 + default_theta = {"global": 1_000_000.0, "local": 10_000.0} vocab_size_per_layer_input: int = 262_144 hidden_size_per_layer_input: int = 256 diff --git a/src/transformers/models/gemma3n/modeling_gemma3n.py b/src/transformers/models/gemma3n/modeling_gemma3n.py index ff674c2eee72..dafdd8bd27a8 100644 --- a/src/transformers/models/gemma3n/modeling_gemma3n.py +++ b/src/transformers/models/gemma3n/modeling_gemma3n.py @@ -1827,8 +1827,10 @@ def project_per_layer_inputs( @auto_docstring(custom_intro="The base Gemma 3n language model with a language modeling head.") class Gemma3nForCausalLM(Gemma3nPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} config: Gemma3nTextConfig def __init__(self, config: Gemma3nTextConfig): diff --git a/src/transformers/models/gemma3n/modular_gemma3n.py b/src/transformers/models/gemma3n/modular_gemma3n.py index 4ce215996fb4..2cbd72a61ed0 100644 --- a/src/transformers/models/gemma3n/modular_gemma3n.py +++ b/src/transformers/models/gemma3n/modular_gemma3n.py @@ -116,17 +116,22 @@ class Gemma3nTextConfig(Gemma3TextConfig): model_type = "gemma3n_text" base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise", - "layers.*.self_attn.k_proj": "colwise", - "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.v_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.q_proj": "colwise_allgather", + "layers.*.self_attn.k_proj": "colwise_allgather", + "layers.*.self_attn.v_proj": "colwise_allgather", + "layers.*.self_attn.o_proj": "vocab_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } + base_model_sp_plan = None + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + default_theta = {"global": 1_000_000.0, "local": 10_000.0} vocab_size: int = 262_400 diff --git a/src/transformers/models/gemma4/configuration_gemma4.py b/src/transformers/models/gemma4/configuration_gemma4.py index fe38dc1739ab..0b4b24b00642 100644 --- a/src/transformers/models/gemma4/configuration_gemma4.py +++ b/src/transformers/models/gemma4/configuration_gemma4.py @@ -125,28 +125,25 @@ class Gemma4TextConfig(PreTrainedConfig): model_type = "gemma4_text" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise", - "layers.*.self_attn.k_proj": "colwise", - "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", + # q/k use allgather because gemma4 has q_norm/k_norm with full-sized weights + # that can't match sharded q/k outputs. + "layers.*.self_attn.q_proj": "colwise_allgather", + "layers.*.self_attn.k_proj": "colwise_allgather", + "layers.*.self_attn.v_proj": "colwise_allgather", + "layers.*.self_attn.o_proj": "vocab_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", - "layers.*.experts.gate_up_proj": "packed_colwise", - "layers.*.experts.down_proj": "rowwise", - "layers.*.experts": "moe_tp_experts", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_ep_plan = { # EP plan for google/gemma-4-26B-A4B-it: do not tp in attention (num_global_key_value_heads=2 too small to partition) "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", "layers.*.router": "ep_router", "layers.*.experts.gate_up_proj": "grouped_gemm", "layers.*.experts.down_proj": "grouped_gemm", - "layers.*.experts": "moe_tp_experts", + "layers.*.experts": "moe_experts_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -154,6 +151,12 @@ class Gemma4TextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 262_144 hidden_size: int = 2304 intermediate_size: int = 9216 @@ -241,12 +244,10 @@ class Gemma4VisionConfig(PreTrainedConfig): "encoder.layers.*.self_attn.q_proj": "colwise", "encoder.layers.*.self_attn.k_proj": "colwise", "encoder.layers.*.self_attn.v_proj": "colwise", - "encoder.layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "encoder.layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "encoder.layers.*.self_attn.o_proj": "rowwise", + "encoder.layers.*.self_attn.o_proj": "rowwise_allreduce", "encoder.layers.*.mlp.gate_proj": "colwise", "encoder.layers.*.mlp.up_proj": "colwise", - "encoder.layers.*.mlp.down_proj": "rowwise", + "encoder.layers.*.mlp.down_proj": "rowwise_allreduce", } default_theta = 100.0 diff --git a/src/transformers/models/gemma4/modeling_gemma4.py b/src/transformers/models/gemma4/modeling_gemma4.py index c620bb383f0f..95757537920c 100644 --- a/src/transformers/models/gemma4/modeling_gemma4.py +++ b/src/transformers/models/gemma4/modeling_gemma4.py @@ -1796,8 +1796,10 @@ def project_per_layer_inputs( @auto_docstring(custom_intro="The base Gemma 4 language model with a language modeling head.") class Gemma4ForCausalLM(Gemma4PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} config: Gemma4TextConfig base_model_prefix = "model" diff --git a/src/transformers/models/gemma4_assistant/modeling_gemma4_assistant.py b/src/transformers/models/gemma4_assistant/modeling_gemma4_assistant.py index 8b7b6b1a8b67..4c300686a0e7 100644 --- a/src/transformers/models/gemma4_assistant/modeling_gemma4_assistant.py +++ b/src/transformers/models/gemma4_assistant/modeling_gemma4_assistant.py @@ -110,6 +110,7 @@ def _init_weights(self, module): class Gemma4AssistantForCausalLM(Gemma4AssistantPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_gather_output"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: Gemma4AssistantConfig): diff --git a/src/transformers/models/glm/configuration_glm.py b/src/transformers/models/glm/configuration_glm.py index 98525012a23b..67275679ab8b 100644 --- a/src/transformers/models/glm/configuration_glm.py +++ b/src/transformers/models/glm/configuration_glm.py @@ -43,9 +43,23 @@ class GlmConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.gate_up_proj": "colwise_gather_output", # we need to replicate here due to the `chunk` operation - "layers.*.mlp.down_proj": "rowwise_split_input", # input is replicated due to the `chunk` operation + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.gate_up_proj": "packed_colwise", # fused gate/up shards stay local for chunk + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_up_proj": "packed_colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -53,6 +67,12 @@ class GlmConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 151552 hidden_size: int = 4096 intermediate_size: int = 13696 diff --git a/src/transformers/models/glm/modeling_glm.py b/src/transformers/models/glm/modeling_glm.py index 712202580943..461703ab81cd 100644 --- a/src/transformers/models/glm/modeling_glm.py +++ b/src/transformers/models/glm/modeling_glm.py @@ -439,8 +439,10 @@ def forward( @auto_docstring class GlmForCausalLM(GlmPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/glm4/configuration_glm4.py b/src/transformers/models/glm4/configuration_glm4.py index f33129607fec..8b2755323f68 100644 --- a/src/transformers/models/glm4/configuration_glm4.py +++ b/src/transformers/models/glm4/configuration_glm4.py @@ -43,9 +43,23 @@ class Glm4Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.gate_up_proj": "colwise_gather_output", # we need to replicate here due to the `chunk` operation - "layers.*.mlp.down_proj": "rowwise_split_input", # input is replicated due to the `chunk` operation + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.gate_up_proj": "colwise_allgather", # we need to replicate here due to the `chunk` operation + "layers.*.mlp.down_proj": "vocab_allreduce", # input is replicated due to the `chunk` operation + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_up_proj": "colwise_allgather", # fused gate/up needs full tensor for chunk + "layers.*.mlp.down_proj": "vocab_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -53,6 +67,12 @@ class Glm4Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 151552 hidden_size: int = 4096 intermediate_size: int = 13696 diff --git a/src/transformers/models/glm4/modeling_glm4.py b/src/transformers/models/glm4/modeling_glm4.py index e99930ae57f6..aa3d6e62a124 100644 --- a/src/transformers/models/glm4/modeling_glm4.py +++ b/src/transformers/models/glm4/modeling_glm4.py @@ -444,8 +444,10 @@ def forward( @auto_docstring class Glm4ForCausalLM(Glm4PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/glm4_moe/configuration_glm4_moe.py b/src/transformers/models/glm4_moe/configuration_glm4_moe.py index a18123e90b33..bd94a492bb66 100644 --- a/src/transformers/models/glm4_moe/configuration_glm4_moe.py +++ b/src/transformers/models/glm4_moe/configuration_glm4_moe.py @@ -57,22 +57,27 @@ class Glm4MoeConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", # NOTE(3outeille): This needs to be right after down_proj in the dict. Otherwise, the pattern model.layers.*.mlp.experts will have priority over model.layers.*.mlp.experts.down_proj which will assign a wrong TP plan. + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + attribute_map = { "num_local_experts": "n_routed_experts", } diff --git a/src/transformers/models/glm4_moe/modeling_glm4_moe.py b/src/transformers/models/glm4_moe/modeling_glm4_moe.py index cc5a564ab86f..730bf70db8ef 100644 --- a/src/transformers/models/glm4_moe/modeling_glm4_moe.py +++ b/src/transformers/models/glm4_moe/modeling_glm4_moe.py @@ -577,8 +577,10 @@ def forward( @auto_docstring class Glm4MoeForCausalLM(Glm4MoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/glm4_moe/modular_glm4_moe.py b/src/transformers/models/glm4_moe/modular_glm4_moe.py index 868018d744b5..41921eb080e6 100644 --- a/src/transformers/models/glm4_moe/modular_glm4_moe.py +++ b/src/transformers/models/glm4_moe/modular_glm4_moe.py @@ -70,22 +70,27 @@ class Glm4MoeConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", # NOTE(3outeille): This needs to be right after down_proj in the dict. Otherwise, the pattern model.layers.*.mlp.experts will have priority over model.layers.*.mlp.experts.down_proj which will assign a wrong TP plan. + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + attribute_map = { "num_local_experts": "n_routed_experts", } diff --git a/src/transformers/models/glm4_moe_lite/configuration_glm4_moe_lite.py b/src/transformers/models/glm4_moe_lite/configuration_glm4_moe_lite.py index a9518ed9c5d3..cc0c1fc67ffe 100644 --- a/src/transformers/models/glm4_moe_lite/configuration_glm4_moe_lite.py +++ b/src/transformers/models/glm4_moe_lite/configuration_glm4_moe_lite.py @@ -53,21 +53,25 @@ class Glm4MoeLiteConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { "layers.*.self_attn.q_b_proj": "colwise", - "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj", "layers.*.self_attn.kv_b_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + attribute_map = { "num_local_experts": "n_routed_experts", "head_dim": "qk_rope_head_dim", diff --git a/src/transformers/models/glm4_moe_lite/modeling_glm4_moe_lite.py b/src/transformers/models/glm4_moe_lite/modeling_glm4_moe_lite.py index 0b8ccc865775..4bab2f0316ae 100644 --- a/src/transformers/models/glm4_moe_lite/modeling_glm4_moe_lite.py +++ b/src/transformers/models/glm4_moe_lite/modeling_glm4_moe_lite.py @@ -651,8 +651,10 @@ def forward( @auto_docstring class Glm4MoeLiteForCausalLM(Glm4MoeLitePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/glm4_moe_lite/modular_glm4_moe_lite.py b/src/transformers/models/glm4_moe_lite/modular_glm4_moe_lite.py index 1f65f44a525a..ea8e5f60e491 100644 --- a/src/transformers/models/glm4_moe_lite/modular_glm4_moe_lite.py +++ b/src/transformers/models/glm4_moe_lite/modular_glm4_moe_lite.py @@ -61,21 +61,25 @@ class Glm4MoeLiteConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { "layers.*.self_attn.q_b_proj": "colwise", - "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj", "layers.*.self_attn.kv_b_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + attribute_map = { "num_local_experts": "n_routed_experts", "head_dim": "qk_rope_head_dim", diff --git a/src/transformers/models/glm4v/configuration_glm4v.py b/src/transformers/models/glm4v/configuration_glm4v.py index 4f151aa38156..dc993ab631d6 100644 --- a/src/transformers/models/glm4v/configuration_glm4v.py +++ b/src/transformers/models/glm4v/configuration_glm4v.py @@ -93,15 +93,22 @@ class Glm4vTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.gate_up_proj": "colwise_gather_output", # we need to replicate here due to the `chunk` operation - "layers.*.mlp.down_proj": "rowwise_split_input", # input is replicated due to the `chunk` operation + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.gate_up_proj": "colwise_allgather", # we need to replicate here due to the `chunk` operation + "layers.*.mlp.down_proj": "vocab_allreduce", # input is replicated due to the `chunk` operation } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + ignore_keys_at_rope_validation = {"mrope_section"} vocab_size: int = 151552 diff --git a/src/transformers/models/glm4v/modular_glm4v.py b/src/transformers/models/glm4v/modular_glm4v.py index 625d130310c6..880a772a537c 100644 --- a/src/transformers/models/glm4v/modular_glm4v.py +++ b/src/transformers/models/glm4v/modular_glm4v.py @@ -138,15 +138,22 @@ class Glm4vTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.gate_up_proj": "colwise_gather_output", # we need to replicate here due to the `chunk` operation - "layers.*.mlp.down_proj": "rowwise_split_input", # input is replicated due to the `chunk` operation + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.gate_up_proj": "colwise_allgather", # we need to replicate here due to the `chunk` operation + "layers.*.mlp.down_proj": "vocab_allreduce", # input is replicated due to the `chunk` operation } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + ignore_keys_at_rope_validation = {"mrope_section"} vocab_size: int = 151552 diff --git a/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py b/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py index 0e4d6a9cb191..73971460b2d4 100644 --- a/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py +++ b/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py @@ -56,16 +56,23 @@ class Glm4vMoeTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + attribute_map = { "num_local_experts": "n_routed_experts", } @@ -98,6 +105,7 @@ class Glm4vMoeTextConfig(PreTrainedConfig): eos_token_id: int | list[int] | None = None pad_token_id: int | None = None base_config_key = "text_config" + ignore_keys_at_rope_validation = {"mrope_section"} router_aux_loss_coef: float = 0.0001 diff --git a/src/transformers/models/glm4v_moe/modular_glm4v_moe.py b/src/transformers/models/glm4v_moe/modular_glm4v_moe.py index b1c1293a00b9..289dac43e160 100644 --- a/src/transformers/models/glm4v_moe/modular_glm4v_moe.py +++ b/src/transformers/models/glm4v_moe/modular_glm4v_moe.py @@ -88,16 +88,23 @@ class Glm4vMoeTextConfig(Glm4MoeConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + ignore_keys_at_rope_validation = {"mrope_section"} vocab_size: int = 151424 diff --git a/src/transformers/models/glm_image/configuration_glm_image.py b/src/transformers/models/glm_image/configuration_glm_image.py index fcd302e35560..e022da23e861 100644 --- a/src/transformers/models/glm_image/configuration_glm_image.py +++ b/src/transformers/models/glm_image/configuration_glm_image.py @@ -106,15 +106,22 @@ class GlmImageTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.gate_up_proj": "colwise_gather_output", # we need to replicate here due to the `chunk` operation - "layers.*.mlp.down_proj": "rowwise_split_input", # input is replicated due to the `chunk` operation + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.gate_up_proj": "colwise_allgather", # we need to replicate here due to the `chunk` operation + "layers.*.mlp.down_proj": "vocab_allreduce", # input is replicated due to the `chunk` operation } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + ignore_keys_at_rope_validation = {"mrope_section"} vocab_size: int = 168064 diff --git a/src/transformers/models/glm_image/modular_glm_image.py b/src/transformers/models/glm_image/modular_glm_image.py index 273861a5b61e..8ce758601acd 100644 --- a/src/transformers/models/glm_image/modular_glm_image.py +++ b/src/transformers/models/glm_image/modular_glm_image.py @@ -133,6 +133,12 @@ class GlmImageTextConfig(Glm4vTextConfig): >>> configuration = model.config ```""" + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 168064 max_position_embeddings: int = 131072 vision_vocab_size: int = 16512 diff --git a/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py index 8f11f42794b3..31512bde976a 100644 --- a/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py @@ -60,24 +60,27 @@ class GlmMoeDsaConfig(PreTrainedConfig): base_model_tp_plan = { "layers.*.self_attn.q_b_proj": "colwise", - "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj", "layers.*.self_attn.kv_b_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } attribute_map = { "num_local_experts": "n_routed_experts", "head_dim": "qk_rope_head_dim", diff --git a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py index 736dcdce32c3..1098096bd4e3 100644 --- a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py @@ -821,8 +821,10 @@ def forward( @auto_docstring class GlmMoeDsaForCausalLM(GlmMoeDsaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py index d909bb97e704..d45602619458 100644 --- a/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py @@ -108,18 +108,21 @@ class GlmMoeDsaConfig(Glm4MoeLiteConfig): base_model_tp_plan = { "layers.*.self_attn.q_b_proj": "colwise", - "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj", "layers.*.self_attn.kv_b_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", } hidden_size: int = 6144 diff --git a/src/transformers/models/glm_ocr/configuration_glm_ocr.py b/src/transformers/models/glm_ocr/configuration_glm_ocr.py index d08710c5e8b5..e1381b8bc79d 100644 --- a/src/transformers/models/glm_ocr/configuration_glm_ocr.py +++ b/src/transformers/models/glm_ocr/configuration_glm_ocr.py @@ -94,15 +94,22 @@ class GlmOcrTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.gate_up_proj": "colwise_gather_output", # we need to replicate here due to the `chunk` operation - "layers.*.mlp.down_proj": "rowwise_split_input", # input is replicated due to the `chunk` operation + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.gate_up_proj": "colwise_allgather", # we need to replicate here due to the `chunk` operation + "layers.*.mlp.down_proj": "vocab_allreduce", # input is replicated due to the `chunk` operation } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + ignore_keys_at_rope_validation = {"mrope_section"} vocab_size: int = 59392 diff --git a/src/transformers/models/glm_ocr/modular_glm_ocr.py b/src/transformers/models/glm_ocr/modular_glm_ocr.py index 782808bb69a9..2e9fc1c25a26 100644 --- a/src/transformers/models/glm_ocr/modular_glm_ocr.py +++ b/src/transformers/models/glm_ocr/modular_glm_ocr.py @@ -82,6 +82,12 @@ class GlmOcrTextConfig(Glm4vTextConfig): >>> configuration = model.config ```""" + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 59392 hidden_size: int = 1024 intermediate_size: int = 4096 diff --git a/src/transformers/models/glmasr/modeling_glmasr.py b/src/transformers/models/glmasr/modeling_glmasr.py index 74c762f5c0bf..df70cfb3884d 100644 --- a/src/transformers/models/glmasr/modeling_glmasr.py +++ b/src/transformers/models/glmasr/modeling_glmasr.py @@ -358,6 +358,7 @@ class GlmAsrForConditionalGeneration(GlmAsrPreTrainedModel, GenerationMixin): _keep_in_fp32_modules_strict = None _supports_attention_backend = True _tp_plan = None + _sp_plan = None _pp_plan = None def __init__(self, config): diff --git a/src/transformers/models/gpt_neox/configuration_gpt_neox.py b/src/transformers/models/gpt_neox/configuration_gpt_neox.py index 256ddca26b49..abd0b4578a72 100644 --- a/src/transformers/models/gpt_neox/configuration_gpt_neox.py +++ b/src/transformers/models/gpt_neox/configuration_gpt_neox.py @@ -47,9 +47,18 @@ class GPTNeoXConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { "layers.*.attention.query_key_value": "colwise", - "layers.*.attention.dense": "rowwise", + "layers.*.attention.dense": "rowwise_allreduce", "layers.*.mlp.dense_h_to_4h": "colwise", - "layers.*.mlp.dense_4h_to_h": "rowwise", + "layers.*.mlp.dense_4h_to_h": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.dense_h_to_4h": "colwise", + "layers.*.mlp.dense_4h_to_h": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_in": (["input_ids"], ["inputs_embeds"]), @@ -58,6 +67,12 @@ class GPTNeoXConfig(PreTrainedConfig): "final_layer_norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 50432 hidden_size: int = 6144 num_hidden_layers: int = 44 diff --git a/src/transformers/models/gpt_neox/modeling_gpt_neox.py b/src/transformers/models/gpt_neox/modeling_gpt_neox.py index 10e4b5922add..cfb81a96227d 100755 --- a/src/transformers/models/gpt_neox/modeling_gpt_neox.py +++ b/src/transformers/models/gpt_neox/modeling_gpt_neox.py @@ -394,7 +394,7 @@ def set_input_embeddings(self, value): ) class GPTNeoXForCausalLM(GPTNeoXPreTrainedModel, GenerationMixin): _tied_weights_keys = {"embed_out.weight": "gpt_neox.embed_in.weight"} - _tp_plan = {"embed_out": "colwise_gather_output"} + _tp_plan = {"embed_out": "colwise_allgather"} _pp_plan = {"embed_out": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/gpt_neox/modular_gpt_neox.py b/src/transformers/models/gpt_neox/modular_gpt_neox.py index f778501b7b38..cabda14021f0 100644 --- a/src/transformers/models/gpt_neox/modular_gpt_neox.py +++ b/src/transformers/models/gpt_neox/modular_gpt_neox.py @@ -341,7 +341,7 @@ def forward( ) class GPTNeoXForCausalLM(GPTNeoXPreTrainedModel, GenerationMixin): _tied_weights_keys = {"embed_out.weight": "gpt_neox.embed_in.weight"} - _tp_plan = {"embed_out": "colwise_gather_output"} + _tp_plan = {"embed_out": "colwise_allgather"} _pp_plan = {"embed_out": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/gpt_oss/configuration_gpt_oss.py b/src/transformers/models/gpt_oss/configuration_gpt_oss.py index 47c029a5bca9..c7cf246fe259 100644 --- a/src/transformers/models/gpt_oss/configuration_gpt_oss.py +++ b/src/transformers/models/gpt_oss/configuration_gpt_oss.py @@ -32,13 +32,18 @@ class GptOssConfig(PreTrainedConfig): "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } base_model_ep_plan = { "layers.*.mlp.router": "ep_router", "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", "layers.*.mlp.experts.gate_up_proj_bias": "grouped_gemm", "layers.*.mlp.experts.down_proj": "grouped_gemm", "layers.*.mlp.experts.down_proj_bias": "grouped_gemm", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts": "moe_experts_allreduce", } num_hidden_layers: int = 36 diff --git a/src/transformers/models/gpt_oss/modeling_gpt_oss.py b/src/transformers/models/gpt_oss/modeling_gpt_oss.py index 484b071f01bc..ebb980bd8a9b 100644 --- a/src/transformers/models/gpt_oss/modeling_gpt_oss.py +++ b/src/transformers/models/gpt_oss/modeling_gpt_oss.py @@ -592,8 +592,10 @@ def load_balancing_loss_func( @auto_docstring class GptOssForCausalLM(GptOssPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/granite/configuration_granite.py b/src/transformers/models/granite/configuration_granite.py index e026cbbe5ff3..1464157b962d 100644 --- a/src/transformers/models/granite/configuration_granite.py +++ b/src/transformers/models/granite/configuration_granite.py @@ -50,10 +50,25 @@ class GraniteConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -61,6 +76,12 @@ class GraniteConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 32000 hidden_size: int = 4096 intermediate_size: int = 11008 diff --git a/src/transformers/models/granite/modeling_granite.py b/src/transformers/models/granite/modeling_granite.py index 934345fe6723..5a266b3acc93 100644 --- a/src/transformers/models/granite/modeling_granite.py +++ b/src/transformers/models/granite/modeling_granite.py @@ -445,8 +445,10 @@ def forward( @auto_docstring class GraniteForCausalLM(GranitePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/granite4_vision/configuration_granite4_vision.py b/src/transformers/models/granite4_vision/configuration_granite4_vision.py index 82c9e6765515..d3fb9129e5d2 100644 --- a/src/transformers/models/granite4_vision/configuration_granite4_vision.py +++ b/src/transformers/models/granite4_vision/configuration_granite4_vision.py @@ -53,10 +53,25 @@ class Granite4VisionTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -64,6 +79,12 @@ class Granite4VisionTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 32000 hidden_size: int = 4096 intermediate_size: int = 11008 diff --git a/src/transformers/models/granitemoe/modeling_granitemoe.py b/src/transformers/models/granitemoe/modeling_granitemoe.py index 5fb53d6afe49..4f79272f164f 100644 --- a/src/transformers/models/granitemoe/modeling_granitemoe.py +++ b/src/transformers/models/granitemoe/modeling_granitemoe.py @@ -626,8 +626,10 @@ def load_balancing_loss_func( @auto_docstring class GraniteMoeForCausalLM(GraniteMoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config: GraniteMoeConfig): super().__init__(config) diff --git a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py index 2e0926f3e5d4..929a82f5ed82 100644 --- a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py +++ b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py @@ -1307,8 +1307,10 @@ def load_balancing_loss_func( @auto_docstring class GraniteMoeHybridForCausalLM(GraniteMoeHybridPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config: GraniteMoeHybridConfig): super().__init__(config) diff --git a/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py b/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py index 71f8c6eaff7d..8d0bc797a57b 100644 --- a/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py +++ b/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py @@ -695,8 +695,10 @@ def load_balancing_loss_func( @auto_docstring class GraniteMoeSharedForCausalLM(GraniteMoeSharedPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config: GraniteMoeSharedConfig): super().__init__(config) diff --git a/src/transformers/models/helium/configuration_helium.py b/src/transformers/models/helium/configuration_helium.py index caa966b58a9d..5aca27c00069 100644 --- a/src/transformers/models/helium/configuration_helium.py +++ b/src/transformers/models/helium/configuration_helium.py @@ -44,10 +44,25 @@ class HeliumConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -55,6 +70,12 @@ class HeliumConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 48000 hidden_size: int = 2560 intermediate_size: int = 7040 diff --git a/src/transformers/models/helium/modeling_helium.py b/src/transformers/models/helium/modeling_helium.py index 8283fcb19e28..555fdc7592d3 100644 --- a/src/transformers/models/helium/modeling_helium.py +++ b/src/transformers/models/helium/modeling_helium.py @@ -423,8 +423,10 @@ def forward( @auto_docstring class HeliumForCausalLM(HeliumPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/higgs_audio_v2/configuration_higgs_audio_v2.py b/src/transformers/models/higgs_audio_v2/configuration_higgs_audio_v2.py index 97823eb79576..d1cd942ae268 100644 --- a/src/transformers/models/higgs_audio_v2/configuration_higgs_audio_v2.py +++ b/src/transformers/models/higgs_audio_v2/configuration_higgs_audio_v2.py @@ -62,10 +62,25 @@ class HiggsAudioV2Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -73,6 +88,12 @@ class HiggsAudioV2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 128256 hidden_size: int = 3072 intermediate_size: int = 8192 diff --git a/src/transformers/models/hrm_text/configuration_hrm_text.py b/src/transformers/models/hrm_text/configuration_hrm_text.py index c8a328906b96..255391399960 100644 --- a/src/transformers/models/hrm_text/configuration_hrm_text.py +++ b/src/transformers/models/hrm_text/configuration_hrm_text.py @@ -65,12 +65,33 @@ class HrmTextConfig(PreTrainedConfig): **{f"{stack}.layers.*.mlp.up_proj": "colwise" for stack in ("L_module", "H_module")}, **{f"{stack}.layers.*.mlp.down_proj": "rowwise" for stack in ("L_module", "H_module")}, } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 151808 hidden_size: int = 1536 intermediate_size: int = 4096 diff --git a/src/transformers/models/hrm_text/modeling_hrm_text.py b/src/transformers/models/hrm_text/modeling_hrm_text.py index 9e10bed4997e..731f3cba1ebb 100644 --- a/src/transformers/models/hrm_text/modeling_hrm_text.py +++ b/src/transformers/models/hrm_text/modeling_hrm_text.py @@ -555,8 +555,10 @@ def forward( @auto_docstring class HrmTextForCausalLM(HrmTextPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/hubert/modeling_hubert.py b/src/transformers/models/hubert/modeling_hubert.py index 44ee98aeffbf..6924271e4736 100755 --- a/src/transformers/models/hubert/modeling_hubert.py +++ b/src/transformers/models/hubert/modeling_hubert.py @@ -27,8 +27,8 @@ from ... import initialization as init from ...activations import ACT2FN +from ...distributed.fsdp import is_fsdp_managed_module from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...masking_utils import create_bidirectional_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer diff --git a/src/transformers/models/hunyuan_v1_dense/modeling_hunyuan_v1_dense.py b/src/transformers/models/hunyuan_v1_dense/modeling_hunyuan_v1_dense.py index d1652d78cbbc..3b2386cf8ab3 100644 --- a/src/transformers/models/hunyuan_v1_dense/modeling_hunyuan_v1_dense.py +++ b/src/transformers/models/hunyuan_v1_dense/modeling_hunyuan_v1_dense.py @@ -461,8 +461,10 @@ def forward( @auto_docstring class HunYuanDenseV1ForCausalLM(HunYuanDenseV1PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/hunyuan_v1_moe/modeling_hunyuan_v1_moe.py b/src/transformers/models/hunyuan_v1_moe/modeling_hunyuan_v1_moe.py index 19779da0528c..db33809260af 100644 --- a/src/transformers/models/hunyuan_v1_moe/modeling_hunyuan_v1_moe.py +++ b/src/transformers/models/hunyuan_v1_moe/modeling_hunyuan_v1_moe.py @@ -550,8 +550,10 @@ def forward( @auto_docstring class HunYuanMoEV1ForCausalLM(HunYuanMoEV1PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/hy_v3/configuration_hy_v3.py b/src/transformers/models/hy_v3/configuration_hy_v3.py index e2eee94b118a..dddb7269f95e 100644 --- a/src/transformers/models/hy_v3/configuration_hy_v3.py +++ b/src/transformers/models/hy_v3/configuration_hy_v3.py @@ -54,18 +54,18 @@ class HYV3Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts.down_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -73,6 +73,12 @@ class HYV3Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 120832 hidden_size: int = 4096 intermediate_size: int = 13312 diff --git a/src/transformers/models/hy_v3/modeling_hy_v3.py b/src/transformers/models/hy_v3/modeling_hy_v3.py index f2d64736b4f0..6ade5c2d278d 100644 --- a/src/transformers/models/hy_v3/modeling_hy_v3.py +++ b/src/transformers/models/hy_v3/modeling_hy_v3.py @@ -544,8 +544,10 @@ def forward( @auto_docstring class HYV3ForCausalLM(HYV3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config: HYV3Config): super().__init__(config) diff --git a/src/transformers/models/hy_v3/modular_hy_v3.py b/src/transformers/models/hy_v3/modular_hy_v3.py index fa0931435197..5c9547a2fe0e 100644 --- a/src/transformers/models/hy_v3/modular_hy_v3.py +++ b/src/transformers/models/hy_v3/modular_hy_v3.py @@ -79,18 +79,18 @@ class HYV3Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts.down_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -98,6 +98,12 @@ class HYV3Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 120832 hidden_size: int = 4096 intermediate_size: int = 13312 diff --git a/src/transformers/models/hyperclovax/configuration_hyperclovax.py b/src/transformers/models/hyperclovax/configuration_hyperclovax.py index 430a56bf0249..1090a4e00fda 100644 --- a/src/transformers/models/hyperclovax/configuration_hyperclovax.py +++ b/src/transformers/models/hyperclovax/configuration_hyperclovax.py @@ -65,10 +65,25 @@ class HyperCLOVAXConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -76,6 +91,12 @@ class HyperCLOVAXConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 32000 hidden_size: int = 4096 intermediate_size: int = 11008 diff --git a/src/transformers/models/hyperclovax/modeling_hyperclovax.py b/src/transformers/models/hyperclovax/modeling_hyperclovax.py index 3608d215bfa9..a69fb46b121a 100644 --- a/src/transformers/models/hyperclovax/modeling_hyperclovax.py +++ b/src/transformers/models/hyperclovax/modeling_hyperclovax.py @@ -452,8 +452,10 @@ def forward( @auto_docstring class HyperCLOVAXForCausalLM(HyperCLOVAXPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/jais2/configuration_jais2.py b/src/transformers/models/jais2/configuration_jais2.py index 7139661d7575..6fa9783effd3 100644 --- a/src/transformers/models/jais2/configuration_jais2.py +++ b/src/transformers/models/jais2/configuration_jais2.py @@ -50,9 +50,23 @@ class Jais2Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -60,6 +74,12 @@ class Jais2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 150272 hidden_size: int = 3328 intermediate_size: int = 26624 diff --git a/src/transformers/models/jais2/modeling_jais2.py b/src/transformers/models/jais2/modeling_jais2.py index 5e6a37c0172d..1fb6c19d9044 100644 --- a/src/transformers/models/jais2/modeling_jais2.py +++ b/src/transformers/models/jais2/modeling_jais2.py @@ -397,8 +397,10 @@ def forward( @auto_docstring class Jais2ForCausalLM(Jais2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/jais2/modular_jais2.py b/src/transformers/models/jais2/modular_jais2.py index 3fbdf2c8cd46..c760e70c4b15 100644 --- a/src/transformers/models/jais2/modular_jais2.py +++ b/src/transformers/models/jais2/modular_jais2.py @@ -34,9 +34,23 @@ class Jais2Config(LlamaConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } vocab_size: int = 150272 diff --git a/src/transformers/models/jamba/modeling_jamba.py b/src/transformers/models/jamba/modeling_jamba.py index d9e3ff7b84ef..ba5a759c4636 100755 --- a/src/transformers/models/jamba/modeling_jamba.py +++ b/src/transformers/models/jamba/modeling_jamba.py @@ -843,8 +843,10 @@ def load_balancing_loss_func( @auto_docstring class JambaForCausalLM(JambaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config: JambaConfig): super().__init__(config) diff --git a/src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py b/src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py index b16274332baf..fbe51b44259d 100644 --- a/src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py +++ b/src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py @@ -874,8 +874,10 @@ def forward( @auto_docstring class KyutaiSpeechToTextForConditionalGeneration(KyutaiSpeechToTextPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} _keep_in_fp32_modules_strict = ["codec_model"] output_modalities = ("audio", "text") diff --git a/src/transformers/models/laguna/configuration_laguna.py b/src/transformers/models/laguna/configuration_laguna.py index 33f939f6db43..50e011ef08be 100644 --- a/src/transformers/models/laguna/configuration_laguna.py +++ b/src/transformers/models/laguna/configuration_laguna.py @@ -61,18 +61,18 @@ class LagunaConfig(PreTrainedConfig): "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.g_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts.down_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -80,6 +80,12 @@ class LagunaConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 100352 hidden_size: int = 2048 intermediate_size: int = 8192 diff --git a/src/transformers/models/laguna/modeling_laguna.py b/src/transformers/models/laguna/modeling_laguna.py index aa4060e77f5f..410fd9a04258 100644 --- a/src/transformers/models/laguna/modeling_laguna.py +++ b/src/transformers/models/laguna/modeling_laguna.py @@ -672,8 +672,10 @@ def load_balancing_loss_func( @auto_docstring class LagunaForCausalLM(LagunaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/laguna/modular_laguna.py b/src/transformers/models/laguna/modular_laguna.py index 945cd40a99b2..7587827acb2c 100644 --- a/src/transformers/models/laguna/modular_laguna.py +++ b/src/transformers/models/laguna/modular_laguna.py @@ -79,18 +79,18 @@ class LagunaConfig(Qwen2MoeConfig): "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.g_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts.down_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", } vocab_size: int = 100352 diff --git a/src/transformers/models/lfm2/modeling_lfm2.py b/src/transformers/models/lfm2/modeling_lfm2.py index ef753e3b2893..95af77e0141b 100644 --- a/src/transformers/models/lfm2/modeling_lfm2.py +++ b/src/transformers/models/lfm2/modeling_lfm2.py @@ -540,8 +540,10 @@ def forward( @auto_docstring class Lfm2ForCausalLM(Lfm2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py b/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py index 0369ae31b8ae..dc97dfb11454 100644 --- a/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py +++ b/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py @@ -630,8 +630,10 @@ def forward( @auto_docstring class Lfm2MoeForCausalLM(Lfm2MoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/llama/configuration_llama.py b/src/transformers/models/llama/configuration_llama.py index 6960a6970592..0f3a0e75f7f6 100644 --- a/src/transformers/models/llama/configuration_llama.py +++ b/src/transformers/models/llama/configuration_llama.py @@ -50,10 +50,25 @@ class LlamaConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -61,6 +76,12 @@ class LlamaConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 32000 hidden_size: int = 4096 intermediate_size: int = 11008 diff --git a/src/transformers/models/llama/modeling_llama.py b/src/transformers/models/llama/modeling_llama.py index 9d659c7c6f08..655e542358f5 100644 --- a/src/transformers/models/llama/modeling_llama.py +++ b/src/transformers/models/llama/modeling_llama.py @@ -428,8 +428,10 @@ def forward( @auto_docstring class LlamaForCausalLM(LlamaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/llama4/configuration_llama4.py b/src/transformers/models/llama4/configuration_llama4.py index 79cfd063f4d4..781de8a57fd0 100644 --- a/src/transformers/models/llama4/configuration_llama4.py +++ b/src/transformers/models/llama4/configuration_llama4.py @@ -45,10 +45,10 @@ class Llama4VisionConfig(PreTrainedConfig): "model.layers.*.self_attn.q_proj": "colwise", "model.layers.*.self_attn.k_proj": "colwise", "model.layers.*.self_attn.v_proj": "colwise", - "model.layers.*.self_attn.o_proj": "rowwise", + "model.layers.*.self_attn.o_proj": "rowwise_allreduce", "vision_adapter.mlp.fc1": "colwise", - "vision_adapter.mlp.fc2": "rowwise", - "patch_embedding.linear": "colwise_gather_output", + "vision_adapter.mlp.fc2": "rowwise_allreduce", + "patch_embedding.linear": "colwise_allgather", } model_type = "llama4_vision_model" base_config_key = "vision_config" @@ -113,26 +113,24 @@ class Llama4TextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.feed_forward.shared_expert.gate_proj": "colwise", "layers.*.feed_forward.shared_expert.up_proj": "colwise", - "layers.*.feed_forward.shared_expert.down_proj": "rowwise", - "layers.*.feed_forward.experts.gate_up_proj": "packed_rowwise", # row because not linear - "layers.*.feed_forward.experts.down_proj": "colwise", # col because not linear + "layers.*.feed_forward.shared_expert.down_proj": "rowwise_allreduce", "layers.*.feed_forward.gate_proj": "colwise", "layers.*.feed_forward.up_proj": "colwise", - "layers.*.feed_forward.down_proj": "rowwise", + "layers.*.feed_forward.down_proj": "rowwise_allreduce", } base_model_ep_plan = { "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.feed_forward.experts.gate_up_proj": "grouped_gemm", # row because not linear "layers.*.feed_forward.experts.down_proj": "grouped_gemm", # col because not linear "layers.*.feed_forward.gate_proj": "colwise", "layers.*.feed_forward.up_proj": "colwise", - "layers.*.feed_forward.down_proj": "rowwise", + "layers.*.feed_forward.down_proj": "rowwise_allreduce", "layers.*.feed_forward.router": "ep_router", } @@ -179,7 +177,7 @@ def __post_init__(self, **kwargs): default_no_rope_layers = [ int((layer_idx + 1) % self.no_rope_layer_interval != 0) for layer_idx in range(self.num_hidden_layers) ] - self.no_rope_layers = self.no_rope_layers if self.no_rope_layers else default_no_rope_layers + self.no_rope_layers = self.no_rope_layers or default_no_rope_layers self.head_dim = self.head_dim if self.head_dim is not None else self.hidden_size // self.num_attention_heads self.moe_layers = ( diff --git a/src/transformers/models/llama4/modeling_llama4.py b/src/transformers/models/llama4/modeling_llama4.py index e064e399e6c0..13dd361ecb49 100644 --- a/src/transformers/models/llama4/modeling_llama4.py +++ b/src/transformers/models/llama4/modeling_llama4.py @@ -590,7 +590,9 @@ class Llama4ForCausalLM(Llama4PreTrainedModel, GenerationMixin): _no_split_modules = ["Llama4TextDecoderLayer"] base_model_prefix = "language_model" _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} config: Llama4TextConfig def __init__(self, config: Llama4TextConfig): diff --git a/src/transformers/models/longcat_flash/configuration_longcat_flash.py b/src/transformers/models/longcat_flash/configuration_longcat_flash.py index 39e5a03338d8..ffeb6de9d01f 100644 --- a/src/transformers/models/longcat_flash/configuration_longcat_flash.py +++ b/src/transformers/models/longcat_flash/configuration_longcat_flash.py @@ -55,16 +55,11 @@ class LongcatFlashConfig(PreTrainedConfig): default_theta = 10000000.0 base_model_tp_plan = { "layers.*.self_attn.*.q_b_proj": "colwise", - "layers.*.self_attn.*.kv_a_proj_with_mqa": "mla_kv_a_proj", "layers.*.self_attn.*.kv_b_proj": "colwise", - "layers.*.self_attn.*.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts.identity_expert": "moe_identity_expert", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.*.o_proj": "rowwise_allreduce", "layers.*.mlps.*.gate_proj": "colwise", "layers.*.mlps.*.up_proj": "colwise", - "layers.*.mlps.*.down_proj": "rowwise", + "layers.*.mlps.*.down_proj": "rowwise_allreduce", } base_model_pp_plan = { @@ -73,6 +68,12 @@ class LongcatFlashConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 131072 hidden_size: int = 6144 num_hidden_layers: int = 56 diff --git a/src/transformers/models/longcat_flash/modeling_longcat_flash.py b/src/transformers/models/longcat_flash/modeling_longcat_flash.py index d5ac6e237742..4e7b3ad287cd 100644 --- a/src/transformers/models/longcat_flash/modeling_longcat_flash.py +++ b/src/transformers/models/longcat_flash/modeling_longcat_flash.py @@ -650,8 +650,10 @@ def forward( @auto_docstring class LongcatFlashForCausalLM(LongcatFlashPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} _keys_to_ignore_on_load_unexpected = [r"model\.mtp.*"] def __init__(self, config): diff --git a/src/transformers/models/minimax/configuration_minimax.py b/src/transformers/models/minimax/configuration_minimax.py index 9a3a0023725e..d60e6d0c2849 100644 --- a/src/transformers/models/minimax/configuration_minimax.py +++ b/src/transformers/models/minimax/configuration_minimax.py @@ -65,16 +65,21 @@ class MiniMaxConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + attribute_map = {"num_experts": "num_local_experts"} vocab_size: int = 32000 diff --git a/src/transformers/models/minimax/modeling_minimax.py b/src/transformers/models/minimax/modeling_minimax.py index 69497f83cad8..7da5e5867ce2 100644 --- a/src/transformers/models/minimax/modeling_minimax.py +++ b/src/transformers/models/minimax/modeling_minimax.py @@ -789,8 +789,10 @@ def load_balancing_loss_func( @auto_docstring class MiniMaxForCausalLM(MiniMaxPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/minimax/modular_minimax.py b/src/transformers/models/minimax/modular_minimax.py index 0bd400458129..8442eebd53e7 100644 --- a/src/transformers/models/minimax/modular_minimax.py +++ b/src/transformers/models/minimax/modular_minimax.py @@ -91,16 +91,21 @@ class MiniMaxConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + attribute_map = {"num_experts": "num_local_experts"} vocab_size: int = 32000 diff --git a/src/transformers/models/minimax_m2/configuration_minimax_m2.py b/src/transformers/models/minimax_m2/configuration_minimax_m2.py index 75acb8b755d7..6193523d95b3 100644 --- a/src/transformers/models/minimax_m2/configuration_minimax_m2.py +++ b/src/transformers/models/minimax_m2/configuration_minimax_m2.py @@ -48,19 +48,37 @@ class MiniMaxM2Config(PreTrainedConfig): model_type = "minimax_m2" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise_gather_output", - "layers.*.self_attn.k_proj": "colwise_gather_output", - "layers.*.self_attn.v_proj": "colwise_gather_output", - "layers.*.self_attn.o_proj": "rowwise_split_input", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.q_proj": "colwise_allgather", + "layers.*.self_attn.k_proj": "colwise_allgather", + "layers.*.self_attn.v_proj": "colwise_allgather", + "layers.*.self_attn.o_proj": "vocab_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise_allgather", + "layers.*.self_attn.k_proj": "colwise_allgather", + "layers.*.self_attn.v_proj": "colwise_allgather", + "layers.*.self_attn.o_proj": "vocab_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts": "moe_experts_allreduce", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + attribute_map = { "num_experts": "num_local_experts", } diff --git a/src/transformers/models/minimax_m2/modeling_minimax_m2.py b/src/transformers/models/minimax_m2/modeling_minimax_m2.py index d19274262810..370ea88c971e 100644 --- a/src/transformers/models/minimax_m2/modeling_minimax_m2.py +++ b/src/transformers/models/minimax_m2/modeling_minimax_m2.py @@ -588,8 +588,10 @@ def load_balancing_loss_func( @auto_docstring class MiniMaxM2ForCausalLM(MiniMaxM2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/minimax_m2/modular_minimax_m2.py b/src/transformers/models/minimax_m2/modular_minimax_m2.py index a9938a555c62..ca1b29125ccf 100644 --- a/src/transformers/models/minimax_m2/modular_minimax_m2.py +++ b/src/transformers/models/minimax_m2/modular_minimax_m2.py @@ -67,19 +67,37 @@ class MiniMaxM2Config(PreTrainedConfig): model_type = "minimax_m2" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise_gather_output", - "layers.*.self_attn.k_proj": "colwise_gather_output", - "layers.*.self_attn.v_proj": "colwise_gather_output", - "layers.*.self_attn.o_proj": "rowwise_split_input", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.q_proj": "colwise_allgather", + "layers.*.self_attn.k_proj": "colwise_allgather", + "layers.*.self_attn.v_proj": "colwise_allgather", + "layers.*.self_attn.o_proj": "vocab_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise_allgather", + "layers.*.self_attn.k_proj": "colwise_allgather", + "layers.*.self_attn.v_proj": "colwise_allgather", + "layers.*.self_attn.o_proj": "vocab_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts": "moe_experts_allreduce", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + attribute_map = { "num_experts": "num_local_experts", } diff --git a/src/transformers/models/ministral/configuration_ministral.py b/src/transformers/models/ministral/configuration_ministral.py index 4ff445e6808a..fa0de0f0671e 100644 --- a/src/transformers/models/ministral/configuration_ministral.py +++ b/src/transformers/models/ministral/configuration_ministral.py @@ -52,10 +52,25 @@ class MinistralConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -63,6 +78,12 @@ class MinistralConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 32000 hidden_size: int = 4096 intermediate_size: int = 14336 diff --git a/src/transformers/models/ministral/modeling_ministral.py b/src/transformers/models/ministral/modeling_ministral.py index af4f7fbeae59..d2faea85b8cc 100644 --- a/src/transformers/models/ministral/modeling_ministral.py +++ b/src/transformers/models/ministral/modeling_ministral.py @@ -429,8 +429,10 @@ def forward( @auto_docstring class MinistralForCausalLM(MinistralPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/ministral3/configuration_ministral3.py b/src/transformers/models/ministral3/configuration_ministral3.py index 76d55710d5ae..d486d2efe796 100644 --- a/src/transformers/models/ministral3/configuration_ministral3.py +++ b/src/transformers/models/ministral3/configuration_ministral3.py @@ -58,16 +58,37 @@ class Ministral3Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } ignore_keys_at_rope_validation = {"llama_4_scaling_beta", "max_position_embeddings"} vocab_size: int = 131072 diff --git a/src/transformers/models/ministral3/modeling_ministral3.py b/src/transformers/models/ministral3/modeling_ministral3.py index 6aacf4c8ce3a..066a71997913 100644 --- a/src/transformers/models/ministral3/modeling_ministral3.py +++ b/src/transformers/models/ministral3/modeling_ministral3.py @@ -412,8 +412,10 @@ def forward( @auto_docstring class Ministral3ForCausalLM(Ministral3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/mistral/configuration_mistral.py b/src/transformers/models/mistral/configuration_mistral.py index c57193f58d7b..0451d2f958a3 100644 --- a/src/transformers/models/mistral/configuration_mistral.py +++ b/src/transformers/models/mistral/configuration_mistral.py @@ -49,10 +49,25 @@ class MistralConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -60,6 +75,12 @@ class MistralConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 32000 hidden_size: int = 4096 intermediate_size: int = 14336 diff --git a/src/transformers/models/mistral/modeling_mistral.py b/src/transformers/models/mistral/modeling_mistral.py index b79dea36c9e9..d94ff3d6a312 100644 --- a/src/transformers/models/mistral/modeling_mistral.py +++ b/src/transformers/models/mistral/modeling_mistral.py @@ -401,8 +401,10 @@ def forward( @auto_docstring class MistralForCausalLM(MistralPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/mistral4/configuration_mistral4.py b/src/transformers/models/mistral4/configuration_mistral4.py index 0e16e0a14f45..c8218ffa5c57 100644 --- a/src/transformers/models/mistral4/configuration_mistral4.py +++ b/src/transformers/models/mistral4/configuration_mistral4.py @@ -47,21 +47,25 @@ class Mistral4Config(PreTrainedConfig): model_type = "mistral4" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } attribute_map = { "num_local_experts": "n_routed_experts", } diff --git a/src/transformers/models/mistral4/modeling_mistral4.py b/src/transformers/models/mistral4/modeling_mistral4.py index 006ddad187bf..4473107418ee 100644 --- a/src/transformers/models/mistral4/modeling_mistral4.py +++ b/src/transformers/models/mistral4/modeling_mistral4.py @@ -640,8 +640,10 @@ def forward( @auto_docstring class Mistral4ForCausalLM(Mistral4PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/mixtral/configuration_mixtral.py b/src/transformers/models/mixtral/configuration_mixtral.py index 240f24411031..41beb623719c 100644 --- a/src/transformers/models/mixtral/configuration_mixtral.py +++ b/src/transformers/models/mixtral/configuration_mixtral.py @@ -42,20 +42,41 @@ class MixtralConfig(PreTrainedConfig): model_type = "mixtral" keys_to_ignore_at_inference = ["past_key_values"] default_theta = 1000000.0 + # TP plan (for inference/generation). base_model_tp_plan = { "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + + # TP + Sequence Parallelism plan (for training). + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts": "moe_experts_allreduce", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + # FSDP2 plan (see Qwen3Config.base_model_fsdp_plan for shape rationale). + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } attribute_map = {"num_experts": "num_local_experts"} vocab_size: int = 32000 diff --git a/src/transformers/models/mixtral/modeling_mixtral.py b/src/transformers/models/mixtral/modeling_mixtral.py index 991851dbadd3..df71b44ebac2 100644 --- a/src/transformers/models/mixtral/modeling_mixtral.py +++ b/src/transformers/models/mixtral/modeling_mixtral.py @@ -580,8 +580,10 @@ def load_balancing_loss_func( @auto_docstring class MixtralForCausalLM(MixtralPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/mixtral/modular_mixtral.py b/src/transformers/models/mixtral/modular_mixtral.py index 139e580fbca7..438005523f7e 100644 --- a/src/transformers/models/mixtral/modular_mixtral.py +++ b/src/transformers/models/mixtral/modular_mixtral.py @@ -335,6 +335,7 @@ def forward( class MixtralForCausalLM(MistralForCausalLM): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/musicflamingo/modeling_musicflamingo.py b/src/transformers/models/musicflamingo/modeling_musicflamingo.py index bea73283f8af..45e3d9e707d8 100644 --- a/src/transformers/models/musicflamingo/modeling_musicflamingo.py +++ b/src/transformers/models/musicflamingo/modeling_musicflamingo.py @@ -202,6 +202,7 @@ class MusicFlamingoForConditionalGeneration(MusicFlamingoPreTrainedModel, Genera _keep_in_fp32_modules_strict = None _supports_attention_backend = True _tp_plan = None + _sp_plan = None _pp_plan = None def __init__(self, config: MusicFlamingoConfig): diff --git a/src/transformers/models/nanochat/configuration_nanochat.py b/src/transformers/models/nanochat/configuration_nanochat.py index 24a0ab7b6d09..12ee29c11ef6 100644 --- a/src/transformers/models/nanochat/configuration_nanochat.py +++ b/src/transformers/models/nanochat/configuration_nanochat.py @@ -46,10 +46,11 @@ class NanoChatConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.fc1": "colwise", - "layers.*.mlp.fc2": "rowwise", + "layers.*.mlp.fc2": "rowwise_allreduce", } + base_model_sp_plan = None vocab_size: int = 50304 hidden_size: int = 768 diff --git a/src/transformers/models/nanochat/modeling_nanochat.py b/src/transformers/models/nanochat/modeling_nanochat.py index 9205b89cd360..a191eb2d5425 100644 --- a/src/transformers/models/nanochat/modeling_nanochat.py +++ b/src/transformers/models/nanochat/modeling_nanochat.py @@ -432,8 +432,10 @@ def forward( @auto_docstring class NanoChatForCausalLM(NanoChatPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/nanochat/modular_nanochat.py b/src/transformers/models/nanochat/modular_nanochat.py index 713cc29b81eb..460cf321150f 100644 --- a/src/transformers/models/nanochat/modular_nanochat.py +++ b/src/transformers/models/nanochat/modular_nanochat.py @@ -198,7 +198,9 @@ def forward( @auto_docstring class NanoChatForCausalLM(Gemma2ForCausalLM): - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} def forward(self, **super_kwargs) -> CausalLMOutputWithPast: r""" diff --git a/src/transformers/models/nllb_moe/modeling_nllb_moe.py b/src/transformers/models/nllb_moe/modeling_nllb_moe.py index d4ef4927ce67..7370722e38b0 100644 --- a/src/transformers/models/nllb_moe/modeling_nllb_moe.py +++ b/src/transformers/models/nllb_moe/modeling_nllb_moe.py @@ -22,9 +22,9 @@ from ... import initialization as init from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache +from ...distributed.fsdp import is_fsdp_managed_module from ...generation import GenerationMixin from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...masking_utils import create_bidirectional_mask, create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer diff --git a/src/transformers/models/olmo/configuration_olmo.py b/src/transformers/models/olmo/configuration_olmo.py index 186cc3a704fb..f9f32c1ed15d 100644 --- a/src/transformers/models/olmo/configuration_olmo.py +++ b/src/transformers/models/olmo/configuration_olmo.py @@ -53,10 +53,25 @@ class OlmoConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -64,6 +79,12 @@ class OlmoConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 50304 hidden_size: int = 4096 intermediate_size: int = 11008 diff --git a/src/transformers/models/olmo/modeling_olmo.py b/src/transformers/models/olmo/modeling_olmo.py index 0a886321bfc3..dad24db7ce06 100644 --- a/src/transformers/models/olmo/modeling_olmo.py +++ b/src/transformers/models/olmo/modeling_olmo.py @@ -425,8 +425,10 @@ def forward( @auto_docstring class OlmoForCausalLM(OlmoPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/olmo2/configuration_olmo2.py b/src/transformers/models/olmo2/configuration_olmo2.py index f879c0b8367f..719db72c42e1 100644 --- a/src/transformers/models/olmo2/configuration_olmo2.py +++ b/src/transformers/models/olmo2/configuration_olmo2.py @@ -53,13 +53,30 @@ class Olmo2Config(PreTrainedConfig): model_type = "olmo2" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.k_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.v_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.o_proj": "rowwise_split_input", # input is replicated due to the added norm on q and k + "layers.*.self_attn.q_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.k_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.v_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.o_proj": "vocab_allreduce", # input is replicated due to the added norm on q and k "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -67,6 +84,12 @@ class Olmo2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 50304 hidden_size: int = 4096 intermediate_size: int = 11008 diff --git a/src/transformers/models/olmo2/modeling_olmo2.py b/src/transformers/models/olmo2/modeling_olmo2.py index 9901212ff737..12d31e551870 100644 --- a/src/transformers/models/olmo2/modeling_olmo2.py +++ b/src/transformers/models/olmo2/modeling_olmo2.py @@ -429,8 +429,10 @@ def forward( @auto_docstring class Olmo2ForCausalLM(Olmo2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/olmo2/modular_olmo2.py b/src/transformers/models/olmo2/modular_olmo2.py index 4ac66c2e4608..3407011afffe 100644 --- a/src/transformers/models/olmo2/modular_olmo2.py +++ b/src/transformers/models/olmo2/modular_olmo2.py @@ -67,13 +67,30 @@ class Olmo2Config(OlmoConfig): model_type = "olmo2" base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.k_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.v_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.o_proj": "rowwise_split_input", # input is replicated due to the added norm on q and k + "layers.*.self_attn.q_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.k_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.v_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.o_proj": "vocab_allreduce", # input is replicated due to the added norm on q and k "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/olmo3/configuration_olmo3.py b/src/transformers/models/olmo3/configuration_olmo3.py index 2f45be450a0b..cf81a4b3e332 100644 --- a/src/transformers/models/olmo3/configuration_olmo3.py +++ b/src/transformers/models/olmo3/configuration_olmo3.py @@ -48,13 +48,30 @@ class Olmo3Config(PreTrainedConfig): model_type = "olmo3" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.k_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.v_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.o_proj": "rowwise_split_input", # input is replicated due to the added norm on q and k + "layers.*.self_attn.q_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.k_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.v_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.o_proj": "vocab_allreduce", # input is replicated due to the added norm on q and k "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -62,6 +79,12 @@ class Olmo3Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 50304 hidden_size: int = 4096 intermediate_size: int = 11008 diff --git a/src/transformers/models/olmo3/modeling_olmo3.py b/src/transformers/models/olmo3/modeling_olmo3.py index 78ade3570f97..8325b651312a 100644 --- a/src/transformers/models/olmo3/modeling_olmo3.py +++ b/src/transformers/models/olmo3/modeling_olmo3.py @@ -433,8 +433,10 @@ def forward( @auto_docstring class Olmo3ForCausalLM(Olmo3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/olmo3/modular_olmo3.py b/src/transformers/models/olmo3/modular_olmo3.py index f5934880c5f5..b587e4a2de30 100644 --- a/src/transformers/models/olmo3/modular_olmo3.py +++ b/src/transformers/models/olmo3/modular_olmo3.py @@ -63,13 +63,30 @@ class Olmo3Config(Olmo2Config): model_type = "olmo3" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.k_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.v_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.o_proj": "rowwise_split_input", # input is replicated due to the added norm on q and k + "layers.*.self_attn.q_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.k_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.v_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.o_proj": "vocab_allreduce", # input is replicated due to the added norm on q and k "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py index 0f7c32e8799e..106fb2010bfc 100644 --- a/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py +++ b/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py @@ -73,20 +73,27 @@ class OlmoHybridConfig(PreTrainedConfig): model_type = "olmo_hybrid" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.k_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.v_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.o_proj": "rowwise_split_input", # input is replicated due to the added norm on q and k + "layers.*.self_attn.q_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.k_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.v_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.o_proj": "vocab_allreduce", # input is replicated due to the added norm on q and k "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } + base_model_sp_plan = None base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 100352 hidden_size: int = 3840 intermediate_size: int = 11008 diff --git a/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py index 5f557dccc12d..081149371067 100644 --- a/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py +++ b/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py @@ -1045,8 +1045,10 @@ def _update_linear_attn_mask(self, attention_mask, past_key_values): @auto_docstring class OlmoHybridForCausalLM(OlmoHybridPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py index ef805ce9498a..df028e534aa3 100644 --- a/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py +++ b/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py @@ -120,14 +120,15 @@ class OlmoHybridConfig(LlamaConfig): model_type = "olmo_hybrid" base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.k_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.v_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.o_proj": "rowwise_split_input", # input is replicated due to the added norm on q and k + "layers.*.self_attn.q_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.k_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.v_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.o_proj": "vocab_allreduce", # input is replicated due to the added norm on q and k "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } + base_model_sp_plan = None vocab_size: int = 100352 hidden_size: int = 3840 diff --git a/src/transformers/models/olmoe/configuration_olmoe.py b/src/transformers/models/olmoe/configuration_olmoe.py index 16bedbe698f8..feedb02bfda4 100644 --- a/src/transformers/models/olmoe/configuration_olmoe.py +++ b/src/transformers/models/olmoe/configuration_olmoe.py @@ -46,13 +46,30 @@ class OlmoeConfig(PreTrainedConfig): # Default tensor parallel plan for base model `Olmoe` base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise_gather_output", # due to the norm, we have to gather - "layers.*.self_attn.k_proj": "colwise_gather_output", # due to the norm, we have to gather - "layers.*.self_attn.v_proj": "colwise_gather_output", # due to the norm, we have to gather - "layers.*.self_attn.o_proj": "rowwise_split_input", # due to the norm, we have to gather - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.q_proj": "colwise_allgather", # due to the norm, we have to gather + "layers.*.self_attn.k_proj": "colwise_allgather", # due to the norm, we have to gather + "layers.*.self_attn.v_proj": "colwise_allgather", # due to the norm, we have to gather + "layers.*.self_attn.o_proj": "vocab_allreduce", # due to the norm, we have to gather + "layers.*.mlp.experts": "moe_experts_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise_allgather", + "layers.*.self_attn.k_proj": "colwise_allgather", + "layers.*.self_attn.v_proj": "colwise_allgather", + "layers.*.self_attn.o_proj": "vocab_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts": "moe_experts_allreduce", + "norm": "activation", + } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", } vocab_size: int = 50304 diff --git a/src/transformers/models/olmoe/modeling_olmoe.py b/src/transformers/models/olmoe/modeling_olmoe.py index 5d89ec741529..10950656c462 100644 --- a/src/transformers/models/olmoe/modeling_olmoe.py +++ b/src/transformers/models/olmoe/modeling_olmoe.py @@ -604,8 +604,10 @@ def load_balancing_loss_func( @auto_docstring class OlmoeForCausalLM(OlmoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/openai_privacy_filter/configuration_openai_privacy_filter.py b/src/transformers/models/openai_privacy_filter/configuration_openai_privacy_filter.py index e7aaefde4bca..5f2933c91984 100644 --- a/src/transformers/models/openai_privacy_filter/configuration_openai_privacy_filter.py +++ b/src/transformers/models/openai_privacy_filter/configuration_openai_privacy_filter.py @@ -56,16 +56,23 @@ class OpenAIPrivacyFilterConfig(PreTrainedConfig): "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } base_model_ep_plan = { "layers.*.mlp.router": "ep_router", "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", "layers.*.mlp.experts.gate_up_proj_bias": "grouped_gemm", "layers.*.mlp.experts.down_proj": "grouped_gemm", "layers.*.mlp.experts.down_proj_bias": "grouped_gemm", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts": "moe_experts_allreduce", } num_hidden_layers: int = 8 num_local_experts: int = 128 + vocab_size: int = 200064 hidden_size: int = 640 intermediate_size: int = 640 diff --git a/src/transformers/models/openai_privacy_filter/modular_openai_privacy_filter.py b/src/transformers/models/openai_privacy_filter/modular_openai_privacy_filter.py index 422235d9da91..7a438fedd042 100644 --- a/src/transformers/models/openai_privacy_filter/modular_openai_privacy_filter.py +++ b/src/transformers/models/openai_privacy_filter/modular_openai_privacy_filter.py @@ -72,6 +72,13 @@ @strict class OpenAIPrivacyFilterConfig(GptOssConfig): model_type = "openai_privacy_filter" + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 200064 hidden_size: int = 640 intermediate_size: int = 640 diff --git a/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py b/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py index 343a22ade814..b9bef13987ac 100644 --- a/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py +++ b/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py @@ -98,10 +98,25 @@ class PaddleOCRTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -109,6 +124,12 @@ class PaddleOCRTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 103424 hidden_size: int = 1024 intermediate_size: int = 3072 diff --git a/src/transformers/models/phi/configuration_phi.py b/src/transformers/models/phi/configuration_phi.py index 2a65f0f16aec..bee22ad6a0b7 100644 --- a/src/transformers/models/phi/configuration_phi.py +++ b/src/transformers/models/phi/configuration_phi.py @@ -49,9 +49,23 @@ class PhiConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.dense": "rowwise", + "layers.*.self_attn.dense": "rowwise_allreduce", "layers.*.mlp.fc1": "colwise", - "layers.*.mlp.fc2": "rowwise", + "layers.*.mlp.fc2": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.dense": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.fc1": "colwise", + "layers.*.mlp.fc2": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -60,14 +74,20 @@ class PhiConfig(PreTrainedConfig): "final_layernorm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 51200 hidden_size: int = 2048 intermediate_size: int = 8192 num_hidden_layers: int = 24 num_attention_heads: int = 32 num_key_value_heads: int | None = None - resid_pdrop: float | int = 0.0 - embd_pdrop: float | int = 0.0 + resid_pdrop: float = 0.0 + embd_pdrop: float = 0.0 attention_dropout: float | int | None = 0.0 hidden_act: str = "gelu_new" max_position_embeddings: int = 2048 diff --git a/src/transformers/models/phi/modeling_phi.py b/src/transformers/models/phi/modeling_phi.py index e3f97a01ee4c..0acf6a8060aa 100644 --- a/src/transformers/models/phi/modeling_phi.py +++ b/src/transformers/models/phi/modeling_phi.py @@ -406,8 +406,10 @@ def forward( @auto_docstring class PhiForCausalLM(PhiPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/phi3/configuration_phi3.py b/src/transformers/models/phi3/configuration_phi3.py index f85502f8205d..d1288d411863 100644 --- a/src/transformers/models/phi3/configuration_phi3.py +++ b/src/transformers/models/phi3/configuration_phi3.py @@ -47,10 +47,22 @@ class Phi3Config(PreTrainedConfig): model_type = "phi3" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.qkv_proj": "colwise_gather_output", # we need to replicate here due to the slicing of qkv - "layers.*.self_attn.o_proj": "rowwise_split_input", # input is replicated due to the slicing of qkv - "layers.*.mlp.gate_up_proj": "colwise_gather_output", # we need to replicate here due to the `chunk` operation - "layers.*.mlp.down_proj": "rowwise_split_input", # input is replicated due to the `chunk` operation + "layers.*.self_attn.qkv_proj": "colwise_allgather", # we need to replicate here due to the slicing of qkv + "layers.*.self_attn.o_proj": "vocab_allreduce", # input is replicated due to the slicing of qkv + "layers.*.mlp.gate_up_proj": "colwise_allgather", # we need to replicate here due to the `chunk` operation + "layers.*.mlp.down_proj": "vocab_allreduce", # input is replicated due to the `chunk` operation + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.qkv_proj": "colwise_allgather", # fused qkv needs full tensor for slicing + "layers.*.self_attn.o_proj": "vocab_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_up_proj": "colwise_allgather", # fused gate/up needs full tensor for chunk + "layers.*.mlp.down_proj": "vocab_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -58,6 +70,12 @@ class Phi3Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 32064 hidden_size: int = 3072 intermediate_size: int = 8192 diff --git a/src/transformers/models/phi3/modeling_phi3.py b/src/transformers/models/phi3/modeling_phi3.py index b07735f8a2e6..f0d6726b73d9 100644 --- a/src/transformers/models/phi3/modeling_phi3.py +++ b/src/transformers/models/phi3/modeling_phi3.py @@ -432,8 +432,10 @@ def forward( @auto_docstring class Phi3ForCausalLM(Phi3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py b/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py index b73741305566..264119a82bec 100644 --- a/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py +++ b/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py @@ -181,10 +181,22 @@ class Phi4MultimodalConfig(PreTrainedConfig): model_type = "phi4_multimodal" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.qkv_proj": "colwise_gather_output", # we need to replicate here due to the slicing of qkv - "layers.*.self_attn.o_proj": "rowwise_split_input", # input is replicated due to the slicing of qkv - "layers.*.mlp.gate_up_proj": "colwise_gather_output", # we need to replicate here due to the `chunk` operation - "layers.*.mlp.down_proj": "rowwise_split_input", # input is replicated due to the `chunk` operation + "layers.*.self_attn.qkv_proj": "colwise_allgather", # we need to replicate here due to the slicing of qkv + "layers.*.self_attn.o_proj": "vocab_allreduce", # input is replicated due to the slicing of qkv + "layers.*.mlp.gate_up_proj": "colwise_allgather", # we need to replicate here due to the `chunk` operation + "layers.*.mlp.down_proj": "vocab_allreduce", # input is replicated due to the `chunk` operation + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.qkv_proj": "colwise_allgather", # fused qkv needs full tensor for slicing + "layers.*.self_attn.o_proj": "vocab_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_up_proj": "colwise_allgather", # fused gate/up needs full tensor for chunk + "layers.*.mlp.down_proj": "vocab_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -192,6 +204,12 @@ class Phi4MultimodalConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 200064 hidden_size: int = 3072 intermediate_size: int = 8192 diff --git a/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py b/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py index 9e6c0339098d..383445a29f8b 100644 --- a/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py +++ b/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py @@ -1596,8 +1596,10 @@ def forward( @auto_docstring class Phi4MultimodalForCausalLM(Phi4MultimodalPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/phimoe/configuration_phimoe.py b/src/transformers/models/phimoe/configuration_phimoe.py index e20f94085be0..0545ffe5ea70 100644 --- a/src/transformers/models/phimoe/configuration_phimoe.py +++ b/src/transformers/models/phimoe/configuration_phimoe.py @@ -47,6 +47,11 @@ class PhimoeConfig(PreTrainedConfig): model_type = "phimoe" keys_to_ignore_at_inference = ["past_key_values"] default_theta = 1000000.0 + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } vocab_size: int = 32064 hidden_size: int = 4096 diff --git a/src/transformers/models/phimoe/modeling_phimoe.py b/src/transformers/models/phimoe/modeling_phimoe.py index 23bc944c522a..d933a6ec69d9 100644 --- a/src/transformers/models/phimoe/modeling_phimoe.py +++ b/src/transformers/models/phimoe/modeling_phimoe.py @@ -772,8 +772,10 @@ def load_balancing_loss_func( @auto_docstring class PhimoeForCausalLM(PhimoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/pi0/modeling_pi0.py b/src/transformers/models/pi0/modeling_pi0.py index bd86b6f10644..da1e4a9c8b35 100644 --- a/src/transformers/models/pi0/modeling_pi0.py +++ b/src/transformers/models/pi0/modeling_pi0.py @@ -222,7 +222,7 @@ def forward( class PI0ForConditionalGeneration(PI0PreTrainedModel): """PI0 model with action projection heads and flow matching.""" - _tp_plan = {"action_out_proj": "colwise_gather_output"} + _tp_plan = {"action_out_proj": "colwise_allgather"} def __init__(self, config: PI0Config): super().__init__(config) diff --git a/src/transformers/models/pi0/modular_pi0.py b/src/transformers/models/pi0/modular_pi0.py index 337ae9406bb4..53472a40ccfe 100644 --- a/src/transformers/models/pi0/modular_pi0.py +++ b/src/transformers/models/pi0/modular_pi0.py @@ -482,7 +482,7 @@ def forward( class PI0ForConditionalGeneration(PI0PreTrainedModel): """PI0 model with action projection heads and flow matching.""" - _tp_plan = {"action_out_proj": "colwise_gather_output"} + _tp_plan = {"action_out_proj": "colwise_allgather"} def __init__(self, config: PI0Config): super().__init__(config) diff --git a/src/transformers/models/qwen2/configuration_qwen2.py b/src/transformers/models/qwen2/configuration_qwen2.py index ae41af7b211f..9106216fb13d 100644 --- a/src/transformers/models/qwen2/configuration_qwen2.py +++ b/src/transformers/models/qwen2/configuration_qwen2.py @@ -47,10 +47,25 @@ class Qwen2Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -58,6 +73,12 @@ class Qwen2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 151936 hidden_size: int = 4096 intermediate_size: int = 22016 diff --git a/src/transformers/models/qwen2/modeling_qwen2.py b/src/transformers/models/qwen2/modeling_qwen2.py index 9263e1d42937..7dcc274a2cf2 100644 --- a/src/transformers/models/qwen2/modeling_qwen2.py +++ b/src/transformers/models/qwen2/modeling_qwen2.py @@ -416,8 +416,10 @@ def forward( @auto_docstring class Qwen2ForCausalLM(Qwen2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/qwen2_5_omni/configuration_qwen2_5_omni.py b/src/transformers/models/qwen2_5_omni/configuration_qwen2_5_omni.py index 081823bf222f..fc90da09bbbf 100644 --- a/src/transformers/models/qwen2_5_omni/configuration_qwen2_5_omni.py +++ b/src/transformers/models/qwen2_5_omni/configuration_qwen2_5_omni.py @@ -157,16 +157,23 @@ class Qwen2_5OmniTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + ignore_keys_at_rope_validation = {"mrope_section"} vocab_size: int = 152064 diff --git a/src/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py b/src/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py index acebc79a3c94..e9eeefc9b9f8 100644 --- a/src/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py +++ b/src/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py @@ -259,16 +259,23 @@ class Qwen2_5OmniTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + ignore_keys_at_rope_validation = {"mrope_section"} vocab_size: int = 152064 diff --git a/src/transformers/models/qwen2_5_vl/configuration_qwen2_5_vl.py b/src/transformers/models/qwen2_5_vl/configuration_qwen2_5_vl.py index 385f2ecad057..da05af794ed5 100644 --- a/src/transformers/models/qwen2_5_vl/configuration_qwen2_5_vl.py +++ b/src/transformers/models/qwen2_5_vl/configuration_qwen2_5_vl.py @@ -91,16 +91,22 @@ class Qwen2_5_VLTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } ignore_keys_at_rope_validation = {"mrope_section"} vocab_size: int = 152064 diff --git a/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py b/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py index 7f961976b58c..329a6ef2a058 100644 --- a/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py +++ b/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py @@ -57,10 +57,10 @@ class Qwen2MoeConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -68,6 +68,12 @@ class Qwen2MoeConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 151936 hidden_size: int = 2048 intermediate_size: int = 5632 diff --git a/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py b/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py index d4150d0a74d7..f7142b4650fb 100644 --- a/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py +++ b/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py @@ -617,8 +617,10 @@ def load_balancing_loss_func( @auto_docstring class Qwen2MoeForCausalLM(Qwen2MoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/qwen2_moe/modular_qwen2_moe.py b/src/transformers/models/qwen2_moe/modular_qwen2_moe.py index deb615c9e7b6..9c5b65f47896 100644 --- a/src/transformers/models/qwen2_moe/modular_qwen2_moe.py +++ b/src/transformers/models/qwen2_moe/modular_qwen2_moe.py @@ -230,7 +230,9 @@ def forward( class Qwen2MoeForCausalLM(MixtralForCausalLM, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py b/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py index 65680d492a78..bb6871f46a05 100644 --- a/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py +++ b/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py @@ -68,16 +68,22 @@ class Qwen2VLTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } ignore_keys_at_rope_validation = {"mrope_section"} vocab_size: int = 152064 diff --git a/src/transformers/models/qwen3/configuration_qwen3.py b/src/transformers/models/qwen3/configuration_qwen3.py index 07ad0bb24b33..ff0cd8f63a23 100644 --- a/src/transformers/models/qwen3/configuration_qwen3.py +++ b/src/transformers/models/qwen3/configuration_qwen3.py @@ -41,17 +41,41 @@ class Qwen3Config(PreTrainedConfig): model_type = "qwen3" keys_to_ignore_at_inference = ["past_key_values"] - # Default tensor parallel plan for base model `Qwen3` + # TP plan (for inference/generation). + # All activations are plain tensors — compatible with KV cache and autoregressive + # decode (seq_len=1). Each rank holds a full copy of activations between layers. base_model_tp_plan = { "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + + # TP + Sequence Parallelism plan (for training). + # Activations between layers are sharded on the sequence dimension (Shard(1)), + # reducing per-rank activation memory by tp_size. In exchange, extra collectives + # (all-gather before attention/MLP, reduce-scatter after) are needed. + # Not compatible with autoregressive decode (because seq_len=1 can't be split across ranks) + # or KV cache (which stores plain tensors). + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -59,6 +83,15 @@ class Qwen3Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Values are sharding strategies; policies (cpu_offload, mixed_precision) + # live on DistributedConfig.fsdp_plan. All entries marked `keep_full_weight` are bundled + # into a single fully_shard([...]) call at apply time so they share one all-gather. + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 151936 hidden_size: int = 4096 intermediate_size: int = 22016 diff --git a/src/transformers/models/qwen3/modeling_qwen3.py b/src/transformers/models/qwen3/modeling_qwen3.py index 91715a33cf9d..2ba669225879 100644 --- a/src/transformers/models/qwen3/modeling_qwen3.py +++ b/src/transformers/models/qwen3/modeling_qwen3.py @@ -441,8 +441,10 @@ def forward( @auto_docstring class Qwen3ForCausalLM(Qwen3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/qwen3/modular_qwen3.py b/src/transformers/models/qwen3/modular_qwen3.py index 73cde6d89a7a..ea6ddb2f7499 100644 --- a/src/transformers/models/qwen3/modular_qwen3.py +++ b/src/transformers/models/qwen3/modular_qwen3.py @@ -108,6 +108,8 @@ def forward( class Qwen3ForCausalLM(Qwen2ForCausalLM): + _fsdp_plan = {"lm_head": "keep_full_weight"} + def forward( self, **super_kwargs: Unpack[TransformersKwargs], diff --git a/src/transformers/models/qwen3_5/configuration_qwen3_5.py b/src/transformers/models/qwen3_5/configuration_qwen3_5.py index ae9eb8f86c6d..04877a3bbdae 100644 --- a/src/transformers/models/qwen3_5/configuration_qwen3_5.py +++ b/src/transformers/models/qwen3_5/configuration_qwen3_5.py @@ -60,12 +60,10 @@ class Qwen3_5TextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -73,6 +71,12 @@ class Qwen3_5TextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 248320 hidden_size: int = 4096 intermediate_size: int = 12288 diff --git a/src/transformers/models/qwen3_5/modeling_qwen3_5.py b/src/transformers/models/qwen3_5/modeling_qwen3_5.py index 587c7cf362fb..b97168db0a30 100644 --- a/src/transformers/models/qwen3_5/modeling_qwen3_5.py +++ b/src/transformers/models/qwen3_5/modeling_qwen3_5.py @@ -1629,8 +1629,10 @@ def forward( @auto_docstring class Qwen3_5ForCausalLM(Qwen3_5PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} config: Qwen3_5TextConfig _keys_to_ignore_on_load_unexpected = [r"^mtp.*", r"^model.visual.*"] diff --git a/src/transformers/models/qwen3_5/modular_qwen3_5.py b/src/transformers/models/qwen3_5/modular_qwen3_5.py index 91457c58f3ce..82d045c0a28b 100644 --- a/src/transformers/models/qwen3_5/modular_qwen3_5.py +++ b/src/transformers/models/qwen3_5/modular_qwen3_5.py @@ -96,12 +96,10 @@ class Qwen3_5TextConfig(Qwen3NextConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } ignore_keys_at_rope_validation = {"mrope_section", "mrope_interleaved"} diff --git a/src/transformers/models/qwen3_5_moe/configuration_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/configuration_qwen3_5_moe.py index f6f9594e0d73..72c2dacfd2c0 100644 --- a/src/transformers/models/qwen3_5_moe/configuration_qwen3_5_moe.py +++ b/src/transformers/models/qwen3_5_moe/configuration_qwen3_5_moe.py @@ -60,15 +60,11 @@ class Qwen3_5MoeTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_expert.gate_proj": "colwise", "layers.*.mlp.shared_expert.up_proj": "colwise", - "layers.*.mlp.shared_expert.down_proj": "rowwise", + "layers.*.mlp.shared_expert.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -76,6 +72,12 @@ class Qwen3_5MoeTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 248320 hidden_size: int = 2048 num_hidden_layers: int = 40 diff --git a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py index dbf459287d0e..295341bdd7b3 100644 --- a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py +++ b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py @@ -1831,8 +1831,10 @@ def load_balancing_loss_func( @auto_docstring class Qwen3_5MoeForCausalLM(Qwen3_5MoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} config: Qwen3_5MoeTextConfig _keys_to_ignore_on_load_unexpected = [r"^mtp.*", r"^model.visual.*"] @@ -1938,7 +1940,9 @@ class Qwen3_5MoeForConditionalGeneration(Qwen3_5MoePreTrainedModel, GenerationMi # Reference: fix gemma3 grad acc #37208 accepts_loss_kwargs = False config: Qwen3_5MoeConfig - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/qwen3_5_moe/modular_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/modular_qwen3_5_moe.py index 70469ea91bc7..58591c6f2648 100644 --- a/src/transformers/models/qwen3_5_moe/modular_qwen3_5_moe.py +++ b/src/transformers/models/qwen3_5_moe/modular_qwen3_5_moe.py @@ -89,15 +89,11 @@ class Qwen3_5MoeTextConfig(Qwen3NextConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_expert.gate_proj": "colwise", "layers.*.mlp.shared_expert.up_proj": "colwise", - "layers.*.mlp.shared_expert.down_proj": "rowwise", + "layers.*.mlp.shared_expert.down_proj": "rowwise_allreduce", } ignore_keys_at_rope_validation = {"mrope_section", "mrope_interleaved"} @@ -248,7 +244,9 @@ def __init__(self, config): class Qwen3_5MoeForConditionalGeneration(Qwen3VLMoeForConditionalGeneration): - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} def forward(self, **super_kwargs): r""" diff --git a/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py b/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py index 9a7d4b4c8b5b..4395b7c622bd 100644 --- a/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py +++ b/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py @@ -57,15 +57,29 @@ class Qwen3MoeConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts": "moe_experts_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } # Expert-only EP plan: only shards MoE experts, not attention. # Attention is left unsharded — FSDP2 handles attention weight distribution. @@ -74,7 +88,7 @@ class Qwen3MoeConfig(PreTrainedConfig): "layers.*.mlp.gate": "ep_router", "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", "layers.*.mlp.experts.down_proj": "grouped_gemm", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts": "moe_experts_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -82,6 +96,12 @@ class Qwen3MoeConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 151936 hidden_size: int = 2048 intermediate_size: int = 6144 diff --git a/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py b/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py index ddf84fc575b7..012e0cac7243 100644 --- a/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py +++ b/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py @@ -609,8 +609,10 @@ def load_balancing_loss_func( @auto_docstring class Qwen3MoeForCausalLM(Qwen3MoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/qwen3_next/configuration_qwen3_next.py b/src/transformers/models/qwen3_next/configuration_qwen3_next.py index bf26179ff3fd..b499291d70e2 100644 --- a/src/transformers/models/qwen3_next/configuration_qwen3_next.py +++ b/src/transformers/models/qwen3_next/configuration_qwen3_next.py @@ -62,18 +62,14 @@ class Qwen3NextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.shared_expert.gate_proj": "colwise", "layers.*.mlp.shared_expert.up_proj": "colwise", - "layers.*.mlp.shared_expert.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.shared_expert.down_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -81,6 +77,12 @@ class Qwen3NextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 151936 hidden_size: int = 2048 intermediate_size: int = 5632 diff --git a/src/transformers/models/qwen3_next/modeling_qwen3_next.py b/src/transformers/models/qwen3_next/modeling_qwen3_next.py index df167395975b..faca94ba5c43 100644 --- a/src/transformers/models/qwen3_next/modeling_qwen3_next.py +++ b/src/transformers/models/qwen3_next/modeling_qwen3_next.py @@ -1088,8 +1088,10 @@ def load_balancing_loss_func( @auto_docstring class Qwen3NextForCausalLM(Qwen3NextPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py b/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py index 4e9bdd35f21f..3c65badeafb4 100644 --- a/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py +++ b/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py @@ -138,18 +138,23 @@ class Qwen3OmniMoeTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + ignore_keys_at_rope_validation = {"mrope_section", "interleaved", "mrope_interleaved"} vocab_size: int = 3584 @@ -263,17 +268,41 @@ class Qwen3OmniMoeTalkerCodePredictorConfig(PreTrainedConfig): model_type = "qwen3_omni_moe_talker_code_predictor" keys_to_ignore_at_inference = ["past_key_values"] - # Default tensor parallel plan for base model `Qwen3OmniMoeTalkerCodePredictor` + # TP plan (for inference/generation). + # All activations are plain tensors — compatible with KV cache and autoregressive + # decode (seq_len=1). Each rank holds a full copy of activations between layers. base_model_tp_plan = { "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + + # TP + Sequence Parallelism plan (for training). + # Activations between layers are sharded on the sequence dimension (Shard(1)), + # reducing per-rank activation memory by tp_size. In exchange, extra collectives + # (all-gather before attention/MLP, reduce-scatter after) are needed. + # Not compatible with autoregressive decode (because seq_len=1 can't be split across ranks) + # or KV cache (which stores plain tensors). + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -281,6 +310,15 @@ class Qwen3OmniMoeTalkerCodePredictorConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Values are sharding strategies; policies (cpu_offload, mixed_precision) + # live on DistributedConfig.fsdp_plan. All entries marked `keep_full_weight` are bundled + # into a single fully_shard([...]) call at apply time so they share one all-gather. + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 2048 hidden_size: int = 1024 intermediate_size: int = 3072 @@ -357,21 +395,35 @@ class Qwen3OmniMoeTalkerTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_ep_plan = { "layers.*.mlp.gate": "ep_router", "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", "layers.*.mlp.experts.down_proj": "grouped_gemm", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts": "moe_experts_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -379,6 +431,12 @@ class Qwen3OmniMoeTalkerTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 3072 hidden_size: int = 1024 intermediate_size: int = 2048 diff --git a/src/transformers/models/qwen3_omni_moe/modeling_qwen3_omni_moe.py b/src/transformers/models/qwen3_omni_moe/modeling_qwen3_omni_moe.py index 05f310a75b11..83875dc88b5d 100644 --- a/src/transformers/models/qwen3_omni_moe/modeling_qwen3_omni_moe.py +++ b/src/transformers/models/qwen3_omni_moe/modeling_qwen3_omni_moe.py @@ -2637,8 +2637,10 @@ def get_input_embeddings(self): @auto_docstring class Qwen3OmniMoeTalkerCodePredictorModelForConditionalGeneration(Qwen3OmniMoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} config_class = Qwen3OmniMoeTalkerCodePredictorConfig base_model_prefix = "talker.code_predictor" _can_record_outputs = { @@ -3021,8 +3023,10 @@ def get_input_embeddings(self): @auto_docstring class Qwen3OmniMoeTalkerForConditionalGeneration(Qwen3OmniMoeThinkerTextPreTrainedModel, GenerationMixin): _tied_weights_keys = {"codec_head": "model.codec_embedding.weight"} - _tp_plan = {"codec_head": "colwise_gather_output"} + _tp_plan = {"codec_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"codec_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} config_class = Qwen3OmniMoeTalkerConfig base_model_prefix = "talker" _no_split_modules = ["Qwen3OmniMoeTalkerCodePredictorModelForConditionalGeneration"] diff --git a/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py b/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py index 9f8351d3f8c4..dd98da222cf3 100644 --- a/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py +++ b/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py @@ -280,18 +280,23 @@ class Qwen3OmniMoeTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + ignore_keys_at_rope_validation = {"mrope_section", "interleaved", "mrope_interleaved"} vocab_size: int = 3584 @@ -394,7 +399,7 @@ class Qwen3OmniMoeTalkerTextConfig(Qwen3MoeConfig): "layers.*.mlp.gate": "ep_router", "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", "layers.*.mlp.experts.down_proj": "grouped_gemm", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts": "moe_experts_allreduce", } vocab_size: int = 3072 @@ -1659,7 +1664,7 @@ def get_input_embeddings(self): class Qwen3OmniMoeTalkerForConditionalGeneration(Qwen3MoeForCausalLM): _tied_weights_keys = {"codec_head": "model.codec_embedding.weight"} - _tp_plan = {"codec_head": "colwise_gather_output"} + _tp_plan = {"codec_head": "colwise_allgather"} _pp_plan = {"codec_head": (["hidden_states"], ["logits"])} config_class = Qwen3OmniMoeTalkerConfig base_model_prefix = "talker" diff --git a/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py b/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py index 6825f5dda929..f49d389a593f 100644 --- a/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py +++ b/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py @@ -59,16 +59,34 @@ class Qwen3VLMoeTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts": "moe_experts_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_ep_plan = { "layers.*.mlp.gate": "ep_router", "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", "layers.*.mlp.experts.down_proj": "grouped_gemm", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts": "moe_experts_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -76,6 +94,12 @@ class Qwen3VLMoeTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 151936 hidden_size: int = 2048 diff --git a/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py b/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py index 11534b395773..555e3d63ccb3 100644 --- a/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py +++ b/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py @@ -86,16 +86,16 @@ class Qwen3VLMoeTextConfig(Qwen3MoeConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_ep_plan = { "layers.*.mlp.gate": "ep_router", "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", "layers.*.mlp.experts.down_proj": "grouped_gemm", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts": "moe_experts_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/seamless_m4t/modeling_seamless_m4t.py b/src/transformers/models/seamless_m4t/modeling_seamless_m4t.py index e3f2d7a6ed26..10b31ee9c3c9 100755 --- a/src/transformers/models/seamless_m4t/modeling_seamless_m4t.py +++ b/src/transformers/models/seamless_m4t/modeling_seamless_m4t.py @@ -24,9 +24,9 @@ from ... import initialization as init from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache +from ...distributed.fsdp import is_fsdp_managed_module from ...generation import GenerationMixin from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...masking_utils import create_bidirectional_mask, create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( diff --git a/src/transformers/models/seamless_m4t_v2/modeling_seamless_m4t_v2.py b/src/transformers/models/seamless_m4t_v2/modeling_seamless_m4t_v2.py index a2b292555fc8..48e85b6499c5 100644 --- a/src/transformers/models/seamless_m4t_v2/modeling_seamless_m4t_v2.py +++ b/src/transformers/models/seamless_m4t_v2/modeling_seamless_m4t_v2.py @@ -24,9 +24,9 @@ from ... import initialization as init from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache +from ...distributed.fsdp import is_fsdp_managed_module from ...generation import GenerationMixin from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...masking_utils import create_bidirectional_mask, create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( diff --git a/src/transformers/models/seed_oss/configuration_seed_oss.py b/src/transformers/models/seed_oss/configuration_seed_oss.py index b1221fcf53ce..363b393dbc4a 100644 --- a/src/transformers/models/seed_oss/configuration_seed_oss.py +++ b/src/transformers/models/seed_oss/configuration_seed_oss.py @@ -48,10 +48,25 @@ class SeedOssConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -59,6 +74,12 @@ class SeedOssConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 155136 hidden_size: int = 4096 intermediate_size: int = 27648 diff --git a/src/transformers/models/seed_oss/modeling_seed_oss.py b/src/transformers/models/seed_oss/modeling_seed_oss.py index 1ebc8f10a272..841765a7c928 100644 --- a/src/transformers/models/seed_oss/modeling_seed_oss.py +++ b/src/transformers/models/seed_oss/modeling_seed_oss.py @@ -429,8 +429,10 @@ def forward( @auto_docstring class SeedOssForCausalLM(SeedOssPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/sew/modeling_sew.py b/src/transformers/models/sew/modeling_sew.py index ea499e63289a..30e527b4c985 100644 --- a/src/transformers/models/sew/modeling_sew.py +++ b/src/transformers/models/sew/modeling_sew.py @@ -28,8 +28,8 @@ from ... import initialization as init from ...activations import ACT2FN +from ...distributed.fsdp import is_fsdp_managed_module from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, CausalLMOutput, SequenceClassifierOutput diff --git a/src/transformers/models/sew/modular_sew.py b/src/transformers/models/sew/modular_sew.py index 312419793a34..f3db3c1bcaaa 100644 --- a/src/transformers/models/sew/modular_sew.py +++ b/src/transformers/models/sew/modular_sew.py @@ -20,8 +20,8 @@ from ... import initialization as init from ...activations import ACT2FN +from ...distributed.fsdp import is_fsdp_managed_module from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...modeling_outputs import BaseModelOutput from ...modeling_utils import PreTrainedModel from ...utils import auto_docstring diff --git a/src/transformers/models/smollm3/configuration_smollm3.py b/src/transformers/models/smollm3/configuration_smollm3.py index f48c979a1dd3..9a14e834db9e 100644 --- a/src/transformers/models/smollm3/configuration_smollm3.py +++ b/src/transformers/models/smollm3/configuration_smollm3.py @@ -58,10 +58,25 @@ class SmolLM3Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -69,6 +84,12 @@ class SmolLM3Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 128256 hidden_size: int = 2048 intermediate_size: int = 11008 diff --git a/src/transformers/models/smollm3/modeling_smollm3.py b/src/transformers/models/smollm3/modeling_smollm3.py index 8d911e414b0f..d82882382bb9 100644 --- a/src/transformers/models/smollm3/modeling_smollm3.py +++ b/src/transformers/models/smollm3/modeling_smollm3.py @@ -445,8 +445,10 @@ def forward( @auto_docstring class SmolLM3ForCausalLM(SmolLM3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/smollm3/modular_smollm3.py b/src/transformers/models/smollm3/modular_smollm3.py index f75017ad2645..69c9d4b51444 100644 --- a/src/transformers/models/smollm3/modular_smollm3.py +++ b/src/transformers/models/smollm3/modular_smollm3.py @@ -74,10 +74,25 @@ class SmolLM3Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -85,6 +100,12 @@ class SmolLM3Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 128256 hidden_size: int = 2048 intermediate_size: int = 11008 diff --git a/src/transformers/models/solar_open/configuration_solar_open.py b/src/transformers/models/solar_open/configuration_solar_open.py index ac0016aa7791..5d9aec3ac7a0 100644 --- a/src/transformers/models/solar_open/configuration_solar_open.py +++ b/src/transformers/models/solar_open/configuration_solar_open.py @@ -41,16 +41,21 @@ class SolarOpenConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + attribute_map = { "num_local_experts": "n_routed_experts", } @@ -81,6 +86,19 @@ class SolarOpenConfig(PreTrainedConfig): eos_token_id: int | list[int] | None = None pad_token_id: int | None = None default_theta = 1_000_000.0 + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts": "moe_experts_allreduce", + "norm": "activation", + } head_dim: int = 128 def __post_init__(self, **kwargs): diff --git a/src/transformers/models/solar_open/modeling_solar_open.py b/src/transformers/models/solar_open/modeling_solar_open.py index 0eb50021ecd6..53abac7c7d07 100644 --- a/src/transformers/models/solar_open/modeling_solar_open.py +++ b/src/transformers/models/solar_open/modeling_solar_open.py @@ -552,8 +552,10 @@ def forward( @auto_docstring class SolarOpenForCausalLM(SolarOpenPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/solar_open/modular_solar_open.py b/src/transformers/models/solar_open/modular_solar_open.py index 90d4f0c389c0..48bba9439e08 100644 --- a/src/transformers/models/solar_open/modular_solar_open.py +++ b/src/transformers/models/solar_open/modular_solar_open.py @@ -47,10 +47,27 @@ class SolarOpenConfig(Glm4MoeConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts": "moe_experts_allreduce", + "norm": "activation", + } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", } vocab_size: int = 196608 diff --git a/src/transformers/models/speecht5/modeling_speecht5.py b/src/transformers/models/speecht5/modeling_speecht5.py index f75fa9dcdcd9..e5d7d706a9ab 100644 --- a/src/transformers/models/speecht5/modeling_speecht5.py +++ b/src/transformers/models/speecht5/modeling_speecht5.py @@ -23,9 +23,9 @@ from ... import initialization as init from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache +from ...distributed.fsdp import is_fsdp_managed_module from ...generation import GenerationMixin from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...masking_utils import create_bidirectional_mask, create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( diff --git a/src/transformers/models/starcoder2/configuration_starcoder2.py b/src/transformers/models/starcoder2/configuration_starcoder2.py index 59efa94fc5f4..88137e8f362f 100644 --- a/src/transformers/models/starcoder2/configuration_starcoder2.py +++ b/src/transformers/models/starcoder2/configuration_starcoder2.py @@ -48,9 +48,23 @@ class Starcoder2Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.c_fc": "colwise", - "layers.*.mlp.c_proj": "rowwise", + "layers.*.mlp.c_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.c_fc": "colwise", + "layers.*.mlp.c_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -58,6 +72,12 @@ class Starcoder2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 49152 hidden_size: int = 3072 intermediate_size: int = 12288 diff --git a/src/transformers/models/starcoder2/modeling_starcoder2.py b/src/transformers/models/starcoder2/modeling_starcoder2.py index 8b89a1d1745c..baabf385da59 100644 --- a/src/transformers/models/starcoder2/modeling_starcoder2.py +++ b/src/transformers/models/starcoder2/modeling_starcoder2.py @@ -409,8 +409,10 @@ def forward( @auto_docstring class Starcoder2ForCausalLM(Starcoder2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/t5gemma/configuration_t5gemma.py b/src/transformers/models/t5gemma/configuration_t5gemma.py index 9de40c832259..8c313bae875e 100644 --- a/src/transformers/models/t5gemma/configuration_t5gemma.py +++ b/src/transformers/models/t5gemma/configuration_t5gemma.py @@ -54,10 +54,25 @@ class T5GemmaModuleConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -65,6 +80,12 @@ class T5GemmaModuleConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 256000 hidden_size: int = 2304 intermediate_size: int = 9216 diff --git a/src/transformers/models/t5gemma/modeling_t5gemma.py b/src/transformers/models/t5gemma/modeling_t5gemma.py index 1f41875c5def..f2ff830bf09c 100644 --- a/src/transformers/models/t5gemma/modeling_t5gemma.py +++ b/src/transformers/models/t5gemma/modeling_t5gemma.py @@ -945,7 +945,7 @@ def forward( class T5GemmaForConditionalGeneration(T5GemmaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.out_proj.weight": "model.decoder.embed_tokens.weight"} - _tp_plan = {"lm_head.out_proj": "colwise_gather_output"} + _tp_plan = {"lm_head.out_proj": "colwise_allgather"} _pp_plan = {"lm_head.out_proj": (["hidden_states"], ["logits"])} def __init__(self, config: T5GemmaConfig): diff --git a/src/transformers/models/t5gemma/modular_t5gemma.py b/src/transformers/models/t5gemma/modular_t5gemma.py index 1c8846ad74b9..f35a06beab33 100644 --- a/src/transformers/models/t5gemma/modular_t5gemma.py +++ b/src/transformers/models/t5gemma/modular_t5gemma.py @@ -84,6 +84,12 @@ class T5GemmaModuleConfig(Gemma2Config): >>> configuration = model.config ```""" + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + is_decoder: bool = False use_bidirectional_attention = AttributeError() @@ -784,7 +790,7 @@ def forward( class T5GemmaForConditionalGeneration(T5GemmaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.out_proj.weight": "model.decoder.embed_tokens.weight"} - _tp_plan = {"lm_head.out_proj": "colwise_gather_output"} + _tp_plan = {"lm_head.out_proj": "colwise_allgather"} _pp_plan = {"lm_head.out_proj": (["hidden_states"], ["logits"])} def __init__(self, config: T5GemmaConfig): diff --git a/src/transformers/models/t5gemma2/configuration_t5gemma2.py b/src/transformers/models/t5gemma2/configuration_t5gemma2.py index d9a9a3f5769f..b3c2ada26f35 100644 --- a/src/transformers/models/t5gemma2/configuration_t5gemma2.py +++ b/src/transformers/models/t5gemma2/configuration_t5gemma2.py @@ -48,12 +48,27 @@ class T5Gemma2TextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -61,6 +76,12 @@ class T5Gemma2TextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 262_208 hidden_size: int = 2304 intermediate_size: int = 9216 @@ -85,6 +106,7 @@ class T5Gemma2TextConfig(PreTrainedConfig): layer_types: list[str] | None = None final_logit_softcapping: float | None = None attn_logit_softcapping: float | None = None + default_theta = {"global": 1_000_000.0, "local": 10_000.0} def __post_init__(self, **kwargs): @@ -223,12 +245,27 @@ class T5Gemma2DecoderConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -236,6 +273,12 @@ class T5Gemma2DecoderConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 262_208 hidden_size: int = 2304 intermediate_size: int = 9216 @@ -260,6 +303,7 @@ class T5Gemma2DecoderConfig(PreTrainedConfig): layer_types: list[str] | None = None final_logit_softcapping: float | None = None attn_logit_softcapping: float | None = None + default_theta = {"global": 1_000_000.0, "local": 10_000.0} def __post_init__(self, **kwargs): diff --git a/src/transformers/models/t5gemma2/modeling_t5gemma2.py b/src/transformers/models/t5gemma2/modeling_t5gemma2.py index 15de07406bfb..1670929c8785 100644 --- a/src/transformers/models/t5gemma2/modeling_t5gemma2.py +++ b/src/transformers/models/t5gemma2/modeling_t5gemma2.py @@ -1183,7 +1183,7 @@ class T5Gemma2ForConditionalGeneration(T5Gemma2PreTrainedModel, GenerationMixin) _tied_weights_keys = { "lm_head.out_proj.weight": "model.encoder.text_model.embed_tokens.weight", } - _tp_plan = {"lm_head.out_proj": "colwise_gather_output"} + _tp_plan = {"lm_head.out_proj": "colwise_allgather"} _pp_plan = {"lm_head.out_proj": (["hidden_states"], ["logits"])} def __init__(self, config: T5Gemma2Config): diff --git a/src/transformers/models/t5gemma2/modular_t5gemma2.py b/src/transformers/models/t5gemma2/modular_t5gemma2.py index a1b80a81ef80..89b926f9d4a7 100644 --- a/src/transformers/models/t5gemma2/modular_t5gemma2.py +++ b/src/transformers/models/t5gemma2/modular_t5gemma2.py @@ -86,6 +86,13 @@ class T5Gemma2TextConfig(Gemma3TextConfig, PreTrainedConfig): """ model_type = "t5gemma2_text" + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + use_bidirectional_attention = AttributeError() def __post_init__(self, **kwargs): @@ -973,7 +980,7 @@ class T5Gemma2ForConditionalGeneration(T5Gemma2PreTrainedModel, GenerationMixin) _tied_weights_keys = { "lm_head.out_proj.weight": "model.encoder.text_model.embed_tokens.weight", } - _tp_plan = {"lm_head.out_proj": "colwise_gather_output"} + _tp_plan = {"lm_head.out_proj": "colwise_allgather"} _pp_plan = {"lm_head.out_proj": (["hidden_states"], ["logits"])} def __init__(self, config: T5Gemma2Config): diff --git a/src/transformers/models/unispeech/modeling_unispeech.py b/src/transformers/models/unispeech/modeling_unispeech.py index 03103760140c..bf4fc0d108f7 100755 --- a/src/transformers/models/unispeech/modeling_unispeech.py +++ b/src/transformers/models/unispeech/modeling_unispeech.py @@ -29,8 +29,8 @@ from ... import initialization as init from ...activations import ACT2FN +from ...distributed.fsdp import is_fsdp_managed_module from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...masking_utils import create_bidirectional_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer diff --git a/src/transformers/models/unispeech_sat/modeling_unispeech_sat.py b/src/transformers/models/unispeech_sat/modeling_unispeech_sat.py index 565852237d06..6cda17570834 100755 --- a/src/transformers/models/unispeech_sat/modeling_unispeech_sat.py +++ b/src/transformers/models/unispeech_sat/modeling_unispeech_sat.py @@ -30,8 +30,8 @@ from ... import initialization as init from ...activations import ACT2FN +from ...distributed.fsdp import is_fsdp_managed_module from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...masking_utils import create_bidirectional_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer diff --git a/src/transformers/models/vaultgemma/configuration_vaultgemma.py b/src/transformers/models/vaultgemma/configuration_vaultgemma.py index a60b7e8edc0c..7ff8c838d362 100644 --- a/src/transformers/models/vaultgemma/configuration_vaultgemma.py +++ b/src/transformers/models/vaultgemma/configuration_vaultgemma.py @@ -53,10 +53,25 @@ class VaultGemmaConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -64,6 +79,12 @@ class VaultGemmaConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 256000 hidden_size: int = 2304 intermediate_size: int = 9216 diff --git a/src/transformers/models/vaultgemma/modeling_vaultgemma.py b/src/transformers/models/vaultgemma/modeling_vaultgemma.py index f0a2e48d20b8..1ddba4af2210 100644 --- a/src/transformers/models/vaultgemma/modeling_vaultgemma.py +++ b/src/transformers/models/vaultgemma/modeling_vaultgemma.py @@ -467,8 +467,10 @@ def forward( @auto_docstring class VaultGemmaForCausalLM(VaultGemmaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/vaultgemma/modular_vaultgemma.py b/src/transformers/models/vaultgemma/modular_vaultgemma.py index 9a4ce67c9a55..65f8b4ab3bf3 100644 --- a/src/transformers/models/vaultgemma/modular_vaultgemma.py +++ b/src/transformers/models/vaultgemma/modular_vaultgemma.py @@ -43,6 +43,12 @@ class VaultGemmaConfig(Gemma2Config): >>> configuration = model.config ```""" + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + use_bidirectional_attention = AttributeError() diff --git a/src/transformers/models/vibevoice_asr/modeling_vibevoice_asr.py b/src/transformers/models/vibevoice_asr/modeling_vibevoice_asr.py index 4eb523fd3218..4e60c72ccd40 100644 --- a/src/transformers/models/vibevoice_asr/modeling_vibevoice_asr.py +++ b/src/transformers/models/vibevoice_asr/modeling_vibevoice_asr.py @@ -264,6 +264,7 @@ class VibeVoiceAsrForConditionalGeneration(VibeVoiceAsrPreTrainedModel, Generati _keep_in_fp32_modules_strict = None _supports_attention_backend = True _tp_plan = None + _sp_plan = None _pp_plan = None def __init__(self, config: VibeVoiceAsrConfig): diff --git a/src/transformers/models/vits/modeling_vits.py b/src/transformers/models/vits/modeling_vits.py index 8ce1411bed5e..3ab719c0921d 100644 --- a/src/transformers/models/vits/modeling_vits.py +++ b/src/transformers/models/vits/modeling_vits.py @@ -23,8 +23,8 @@ from ... import initialization as init from ...activations import ACT2FN +from ...distributed.fsdp import is_fsdp_managed_module from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...masking_utils import create_bidirectional_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, ModelOutput diff --git a/src/transformers/models/voxtral_realtime/configuration_voxtral_realtime.py b/src/transformers/models/voxtral_realtime/configuration_voxtral_realtime.py index b0227b418771..568c6f8748b9 100644 --- a/src/transformers/models/voxtral_realtime/configuration_voxtral_realtime.py +++ b/src/transformers/models/voxtral_realtime/configuration_voxtral_realtime.py @@ -30,10 +30,10 @@ class VoxtralRealtimeTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -41,6 +41,12 @@ class VoxtralRealtimeTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 32000 hidden_size: int = 4096 intermediate_size: int = 14336 diff --git a/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py b/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py index eb0ab6a2cef0..d5d2a92f18f7 100644 --- a/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py +++ b/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py @@ -830,8 +830,10 @@ def forward( @auto_docstring class VoxtralRealtimeTextForCausalLM(VoxtralRealtimeTextPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/wav2vec2/modeling_wav2vec2.py b/src/transformers/models/wav2vec2/modeling_wav2vec2.py index 274a03365710..6594e813a5a9 100755 --- a/src/transformers/models/wav2vec2/modeling_wav2vec2.py +++ b/src/transformers/models/wav2vec2/modeling_wav2vec2.py @@ -26,8 +26,8 @@ from ... import initialization as init from ...activations import ACT2FN +from ...distributed.fsdp import is_fsdp_managed_module from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...masking_utils import create_bidirectional_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer diff --git a/src/transformers/models/wav2vec2_bert/modeling_wav2vec2_bert.py b/src/transformers/models/wav2vec2_bert/modeling_wav2vec2_bert.py index 6023d856798b..1faa3b667b73 100644 --- a/src/transformers/models/wav2vec2_bert/modeling_wav2vec2_bert.py +++ b/src/transformers/models/wav2vec2_bert/modeling_wav2vec2_bert.py @@ -14,8 +14,8 @@ from ... import initialization as init from ...activations import ACT2FN +from ...distributed.fsdp import is_fsdp_managed_module from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...masking_utils import create_bidirectional_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( diff --git a/src/transformers/models/wav2vec2_bert/modular_wav2vec2_bert.py b/src/transformers/models/wav2vec2_bert/modular_wav2vec2_bert.py index 710e7a64cea2..a616f62059ef 100644 --- a/src/transformers/models/wav2vec2_bert/modular_wav2vec2_bert.py +++ b/src/transformers/models/wav2vec2_bert/modular_wav2vec2_bert.py @@ -6,8 +6,8 @@ from ... import initialization as init from ...activations import ACT2FN +from ...distributed.fsdp import is_fsdp_managed_module from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...masking_utils import create_bidirectional_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( diff --git a/src/transformers/models/wav2vec2_conformer/modeling_wav2vec2_conformer.py b/src/transformers/models/wav2vec2_conformer/modeling_wav2vec2_conformer.py index 9f35e5db42ed..aa454e20da48 100644 --- a/src/transformers/models/wav2vec2_conformer/modeling_wav2vec2_conformer.py +++ b/src/transformers/models/wav2vec2_conformer/modeling_wav2vec2_conformer.py @@ -15,8 +15,8 @@ from ... import initialization as init from ...activations import ACT2FN +from ...distributed.fsdp import is_fsdp_managed_module from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( BaseModelOutput, diff --git a/src/transformers/models/wav2vec2_conformer/modular_wav2vec2_conformer.py b/src/transformers/models/wav2vec2_conformer/modular_wav2vec2_conformer.py index f02ce539d228..ce34abedbdb4 100644 --- a/src/transformers/models/wav2vec2_conformer/modular_wav2vec2_conformer.py +++ b/src/transformers/models/wav2vec2_conformer/modular_wav2vec2_conformer.py @@ -6,8 +6,8 @@ from ... import initialization as init from ...activations import ACT2FN +from ...distributed.fsdp import is_fsdp_managed_module from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, Wav2Vec2BaseModelOutput from ...modeling_utils import PreTrainedModel diff --git a/src/transformers/models/wavlm/modeling_wavlm.py b/src/transformers/models/wavlm/modeling_wavlm.py index 18440ebf7d25..024c889ed110 100755 --- a/src/transformers/models/wavlm/modeling_wavlm.py +++ b/src/transformers/models/wavlm/modeling_wavlm.py @@ -15,8 +15,8 @@ from ... import initialization as init from ...activations import ACT2FN +from ...distributed.fsdp import is_fsdp_managed_module from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( BaseModelOutput, diff --git a/src/transformers/models/wavlm/modular_wavlm.py b/src/transformers/models/wavlm/modular_wavlm.py index b3329e64913d..77a58d986ab3 100644 --- a/src/transformers/models/wavlm/modular_wavlm.py +++ b/src/transformers/models/wavlm/modular_wavlm.py @@ -5,8 +5,8 @@ import torch.nn.functional as F from ... import initialization as init +from ...distributed.fsdp import is_fsdp_managed_module from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, Wav2Vec2BaseModelOutput from ...modeling_utils import PreTrainedModel diff --git a/src/transformers/models/youtu/configuration_youtu.py b/src/transformers/models/youtu/configuration_youtu.py index 6d9f2cef1f96..6210f3c5f42b 100644 --- a/src/transformers/models/youtu/configuration_youtu.py +++ b/src/transformers/models/youtu/configuration_youtu.py @@ -53,13 +53,19 @@ class YoutuConfig(PreTrainedConfig): base_model_tp_plan = { "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } attribute_map = {} vocab_size: int = 128256 @@ -86,6 +92,7 @@ class YoutuConfig(PreTrainedConfig): rope_interleave: bool | None = True attention_bias: bool = False attention_dropout: float | int | None = 0.0 + base_model_sp_plan = None embedding_initializer_range: float | None = None def __post_init__(self, **kwargs): @@ -103,5 +110,20 @@ def __post_init__(self, **kwargs): self.head_dim = self.qk_rope_head_dim super().__post_init__(**kwargs) + def convert_rope_params_to_dict(self, **kwargs): + rope_scaling = kwargs.pop("rope_scaling", None) + self.rope_parameters = rope_scaling or self.rope_parameters + self.rope_parameters = self.rope_parameters if self.rope_parameters is not None else {} + + # Standardize and validate the correctness of rotary position embeddings parameters + self.rope_parameters.setdefault("rope_theta", kwargs.pop("rope_theta", self.default_theta)) + self.standardize_rope_params() + + # Convert to float because RoPE fn expect a float. Models on the hub were saved as int + for key in ["beta_fast", "beta_slow", "factor"]: + if key in self.rope_parameters: + self.rope_parameters[key] = float(self.rope_parameters[key]) + return kwargs + __all__ = ["YoutuConfig"] diff --git a/src/transformers/models/youtu/modeling_youtu.py b/src/transformers/models/youtu/modeling_youtu.py index d40bef358da6..8751f452fd28 100644 --- a/src/transformers/models/youtu/modeling_youtu.py +++ b/src/transformers/models/youtu/modeling_youtu.py @@ -533,8 +533,10 @@ def forward( @auto_docstring class YoutuForCausalLM(YoutuPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/youtu/modular_youtu.py b/src/transformers/models/youtu/modular_youtu.py index b2de3a2df0a5..f3218061670c 100644 --- a/src/transformers/models/youtu/modular_youtu.py +++ b/src/transformers/models/youtu/modular_youtu.py @@ -62,8 +62,9 @@ class YoutuConfig(DeepseekV3Config): base_model_tp_plan = { "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } + base_model_sp_plan = None attribute_map = {} vocab_size: int = 128256 diff --git a/src/transformers/testing_utils.py b/src/transformers/testing_utils.py index 87236fb819ab..1a3c40380c48 100644 --- a/src/transformers/testing_utils.py +++ b/src/transformers/testing_utils.py @@ -262,7 +262,10 @@ if is_torch_available(): import torch + import torch.distributed as dist + import torch.multiprocessing as mp from safetensors.torch import load_file + from torch.distributed.device_mesh import init_device_mesh from .modeling_utils import FLASH_ATTN_KERNEL_FALLBACK, PreTrainedModel @@ -316,6 +319,7 @@ def parse_int_from_env(key, default=None): _run_agent_tests = parse_flag_from_env("RUN_AGENT_TESTS", default=False) _run_training_tests = parse_flag_from_env("RUN_TRAINING_TESTS", default=True) _run_tensor_parallel_tests = parse_flag_from_env("RUN_TENSOR_PARALLEL_TESTS", default=True) +_run_fsdp_tests = parse_flag_from_env("RUN_FSDP_TESTS", default=True) def is_staging_test(test_case): @@ -382,6 +386,22 @@ def is_training_test(test_case): return pytest.mark.is_training_test()(test_case) +def is_training_distributed_test(test_case): + """ + Decorator marking a test as a training distributed test. If RUN_TRAINING_DISTRIBUTED_TESTS is set to a falsy value, those tests will be + skipped. + """ + if not _run_training_tests: + return unittest.skip(reason="test is training distributed test")(test_case) + else: + try: + import pytest # We don't need a hard dependency on pytest in the main library + except ImportError: + return test_case + else: + return pytest.mark.is_training_distributed_test()(test_case) + + def is_tensor_parallel_test(test_case): """ Decorator marking a test as a tensor parallel test. If RUN_TENSOR_PARALLEL_TESTS is set to a falsy value, those @@ -398,6 +418,22 @@ def is_tensor_parallel_test(test_case): return pytest.mark.is_tensor_parallel_test()(test_case) +def is_fsdp_test(test_case): + """ + Decorator marking a test as an FSDP test. If RUN_FSDP_TESTS is set to a falsy value, those tests will be + skipped. + """ + if not _run_fsdp_tests: + return unittest.skip(reason="test is fsdp test")(test_case) + else: + try: + import pytest # We don't need a hard dependency on pytest in the main library + except ImportError: + return test_case + else: + return pytest.mark.is_fsdp_test()(test_case) + + def slow(test_case): """ Decorator marking a test as slow. @@ -4193,6 +4229,45 @@ def read_json_file(file): # ============================================================================= +def global_wrapper(rank, func, fsdp_size, tp_size, port, func_args, func_kwargs): + def setup_dist_env(rank, world_size, port): + os.environ["WORLD_SIZE"] = str(world_size) + os.environ["RANK"] = str(rank) + os.environ["LOCAL_RANK"] = str(rank) + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = str(port) + + world_size = fsdp_size * tp_size + setup_dist_env(rank, world_size, port) + + dist.init_process_group(backend="gloo", rank=rank, world_size=world_size) + + # NOTE(3outeille): if want to handle DataParallel, create dp_replicate dims (do not mixed with dp_shard which is for FSDP) + # NOTE(3outeille): if other parallelism is added, order matters, it should be ["pp", "ddp", "fsdp", "cp", "tp"] + # TODO(3outeille): figure out EP + # from less costly to most costly (internode to intranode) + dims, names = [fsdp_size, tp_size], ["fsdp", "tp"] + mesh = init_device_mesh("cpu", dims, mesh_dim_names=names) + + func(mesh, *func_args, **func_kwargs) + + dist.barrier() + dist.destroy_process_group() + + +def init_distributed(fsdp_size: int = 1, tp_size: int = 1): + def _init_distributed(func): + def wrapper(*args, **kwargs): + world_size = fsdp_size * tp_size + port = get_torch_dist_unique_port() + spawn_args = (func, fsdp_size, tp_size, port, args, kwargs) + mp.spawn(global_wrapper, args=spawn_args, nprocs=world_size) + + return wrapper + + return _init_distributed + + # ANSI color codes for terminal output class Colors: """ANSI color codes for terminal output formatting.""" @@ -4232,8 +4307,9 @@ class ColoredFormatter(logging.Formatter): # Loggers that should be dimmed (less important/verbose) DIMMED_LOGGERS = {"httpx", "httpcore", "urllib3", "requests"} - def __init__(self, fmt: str | None = None, datefmt: str | None = None): + def __init__(self, fmt: str | None = None, datefmt: str | None = None, rank_prefix: str = ""): super().__init__(fmt, datefmt) + self.rank_prefix = rank_prefix def format(self, record: logging.LogRecord) -> str: # Check if this logger should be dimmed @@ -4243,7 +4319,7 @@ def format(self, record: logging.LogRecord) -> str: # Dim the entire log line for httpx and similar timestamp = self.formatTime(record, self.datefmt) message = record.getMessage() - return f"{Colors.DIM}{timestamp} - {record.name} - {record.levelname:8} - {message}{Colors.RESET}" + return f"{Colors.DIM}{timestamp} - {record.name} - {record.levelname:8} - {self.rank_prefix}{message}{Colors.RESET}" # Get color for this level color = self.LEVEL_COLORS.get(record.levelno, Colors.RESET) @@ -4261,7 +4337,7 @@ def format(self, record: logging.LogRecord) -> str: # Get message message = record.getMessage() - return f"{colored_time} - {colored_name} - {colored_levelname} - {message}" + return f"{colored_time} - {colored_name} - {colored_levelname} - {self.rank_prefix}{message}" _warn_once_logged: set[str] = set() @@ -4272,26 +4348,34 @@ def init_test_logger() -> logging.Logger: Uses a named logger instead of root logger to avoid conflicts with pytest-xdist parallel execution. Uses stderr instead of stdout to avoid deadlocks with pytest-xdist output capture. + Automatically includes rank in log format when distributed is initialized. """ logger = logging.getLogger("transformers.training_test") logger.setLevel(logging.INFO) - # Only add handler if not already present (avoid duplicate handlers on repeated calls) - if not logger.handlers: - # Use stderr instead of stdout - pytest-xdist captures stdout which can cause deadlocks - ch = logging.StreamHandler(sys.stderr) - ch.setLevel(logging.INFO) + # Clear existing handlers to update format (e.g., when dist becomes initialized) + logger.handlers.clear() - # Use colored formatter if terminal supports it, plain otherwise - if sys.stderr.isatty(): - formatter = ColoredFormatter(datefmt="%Y-%m-%d %H:%M:%S") - else: - formatter = logging.Formatter( - "%(asctime)s - %(name)s - %(levelname)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S" - ) + # Use stderr instead of stdout - pytest-xdist captures stdout which can cause deadlocks + ch = logging.StreamHandler(sys.stderr) + ch.setLevel(logging.INFO) + + # Build format string - include rank if distributed is initialized + rank_prefix = "" + if is_torch_available() and dist.is_initialized(): + rank = dist.get_rank() + rank_prefix = f"[rank{rank}] " + + # Use colored formatter if terminal supports it, plain otherwise + if sys.stderr.isatty(): + formatter = ColoredFormatter(datefmt="%Y-%m-%d %H:%M:%S", rank_prefix=rank_prefix) + else: + formatter = logging.Formatter( + f"%(asctime)s - %(name)s - %(levelname)s - {rank_prefix}%(message)s", datefmt="%Y-%m-%d %H:%M:%S" + ) - ch.setFormatter(formatter) - logger.addHandler(ch) + ch.setFormatter(formatter) + logger.addHandler(ch) logger.propagate = False # Don't propagate to root logger to avoid duplicate output return logger diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index 4ba6e516932b..a3973a3ada4c 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -55,6 +55,7 @@ from .configuration_utils import PreTrainedConfig from .data.data_collator import DataCollator, DataCollatorWithPadding, default_data_collator from .debug_utils import DebugOption, DebugUnderflowOverflow +from .distributed.fsdp import get_fsdp_ckpt_kwargs, update_fsdp_plugin_peft from .feature_extraction_sequence_utils import SequenceFeatureExtractor from .feature_extraction_utils import FeatureExtractionMixin from .hyperparameter_search import ALL_HYPERPARAMETER_SEARCH_BACKENDS, default_hp_search_backend @@ -66,7 +67,6 @@ is_deepspeed_available, propagate_args_to_deepspeed, ) -from .integrations.fsdp import get_fsdp_ckpt_kwargs, update_fsdp_plugin_peft from .integrations.liger import apply_liger_kernel from .integrations.neftune import activate_neftune, deactivate_neftune from .integrations.peft import MIN_PEFT_VERSION @@ -2394,9 +2394,10 @@ def get_cp_size(self) -> int: def get_tp_size(self) -> int: """Get the tensor parallel size from either the model or DeepSpeed config.""" - # 1. Check model.tp_size first - if (model_tp := getattr(self.model, "_tp_size", None)) is not None: - return model_tp + # TODO: adapt it cleaner with distributed config once distributed api is stable + dc = getattr(getattr(self.model, "config", None), "distributed_config", None) + if dc is not None and dc.tp_size is not None: + return dc.tp_size # 2. Fall back to DeepSpeed config if enabled if self.is_deepspeed_enabled and (deepspeed_config := getattr(self.args, "hf_deepspeed_config", None)): diff --git a/src/transformers/trainer_seq2seq.py b/src/transformers/trainer_seq2seq.py index ada588adbd21..c3907ed2556c 100644 --- a/src/transformers/trainer_seq2seq.py +++ b/src/transformers/trainer_seq2seq.py @@ -22,9 +22,9 @@ from torch import nn from torch.utils.data import Dataset +from .distributed.fsdp import is_fsdp_managed_module from .generation.configuration_utils import GenerationConfig from .integrations.deepspeed import is_deepspeed_zero3_enabled -from .integrations.fsdp import is_fsdp_managed_module from .trainer import Trainer from .utils import is_datasets_available, logging diff --git a/tests/causal_lm_tester.py b/tests/causal_lm_tester.py index 6b94a520d4f2..9df6b77447ae 100644 --- a/tests/causal_lm_tester.py +++ b/tests/causal_lm_tester.py @@ -30,6 +30,7 @@ ) from .test_configuration_common import ConfigTester +from .test_fsdp_mixin import FSDPTesterMixin from .test_modeling_common import ( GenerationTesterMixin, ModelTesterMixin, @@ -271,7 +272,12 @@ def prepare_config_and_inputs_for_common(self): @require_torch class CausalLMModelTest( - ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin, TrainingTesterMixin, TensorParallelTesterMixin + ModelTesterMixin, + GenerationTesterMixin, + PipelineTesterMixin, + TrainingTesterMixin, + TensorParallelTesterMixin, + FSDPTesterMixin, ): model_tester_class = None all_model_classes = None diff --git a/tests/tensor_parallel/test_tensor_parallel.py b/tests/tensor_parallel/test_tensor_parallel.py index 91770b683e45..e81aa342d635 100644 --- a/tests/tensor_parallel/test_tensor_parallel.py +++ b/tests/tensor_parallel/test_tensor_parallel.py @@ -11,56 +11,12 @@ # 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 math import warnings -from types import SimpleNamespace - -import torch from transformers import AutoModelForCausalLM -from transformers.integrations.tensor_parallel import ( - ColwiseParallel, - EmbeddingParallel, - GroupedGemmParallel, - PackedColwiseParallel, - PackedRowwiseParallel, - RowwiseParallel, - get_packed_weights, - repack_weights, -) from transformers.testing_utils import TestCasePlus, is_tensor_parallel_test -@is_tensor_parallel_test -class TestTensorParallelUtils(TestCasePlus): - def test_packed_unpacked_conversion(self): - WORLD_SIZE = 2 - PACKED_BLOCK_SIZE = 800 - SHARDING_DIM = 2 - NUM_BLOCKS = 2 - - original_packed_weights = torch.randn(4, 512, 2 * PACKED_BLOCK_SIZE) - original_packed_weights.get_dtype = lambda: "F32" # get_packed_weights expects PySlice object - empty_param = torch.empty(4, 512, 2 * PACKED_BLOCK_SIZE) - - class MockDeviceMesh: - def size(self): - return WORLD_SIZE - - mock_mesh = ( - MockDeviceMesh() - ) # get_packed_weights only calls `.size()`, do this to avoid doing actual distributed run - - packed_weights_0 = get_packed_weights(original_packed_weights, empty_param, mock_mesh, 0, SHARDING_DIM) - packed_weights_1 = get_packed_weights(original_packed_weights, empty_param, mock_mesh, 1, SHARDING_DIM) - - # simulate all gather of sharded weights - packed_weights = torch.cat([packed_weights_0, packed_weights_1], dim=SHARDING_DIM) - unpacked_weights = repack_weights(packed_weights, SHARDING_DIM, WORLD_SIZE, NUM_BLOCKS) - - assert torch.allclose(unpacked_weights, original_packed_weights) - - @is_tensor_parallel_test class TestTensorParallelProperties(TestCasePlus): def test_tp_plan_property_setter_getter(self): @@ -90,18 +46,6 @@ def test_tp_plan_property_setter_getter(self): } self.assertEqual(model.tp_plan, expected_plan) - def test_tp_plan_validation_invalid_style(self): - """Test that invalid parallel styles are rejected.""" - model_id = "hf-internal-testing/tiny-random-LlamaForCausalLM" - model = AutoModelForCausalLM.from_pretrained(model_id, dtype="auto") - - # Test invalid parallel style - with self.assertRaises(ValueError) as context: - model.tp_plan = {"layers.*.self_attn.q_proj": "invalid_style"} - - self.assertIn("Unsupported tensor parallel style 'invalid_style'", str(context.exception)) - self.assertIn("Supported styles are", str(context.exception)) - def test_tp_plan_validation_nonexistent_layer_warning(self): """Test that warnings are issued for non-existent layer patterns.""" @@ -165,256 +109,3 @@ def test_tp_plan_none_handling(self): # Test setting a plan after None model.tp_plan = {"model.layers.*.self_attn.q_proj": "colwise"} self.assertEqual(model.tp_plan, {"model.layers.*.self_attn.q_proj": "colwise"}) - - -@is_tensor_parallel_test -class TestTensorParallelLayer(TestCasePlus): - class MockDeviceMesh: - def __init__(self, world_size, rank): - self.world_size = world_size - self.rank = rank - self.shape = (world_size,) - - def size(self): - return self.world_size - - def get_local_rank(self): - return self.rank - - def test_colwise_get_expected_sharded_shape(self): - world_size = 3 - size = 10 # not divisible by world_size to test edge case - empty_param_2d = torch.empty(size, 32) - empty_param_1d = torch.empty((size,)) - step = math.ceil(size / world_size) - - for rank in range(world_size): - for empty_param in [empty_param_2d, empty_param_1d]: - device_mesh = self.MockDeviceMesh(world_size=world_size, rank=rank) - layer = ColwiseParallel(device_mesh=device_mesh, rank=rank, empty_param=empty_param) - - begin = rank * step - end = min(begin + step, size) - ground_truth = (end - begin,) + empty_param.shape[1:] - expected_shape = layer.get_expected_sharded_shape(empty_param.shape) - self.assertEqual( - expected_shape, ground_truth, f"Rank {rank} expected shape {ground_truth} but got {expected_shape}" - ) - - def test_rowwise_get_expected_sharded_shape(self): - world_size = 3 - size = 10 # not divisible by world_size to test edge case - empty_param_2d = torch.empty(32, size) - empty_param_1d = torch.empty((size,)) - step = math.ceil(size / world_size) - - for rank in range(world_size): - device_mesh = self.MockDeviceMesh(world_size=world_size, rank=rank) - - # 2D: shards on dim -1 (input features) - layer = RowwiseParallel(device_mesh=device_mesh, rank=rank, empty_param=empty_param_2d) - begin = rank * step - end = min(begin + step, size) - ground_truth = empty_param_2d.shape[:-1] + (end - begin,) - expected_shape = layer.get_expected_sharded_shape(empty_param_2d.shape) - self.assertEqual( - expected_shape, ground_truth, f"Rank {rank} expected shape {ground_truth} but got {expected_shape}" - ) - - # 1D bias: NOT sharded - layer = RowwiseParallel(device_mesh=device_mesh, rank=rank, empty_param=empty_param_1d) - self.assertEqual(layer.get_expected_sharded_shape(empty_param_1d.shape), empty_param_1d.shape) - - def test_embedding_get_expected_sharded_shape(self): - world_size = 3 - size = 10 # not divisible by world_size to test edge case; same size on both dims so step applies to both - empty_param = torch.empty(size, size) - step = math.ceil(size / world_size) - - for rank in range(world_size): - device_mesh = self.MockDeviceMesh(world_size=world_size, rank=rank) - begin = rank * step - end = min(begin + step, size) - - # embedding_dim_sharding=0: shards dim 0 (vocab) - layer = EmbeddingParallel( - device_mesh=device_mesh, rank=rank, empty_param=empty_param, embedding_dim_sharding=0 - ) - ground_truth = (end - begin,) + empty_param.shape[1:] - expected_shape = layer.get_expected_sharded_shape(empty_param.shape) - self.assertEqual( - expected_shape, ground_truth, f"Rank {rank} expected shape {ground_truth} but got {expected_shape}" - ) - - # embedding_dim_sharding=1: shards dim 1 (embedding dim) - layer = EmbeddingParallel( - device_mesh=device_mesh, rank=rank, empty_param=empty_param, embedding_dim_sharding=1 - ) - ground_truth = empty_param.shape[:1] + (end - begin,) + empty_param.shape[2:] - expected_shape = layer.get_expected_sharded_shape(empty_param.shape) - self.assertEqual( - expected_shape, ground_truth, f"Rank {rank} expected shape {ground_truth} but got {expected_shape}" - ) - - def test_grouped_gemm_get_expected_sharded_shape(self): - world_size = 3 - size = 9 # must be divisible by world_size (GroupedGemm requires it) - empty_param = torch.empty(size, 16, 32) - step = math.ceil(size / world_size) - - for rank in range(world_size): - device_mesh = self.MockDeviceMesh(world_size=world_size, rank=rank) - layer = GroupedGemmParallel(device_mesh=device_mesh, rank=rank, empty_param=empty_param) - begin = rank * step - end = min(begin + step, size) - ground_truth = (end - begin,) + empty_param.shape[1:] - expected_shape = layer.get_expected_sharded_shape(empty_param.shape) - self.assertEqual( - expected_shape, ground_truth, f"Rank {rank} expected shape {ground_truth} but got {expected_shape}" - ) - - def test_colwise_update_module_attributes(self): - device_mesh = self.MockDeviceMesh(world_size=4, rank=0) - - # gather_output=False (default): out_features is updated - module = torch.nn.Linear(32, 16) - layer = ColwiseParallel(device_mesh=device_mesh, rank=0, empty_param=torch.empty(16, 32)) - layer.update_module_attributes(module) - self.assertEqual(module.out_features, 4) - - # gather_output=True: out_features is NOT updated - module = torch.nn.Linear(32, 16) - layer = ColwiseParallel(device_mesh=device_mesh, rank=0, empty_param=torch.empty(16, 32), gather_output=True) - layer.update_module_attributes(module) - self.assertEqual(module.out_features, 16) - - def test_rowwise_update_module_attributes(self): - device_mesh = self.MockDeviceMesh(world_size=4, rank=0) - - module = torch.nn.Linear(32, 16) - layer = RowwiseParallel(device_mesh=device_mesh, rank=0, empty_param=torch.empty(16, 32)) - layer.update_module_attributes(module) - self.assertEqual(module.in_features, 8) - - def test_embedding_update_module_attributes(self): - device_mesh = self.MockDeviceMesh(world_size=4, rank=0) - - # embedding_dim_sharding=0: num_embeddings is updated - module = torch.nn.Embedding(32, 16) - layer = EmbeddingParallel( - device_mesh=device_mesh, rank=0, empty_param=torch.empty(32, 16), embedding_dim_sharding=0 - ) - layer.update_module_attributes(module) - self.assertEqual(module.num_embeddings, 8) - self.assertEqual(module.embedding_dim, 16) - - # embedding_dim_sharding=1: embedding_dim is updated - module = torch.nn.Embedding(32, 16) - layer = EmbeddingParallel( - device_mesh=device_mesh, rank=0, empty_param=torch.empty(32, 16), embedding_dim_sharding=1 - ) - layer.update_module_attributes(module) - self.assertEqual(module.num_embeddings, 32) - self.assertEqual(module.embedding_dim, 4) - - def test_grouped_gemm_update_module_attributes(self): - device_mesh = self.MockDeviceMesh(world_size=4, rank=0) - - # There is no torch module with num_experts attribute, it is more at the Transformers level, - # so just use a SimpleNamespace to test that the attribute is updated correctly. - module = SimpleNamespace(num_experts=8) - layer = GroupedGemmParallel(device_mesh=device_mesh, rank=0, empty_param=torch.empty(8, 16, 32)) - layer.update_module_attributes(module) - self.assertEqual(module.num_experts, 2) - - def test_update_module_attributes_missing_attribute(self): - device_mesh = self.MockDeviceMesh(world_size=4, rank=0) - module = SimpleNamespace(random_attr=123) - for cls in [ColwiseParallel, RowwiseParallel, GroupedGemmParallel]: - layer = cls(device_mesh=device_mesh, rank=0, empty_param=torch.empty(16, 32)) - layer.update_module_attributes(module) - - self.assertEqual( - module.__dict__, - {"random_attr": 123}, - "update_module_attributes should not modify attributes that don't exist", - ) - - def test_shard_tensor_shape_consistency(self): - """ - Test that shard_tensor returns tensors of the expected shape for different parallel styles and ranks. - """ - WORLD_SIZE = 4 - cases = [ - (ColwiseParallel, (16, 32), {}), - (ColwiseParallel, (16, 32), {"gather_output": True}), - (ColwiseParallel, (16,), {}), - (RowwiseParallel, (16, 32), {}), - (RowwiseParallel, (32,), {}), - (EmbeddingParallel, (32, 16), {"embedding_dim_sharding": 0}), - (EmbeddingParallel, (32, 16), {"embedding_dim_sharding": 1}), - ] - for cls, shape, kwargs in cases: - for rank in range(WORLD_SIZE): - device_mesh = self.MockDeviceMesh(world_size=WORLD_SIZE, rank=rank) - layer = cls(device_mesh=device_mesh, rank=rank, empty_param=torch.empty(*shape), **kwargs) - - full_tensor = torch.randn(*shape) - sharded = layer.shard_tensor(full_tensor) - expected = layer.get_expected_sharded_shape(shape) - - self.assertEqual(tuple(sharded.shape), expected, f"{cls.__name__} rank={rank} shape={shape}") - - def test_packed_colwise_shard_tensor(self): - WORLD_SIZE = 2 - # 3D empty_param - empty = torch.empty(2, 16, 64) - - # Packed vs unpacked path is determined by checking the following: - # input.dim() == get_expected_sharded_shape(empty_param).dim() - - # Packed - full_packed = torch.randn(2, 16, 64) - full_packed.get_dtype = lambda: "F32" - for rank in range(WORLD_SIZE): - device_mesh = self.MockDeviceMesh(world_size=WORLD_SIZE, rank=rank) - layer = PackedColwiseParallel(device_mesh=device_mesh, rank=rank, empty_param=empty) - sharded = layer.shard_tensor(full_packed) - expected_shape = (2, 8, 64) # last dim is packed size, middle dim is sharded - self.assertEqual(sharded.shape, expected_shape) - - # Unpacked - full_unpacked = torch.randn(16, 64) - for rank in range(WORLD_SIZE): - device_mesh = self.MockDeviceMesh(world_size=WORLD_SIZE, rank=rank) - layer = PackedColwiseParallel(device_mesh=device_mesh, rank=rank, empty_param=empty) - sharded = layer.shard_tensor(full_unpacked) - expected_shape = (8, 64) # last dim is not packed, so just sharded - self.assertEqual(sharded.shape, expected_shape) - - def test_packed_rowwise_shard_tensor(self): - WORLD_SIZE = 2 - # empty_param last dim = 64 signals the packed size (2 * 32) - empty = torch.empty(16, 64) - - # Packed vs unpacked path is determined by checking the following: - # input.shape[-1] < empty_param.shape[-1] - - # Packed - full_packed = torch.randn(16, 64) - full_packed.get_dtype = lambda: "F32" - for rank in range(WORLD_SIZE): - device_mesh = self.MockDeviceMesh(world_size=WORLD_SIZE, rank=rank) - layer = PackedRowwiseParallel(device_mesh=device_mesh, rank=rank, empty_param=empty) - sharded = layer.shard_tensor(full_packed) - expected_shape = (16, 32) # last dim is packed size, sharded - self.assertEqual(sharded.shape, expected_shape) - - # Unpacked - full_unpacked = torch.randn(16, 32) - for rank in range(WORLD_SIZE): - device_mesh = self.MockDeviceMesh(world_size=WORLD_SIZE, rank=rank) - layer = PackedRowwiseParallel(device_mesh=device_mesh, rank=rank, empty_param=empty) - sharded = layer.shard_tensor(full_unpacked) - expected_shape = (16, 16) # last dim is not packed, so just sharded - self.assertEqual(sharded.shape, expected_shape) diff --git a/tests/test_distributed_config.py b/tests/test_distributed_config.py new file mode 100644 index 000000000000..3cf5b4a914db --- /dev/null +++ b/tests/test_distributed_config.py @@ -0,0 +1,93 @@ +import json +import tempfile + +from transformers.distributed import DistributedConfig + + +class TestDistributedConfig: + def test_2d_parallelism(self): + dc = DistributedConfig(tp_size=2, fsdp_size=2) + assert dc.tp_size == 2 + assert dc.fsdp_size == 2 + assert dc.tp_plan is None + assert dc.fsdp_plan is None + + def test_tp_only_defaults_fsdp_to_1(self): + dc = DistributedConfig(tp_size=4) + assert dc.tp_size == 4 + assert dc.fsdp_size == 1 + assert dc.tp_plan is None + + def test_fsdp_only_defaults_tp_to_1(self): + dc = DistributedConfig(fsdp_size=4) + assert dc.tp_size == 1 + assert dc.fsdp_size == 4 + assert dc.fsdp_plan is None + assert dc.tp_plan is None + + def test_empty_config(self): + dc = DistributedConfig() + assert dc.tp_size is None + assert dc.fsdp_size is None + assert dc.tp_plan is None + assert dc.fsdp_plan is None + + def test_from_dict(self): + dc = DistributedConfig.from_dict({"tp_size": 2, "fsdp_size": 4}) + assert dc.tp_size == 2 + assert dc.fsdp_size == 4 + assert dc.tp_plan is None + + def test_from_dict_ignores_unknown_keys(self): + dc = DistributedConfig.from_dict({"tp_size": 2, "unknown_key": 42}) + assert dc.tp_size == 2 + assert not hasattr(dc, "unknown_key") + + def test_from_dict_kwargs_override(self): + dc = DistributedConfig.from_dict({"tp_size": 2}, tp_size=8) + assert dc.tp_size == 8 + + def test_to_dict(self): + dc = DistributedConfig(tp_size=2, fsdp_size=4) + d = dc.to_dict() + assert d == { + "tp_size": 2, + "tp_plan": None, + "enable_sequence_parallel": False, + "enable_expert_parallel": False, + "fsdp_size": 4, + "fsdp_plan": None, + } + + def test_to_dict_is_a_copy(self): + dc = DistributedConfig(tp_plan={"layer": "colwise"}) + d = dc.to_dict() + d["tp_plan"]["layer"] = "rowwise" + assert dc.tp_plan["layer"] == "colwise" + + def test_to_json_string(self): + dc = DistributedConfig(tp_size=2, fsdp_size=2) + s = dc.to_json_string() + parsed = json.loads(s) + assert parsed["tp_size"] == 2 + assert parsed["fsdp_size"] == 2 + + def test_to_json_file(self): + dc = DistributedConfig(tp_size=4) + with tempfile.NamedTemporaryFile(mode="r", suffix=".json", delete=False) as f: + dc.to_json_file(f.name) + f.seek(0) + parsed = json.load(f) + assert parsed["tp_size"] == 4 + assert parsed["tp_plan"] is None + + def test_roundtrip_dict(self): + original = DistributedConfig(tp_size=2, fsdp_size=4) + restored = DistributedConfig.from_dict(original.to_dict()) + assert original == restored + + def test_repr(self): + dc = DistributedConfig(tp_size=2) + r = repr(dc) + assert "DistributedConfig" in r + assert '"tp_size": 2' in r diff --git a/tests/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py new file mode 100644 index 000000000000..1a7032489224 --- /dev/null +++ b/tests/test_fsdp_mixin.py @@ -0,0 +1,798 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +"""FSDP tester mixin for model tests.""" + +import json +import logging +import os +import socket +import sys +import tempfile +import time +import traceback +from abc import ABC, abstractmethod + +from transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM, is_torch_available +from transformers.testing_utils import ( + backend_device_count, + backend_empty_cache, + backend_torch_accelerator_module, + init_test_logger, + is_fsdp_test, + require_fsdp, +) +from transformers.trainer_utils import set_seed + + +logger = logging.getLogger("transformers.training_test") + + +if is_torch_available(): + import torch + import torch.distributed as dist + import torch.multiprocessing as mp + from torch.nn.parallel import DistributedDataParallel as DDP + + from transformers.distributed import DistributedConfig + from transformers.distributed.fsdp import apply_fully_shard_data_parallel, initialize_fsdp + from transformers.distributed.tensor_parallel import replace_layer_number_by_wildcard + from transformers.distributed.utils import ( + gather_full_state_dict, + load_optimizer_distributed, + save_optimizer_distributed, + ) + + +# ============================================================================= +# Constants +# ============================================================================= + +BATCH_SIZE = 2 +SEQ_LEN = 64 +NUM_STEPS = 20 +LR = 3e-4 +SEED = 42 +FSDP_TOP_MODEL_NAMES = { + # FSDP coverage is gated on models declaring `base_model_fsdp_plan` (config) + + # class-level `_fsdp_plan` (head class). Listed here are models whose test + # class extends `CausalLMModelTest` (so they pick up the FSDP mixin) and + # which use the standard embed_tokens / layers.* / norm naming. + # Dense + "llama", + "mistral", + "qwen3", + "phi", + "olmo3", + "gemma2", + # MoE + "mixtral", + "qwen3_moe", + "qwen2_moe", + "qwen3_5_moe", + "deepseek_v2", + "gpt_oss", + "glm4_moe_lite", +} + + +# ============================================================================= +# Distributed helpers (top-level for pickling by mp.spawn) +# ============================================================================= + + +def _get_distributed_device_type(): + device_type = torch._C._get_accelerator().type + return "cpu" if device_type == "mps" else device_type + + +def _get_distributed_backend(): + backend_map = {"cpu": "gloo", "cuda": "nccl", "xpu": "xccl", "hpu": "hccl"} + return backend_map.get(_get_distributed_device_type(), "gloo") + + +def _get_rank_device(rank): + device_type = _get_distributed_device_type() + if device_type == "cpu": + return torch.device("cpu") + return torch.device(device_type, rank) + + +def _set_rank_device(rank): + accelerator_module = backend_torch_accelerator_module(_get_distributed_device_type()) + if accelerator_module is not None and hasattr(accelerator_module, "set_device"): + accelerator_module.set_device(rank) + + +def _get_accelerator_rng_state(): + accelerator_module = backend_torch_accelerator_module(_get_distributed_device_type()) + if accelerator_module is None or not hasattr(accelerator_module, "get_rng_state"): + return None + return accelerator_module.get_rng_state() + + +def _set_accelerator_rng_state(rng_state): + accelerator_module = backend_torch_accelerator_module(_get_distributed_device_type()) + if rng_state is not None and accelerator_module is not None and hasattr(accelerator_module, "set_rng_state"): + accelerator_module.set_rng_state(rng_state) + + +def _get_available_fsdp_workers(): + if _get_distributed_device_type() == "cpu": + return os.cpu_count() or 1 + return backend_device_count(_get_distributed_device_type()) + + +def _fsdp_global_wrapper(rank, test_name, func, func_args, func_kwargs, world_size, port, results_file): + os.environ["WORLD_SIZE"] = str(world_size) + os.environ["RANK"] = str(rank) + os.environ["LOCAL_RANK"] = str(rank) + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = str(port) + + _set_determinism(SEED) + + dist.init_process_group(backend=_get_distributed_backend(), rank=rank, world_size=world_size) + _set_rank_device(rank) + + if rank == 0: + start_time = time.perf_counter() + print(f"[FSDP] Starting test: {test_name}", flush=True) + + error = None + try: + func(rank, *func_args, **func_kwargs) + except Exception as e: + error = f"{type(e).__name__}: {e}\n{traceback.format_exc()}" + + error_flag = torch.tensor([1 if error else 0], device=_get_rank_device(rank)) + dist.all_reduce(error_flag, op=dist.ReduceOp.MAX) + any_failed = error_flag.item() > 0 + + if rank == 0: + elapsed = time.perf_counter() - start_time + status = "FAIL" if any_failed else "PASS" + output_stream = sys.stderr if any_failed else sys.stdout + print(f"[FSDP] {status} test: {test_name} ({elapsed:.1f}s)", file=output_stream, flush=True) + with open(results_file, "w") as f: + json.dump({"error": error or ("Failed on another rank" if any_failed else None)}, f) + + backend_empty_cache(_get_distributed_device_type()) + dist.barrier() + dist.destroy_process_group() + + +def _set_determinism(seed): + torch.use_deterministic_algorithms(True) + if _get_distributed_device_type() == "cuda" and torch.cuda.is_available(): + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + set_seed(seed) + + +# ============================================================================= +# Training & comparison helpers (top-level for pickling) +# ============================================================================= + + +def _build_repeated_training_batches(config, device, num_steps): + """Create one deterministic batch and reuse it across steps.""" + generator = torch.Generator(device=device) + generator.manual_seed(SEED) + input_ids = torch.randint(0, config.vocab_size, (BATCH_SIZE, SEQ_LEN), device=device, generator=generator) + labels = input_ids.clone() + return [(input_ids, labels)] * num_steps + + +def _create_shared_tmpdir(rank): + if rank == 0: + tmpdir_obj = tempfile.TemporaryDirectory() + tmpdir = tmpdir_obj.name + tmpdir_list = [tmpdir] + else: + tmpdir_obj = None + tmpdir_list = [None] + dist.broadcast_object_list(tmpdir_list, src=0) + return tmpdir_list[0], tmpdir_obj + + +def _gather_ddp_state_dict(model): + # Only rank 0 returns data to match gather_full_state_dict semantics, so the + # downstream DDP-vs-FSDP comparison runs once on rank 0 instead of N times. + if dist.get_rank() != 0: + return {} + return {k: v.clone().detach().cpu() for k, v in model.module.state_dict().items()} + + +def _resolve_fsdp_plan_paths(model): + """Expand model._fsdp_plan into (paths, strategy) entries. + + Wildcard keys are expanded via ``replace_layer_number_by_wildcard``. When + weights are tied, the standalone embed_tokens entry is skipped and any + ``"lm_head"`` keep entry is rewritten to the tied source path (the keep + group will wrap the shared parameter once). + """ + plan = model._fsdp_plan + + input_embed = model.get_input_embeddings() + output_embed = model.get_output_embeddings() + weights_tied = ( + input_embed is not None + and output_embed is not None + and hasattr(input_embed, "weight") + and hasattr(output_embed, "weight") + and input_embed.weight is output_embed.weight + ) + tied_source = None + if weights_tied: + for name, mod in model.named_modules(): + if mod is input_embed: + tied_source = name + break + + name_to_module = dict(model.named_modules()) + entries: list[tuple[list[str], str]] = [] + for key, strategy in plan.items(): + if weights_tied and key == tied_source: + continue + if weights_tied and key == "lm_head" and strategy == "keep_full_weight": + entries.append(([tied_source], strategy)) + continue + + if key in name_to_module: + entries.append(([key], strategy)) + continue + matched = [name for name in name_to_module if replace_layer_number_by_wildcard(name) == key] + if matched: + entries.append((matched, strategy)) + return entries + + +def _build_manual_fsdp_plan(config, device, policy_options=None): + """Build a manual FSDP2 plan by expanding model._fsdp_plan.""" + policy_options = policy_options or [] + set_seed(SEED) + model = AutoModelForCausalLM.from_config(config).to(device) + + module_plan: dict[str, list[str]] = {} + for paths, strategy in _resolve_fsdp_plan_paths(model): + for path in paths: + module_plan[path] = [strategy, *policy_options] + + del model + return {"modules": module_plan} + + +def _save_init_pretrained(rank, config, dtype): + """Save a deterministic initial model to a shared tmpdir for from_pretrained loading.""" + tmpdir, tmpdir_obj = _create_shared_tmpdir(rank) + if rank == 0: + set_seed(SEED) + model = AutoModelForCausalLM.from_config(config).to(dtype) + model.save_pretrained(tmpdir) + del model + dist.barrier() + return tmpdir, tmpdir_obj + + +def _save_training_state(model, optimizer, training_state_dir): + """Save optimizer (canonical DCP path) plus per-rank RNG for resume.""" + save_optimizer_distributed(model, optimizer, os.path.join(training_state_dir, "optim")) + rng = {"cpu": torch.get_rng_state()} + accel = _get_accelerator_rng_state() + if accel is not None: + rng["accel"] = accel + torch.save(rng, os.path.join(training_state_dir, f"rng_rank{dist.get_rank()}.pt")) + + +def _load_training_state(model, optimizer, training_state_dir): + """Inverse of `_save_training_state`.""" + load_optimizer_distributed(model, optimizer, os.path.join(training_state_dir, "optim")) + rng = torch.load(os.path.join(training_state_dir, f"rng_rank{dist.get_rank()}.pt"), weights_only=False) + torch.set_rng_state(rng["cpu"]) + if "accel" in rng: + _set_accelerator_rng_state(rng["accel"]) + + +def train_ddp(rank, batches, lr, device, dtype, init_model_dir): + _set_determinism(SEED) + model = AutoModelForCausalLM.from_pretrained(init_model_dir, torch_dtype=dtype, attn_implementation="eager").to( + device + ) + # MoE/conditional-routing variants) may not use all params on + # every step, and DDP would otherwise fail. Specifying find_unused_parameters=True allows running backward on a subgraph of the model. + ddp_kwargs = {"find_unused_parameters": True} + if device.type != "cpu": + ddp_kwargs["device_ids"] = [rank] + ddp_model = DDP(model, **ddp_kwargs) + ddp_model.train() + optimizer = torch.optim.Adam(ddp_model.parameters(), lr=lr) + + losses, grad_norms = [], [] + for input_ids, labels in batches: + optimizer.zero_grad() + output = ddp_model(input_ids=input_ids, labels=labels, use_cache=False) + loss = output.loss + loss.backward() + grad_norm = torch.nn.utils.clip_grad_norm_(ddp_model.parameters(), max_norm=float("inf")) + optimizer.step() + + losses.append(loss.detach().item()) + grad_norms.append(grad_norm) + + state_dict = _gather_ddp_state_dict(ddp_model) + + del optimizer, ddp_model, model + backend_empty_cache(_get_distributed_device_type()) + dist.barrier() + + return losses, grad_norms, state_dict + + +def train_fsdp2( + rank, + batches, + lr, + dtype, + init_model_dir, + checkpoint_step, + fsdp_plan, +): + # -- Phase 1: Pre-checkpoint run -- train only the first `checkpoint_step` steps, then save + _set_determinism(SEED) + distributed_config = DistributedConfig(fsdp_size=dist.get_world_size(), fsdp_plan=fsdp_plan) + pre_ckpt_model = AutoModelForCausalLM.from_pretrained( + init_model_dir, + torch_dtype=dtype, + distributed_config=distributed_config, + attn_implementation="eager", + ) + pre_ckpt_model.train() + pre_ckpt_optimizer = torch.optim.Adam(pre_ckpt_model.parameters(), lr=lr) + + pre_ckpt_losses, pre_ckpt_grad_norms = [], [] + for step in range(0, checkpoint_step): + input_ids, labels = batches[step] + pre_ckpt_optimizer.zero_grad() + output = pre_ckpt_model(input_ids=input_ids, labels=labels, use_cache=False) + loss = output.loss + loss.backward() + grad_norm = torch.nn.utils.clip_grad_norm_(pre_ckpt_model.parameters(), max_norm=float("inf")) + pre_ckpt_optimizer.step() + + pre_ckpt_losses.append(loss.detach().item()) + pre_ckpt_grad_norms.append(grad_norm) + + # -- Phase 2: Save checkpoint, then load into a fresh model + # tmpdir/ + # model/ <- HF safetensors via save_pretrained (DCP + consolidation) + # training_state/ <- distcp (optimizer + RNG) + tmpdir, tmpdir_obj = _create_shared_tmpdir(rank) + try: + model_dir = os.path.join(tmpdir, "model") + training_state_dir = os.path.join(tmpdir, "training_state") + + pre_ckpt_model.save_pretrained(model_dir, is_main_process=(rank == 0)) + _save_training_state(pre_ckpt_model, pre_ckpt_optimizer, training_state_dir) + dist.barrier() + + # Intentionally scramble RNG to prove checkpoint restore works + _set_determinism(SEED + 1234) + resumed_model = AutoModelForCausalLM.from_pretrained( + model_dir, + torch_dtype=dtype, + distributed_config=distributed_config, + attn_implementation="eager", + ) + resumed_model.train() + resumed_optimizer = torch.optim.Adam(resumed_model.parameters(), lr=lr) + + _load_training_state(resumed_model, resumed_optimizer, training_state_dir) + dist.barrier() + finally: + if rank == 0: + tmpdir_obj.cleanup() + + # -- Phase 3: Post-checkpoint run -- continue training the remaining steps from the resumed model + post_ckpt_losses, post_ckpt_grad_norms = [], [] + for step in range(checkpoint_step, len(batches)): + input_ids, labels = batches[step] + resumed_optimizer.zero_grad() + output = resumed_model(input_ids=input_ids, labels=labels, use_cache=False) + loss = output.loss + loss.backward() + grad_norm = torch.nn.utils.clip_grad_norm_(resumed_model.parameters(), max_norm=float("inf")) + resumed_optimizer.step() + + post_ckpt_losses.append(loss.detach().item()) + post_ckpt_grad_norms.append(grad_norm) + + combined_losses = pre_ckpt_losses + post_ckpt_losses + combined_grad_norms = pre_ckpt_grad_norms + post_ckpt_grad_norms + combined_state_dict = gather_full_state_dict(resumed_model) + + return combined_losses, combined_grad_norms, combined_state_dict + + +# ============================================================================= +# Distributed test implementations (top-level for pickling by mp.spawn) +# ============================================================================= +def _test_fsdp2_save_load_impl(rank, config_class, config_dict): + """Train FSDP2 model, save via save_pretrained, load via from_pretrained, compare state dicts.""" + init_test_logger() + + device = _get_rank_device(rank) + config = config_class.from_dict(config_dict) + + batches = _build_repeated_training_batches(config, device, 3) + + distributed_config = DistributedConfig(fsdp_size=dist.get_world_size()) + + init_tmpdir, init_tmpdir_obj = _save_init_pretrained(rank, config, torch.float32) + try: + _set_determinism(SEED) + model = AutoModelForCausalLM.from_pretrained( + init_tmpdir, + distributed_config=distributed_config, + attn_implementation="eager", + ) + dist.barrier() + finally: + if rank == 0 and init_tmpdir_obj is not None: + init_tmpdir_obj.cleanup() + model.train() + optimizer = torch.optim.Adam(model.parameters(), lr=LR) + + for input_ids, labels in batches: + optimizer.zero_grad() + output = model(input_ids=input_ids, labels=labels, use_cache=False) + output.loss.backward() + optimizer.step() + + state_dict_before = gather_full_state_dict(model) + + tmpdir, tmpdir_obj = _create_shared_tmpdir(rank) + try: + model.save_pretrained(tmpdir, is_main_process=(rank == 0)) + dist.barrier() + + new_model = AutoModelForCausalLM.from_pretrained( + tmpdir, + distributed_config=distributed_config, + attn_implementation="eager", + ) + dist.barrier() + finally: + if rank == 0: + tmpdir_obj.cleanup() + + state_dict_after = gather_full_state_dict(new_model) + + for key in state_dict_before: + assert key in state_dict_after, f"Key {key} missing after load" + torch.testing.assert_close( + state_dict_before[key], + state_dict_after[key], + rtol=0, + atol=0, + msg=f"Weight mismatch for {key} after save/load", + ) + + if rank == 0: + logger.debug(f"FSDP2 save/load test passed: all {len(state_dict_before)} parameters match exactly.") + + +def _test_fsdp2_sharding_structure_impl(rank, config_class, config_dict, tie_word_embeddings): + """ + Verify that apply_fully_shard_data_parallel(fsdp_plan=None) wraps exactly the right modules. + + Expected FSDP targets: + UNTIED TIED + ────── ──── + 1. embed_tokens (reshard=True) 1. (skip — embed goes to step 3) + 2. layers[i] (reshard=True) 2. layers[i] (reshard=True) + 3. [norm, lm_head] (reshard=False) 3. [norm, embed_tokens] (reshard=False) + 4. root 4. root + """ + init_test_logger() + + config = config_class.from_dict(config_dict) + config.tie_word_embeddings = tie_word_embeddings + + auto_plan = None + device_map, device_mesh, _ = initialize_fsdp(fsdp_plan={}) + + set_seed(SEED) + model = AutoModelForCausalLM.from_config(config).to(device_map) + + # Expected FSDP targets come from model._fsdp_plan: every resolved path gets a + # fully_shard call (keep_full_weight entries are bundled into one group, but each + # member still appears as an FSDP-wrapped module in named_modules), plus the root. + expected_targets = {""} + for paths, _strategy in _resolve_fsdp_plan_paths(model): + expected_targets.update(paths) + + model = apply_fully_shard_data_parallel(model, device_mesh, fsdp_plan=auto_plan) + + actual_targets = {name for name, module in model.named_modules() if type(module).__name__.startswith("FSDP")} + + if rank == 0: + logger.debug(f" Weights tied: {config.tie_word_embeddings}") + logger.debug(f" Expected FSDP targets: {sorted(expected_targets)}") + logger.debug(f" Actual FSDP targets: {sorted(actual_targets)}") + + missing = expected_targets - actual_targets + extra = actual_targets - expected_targets + assert not missing and not extra, ( + f"FSDP target mismatch.\n" + f" Missing (expected but not wrapped): {sorted(missing)}\n" + f" Extra (wrapped but not expected): {sorted(extra)}" + ) + + if rank == 0: + logger.debug(f" FSDP sharding structure OK ({len(actual_targets)} targets)") + + +def _test_fsdp2_plan_vs_ddp_impl( + rank, config_class, config_dict, tie_word_embeddings, plan_mode, policy_options=None, dtype=None +): + """Validate DDP-vs-FSDP2 trace matching for either auto or manual plan mode.""" + init_test_logger() + + if dtype is None: + dtype = torch.float32 + + policy_options = policy_options or [] + assert "mixed_precision" not in policy_options, ( + "Use the mixed-precision specific tests when enabling mixed_precision policy." + ) + + device = _get_rank_device(rank) + config = config_class.from_dict(config_dict) + config.tie_word_embeddings = tie_word_embeddings + + if plan_mode == "auto": + fsdp_plan = { + "cpu_offload": "cpu_offload" in policy_options, + "mixed_precision": "mixed_precision" in policy_options, + } + test_label = f"FSDP2(auto{'+policies' if policy_options else ''})" + elif plan_mode == "manual": + fsdp_plan = _build_manual_fsdp_plan(config, device, policy_options=policy_options) + test_label = f"FSDP2(manual{'+policies' if policy_options else ''})" + else: + raise ValueError(f"Unsupported plan_mode '{plan_mode}'. Expected 'auto' or 'manual'.") + + checkpoint_step = NUM_STEPS // 2 + init_model_dir, init_tmpdir_obj = _save_init_pretrained(rank, config, dtype) + batches = _build_repeated_training_batches(config, device, NUM_STEPS) + try: + ddp_losses, ddp_grad_norms, ddp_state_dict = train_ddp(rank, batches, LR, device, dtype, init_model_dir) + + fsdp_losses, fsdp_grad_norms, fsdp_state_dict = train_fsdp2( + rank, + batches, + LR, + dtype, + init_model_dir=init_model_dir, + checkpoint_step=checkpoint_step, + fsdp_plan=fsdp_plan, + ) + finally: + if rank == 0 and init_tmpdir_obj is not None: + init_tmpdir_obj.cleanup() + + for step in range(len(ddp_losses)): + torch.testing.assert_close( + torch.tensor(ddp_losses[step]), + torch.tensor(fsdp_losses[step]), + rtol=1e-5, + atol=1e-5, + msg=f"Loss mismatch at step {step}: DDP={ddp_losses[step]}, {test_label}={fsdp_losses[step]}", + ) + torch.testing.assert_close( + torch.tensor(ddp_grad_norms[step]), + torch.tensor(fsdp_grad_norms[step]), + rtol=1e-5, + atol=1e-5, + msg=f"Grad norm mismatch at step {step}: DDP={ddp_grad_norms[step]}, {test_label}={fsdp_grad_norms[step]}", + ) + + for key in ddp_state_dict: + assert key in fsdp_state_dict, f"Key {key} missing from {test_label} state dict" + torch.testing.assert_close( + ddp_state_dict[key], + fsdp_state_dict[key], + rtol=1e-5, + atol=1e-5, + msg=f"Weight mismatch for {key}: DDP vs {test_label}", + ) + + if rank == 0: + logger.debug(f"DDP and {test_label} comparison checks passed.") + + +# ============================================================================= +# Mixin class +# ============================================================================= + + +class FSDPTesterMixin(ABC): + fsdp_nproc_per_node: int = 2 + skip_fsdp_tests: bool = False + + @property + @abstractmethod + def model_tester(self): + """The model tester instance (e.g., CausalLMModelTester).""" + ... + + def _skip_if_fsdp_disabled(self): + if self.skip_fsdp_tests: + self.skipTest("FSDP tests disabled for this model (skip_fsdp_tests=True)") + + def _skip_if_insufficient_devices(self): + self._skip_if_fsdp_disabled() + available_workers = _get_available_fsdp_workers() + if available_workers < self.fsdp_nproc_per_node: + self.skipTest(f"Need at least {self.fsdp_nproc_per_node} FSDP workers, have {available_workers}") + + def _get_fsdp_model_name(self): + module_parts = self.__class__.__module__.split(".") + if len(module_parts) >= 3 and module_parts[0] == "tests" and module_parts[1] == "models": + return module_parts[2] + return None + + def _skip_if_fsdp_model_not_selected(self): + model_name = self._get_fsdp_model_name() + if model_name not in FSDP_TOP_MODEL_NAMES: + model_label = model_name or self.__class__.__module__ + self.skipTest( + "FSDP mixin coverage is currently limited to the top-10 dense and top-10 MoE model suites " + f"(skipping {model_label})." + ) + + def _create_model_on_meta(self, config): + """Instantiate a model on the meta device (no memory allocated).""" + auto_classes = [AutoModelForCausalLM, AutoModelForSeq2SeqLM] + for auto_cls in auto_classes: + try: + with torch.device("meta"): + return auto_cls.from_config(config) + except Exception: + continue + self.skipTest(f"Cannot instantiate model with any Auto class for config {type(config).__name__}") + + def _get_tiny_config(self): + """Get config class and serialized dict for passing to spawned processes.""" + config = self.model_tester.get_config() + config.vocab_size = 256 + config.hidden_size = 64 + config.intermediate_size = 128 + if hasattr(config, "ffn_config"): + # Keep nested FFN projections consistent with resized hidden size. + if hasattr(config.ffn_config, "ffn_hidden_size"): + config.ffn_config.ffn_hidden_size = config.hidden_size + if hasattr(config.ffn_config, "hidden_size"): + config.ffn_config.hidden_size = config.intermediate_size + if hasattr(config, "num_attention_heads"): + config.num_attention_heads = 4 + if hasattr(config, "num_key_value_heads"): + config.num_key_value_heads = 4 + if hasattr(config, "vocab_size_per_layer_input"): + config.vocab_size_per_layer_input = config.vocab_size + # `to_diff_dict()` avoids nested config pollution (e.g. DBRX ffn_config receiving + # generic PretrainedConfig keys that its constructor rejects). + config_dict = config.to_diff_dict() + return type(config), config_dict + + def _run_fsdp2_distributed_test(self, test_name, test_impl, *test_args, **test_kwargs): + self._skip_if_fsdp_model_not_selected() + self._skip_if_insufficient_devices() + + config_class, config_dict = self._get_tiny_config() + func_args = (config_class, config_dict, *test_args) + + results_file = tempfile.mktemp(suffix=".json") + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("", 0)) + port = s.getsockname()[1] + + try: + mp.spawn( + _fsdp_global_wrapper, + args=(test_name, test_impl, func_args, test_kwargs, self.fsdp_nproc_per_node, port, results_file), + nprocs=self.fsdp_nproc_per_node, + ) + + with open(results_file) as f: + result = json.load(f) + finally: + if os.path.exists(results_file): + os.unlink(results_file) + + if result["error"] is not None: + self.fail(f"FSDP test '{test_name}' failed:\n{result['error']}") + + # ========================================================================= + # Test: get_transformer_block_classes (CPU, meta device) + # ========================================================================= + + @is_fsdp_test + def test_fsdp_plan_declared(self): + """The model exposes a non-empty `_fsdp_plan` derived from config + class-level overrides.""" + self._skip_if_fsdp_disabled() + self._skip_if_fsdp_model_not_selected() + start_time = time.perf_counter() + logger.info("[FSDP] Starting test: test_fsdp_plan_declared") + status = "FAIL" + try: + config = self.model_tester.get_config() + model = self._create_model_on_meta(config) + + plan = getattr(model, "_fsdp_plan", None) + self.assertTrue(plan, f"No _fsdp_plan declared for {type(model).__name__}") + status = "PASS" + finally: + logger.info("[FSDP] %s test: test_fsdp_plan_declared (%.1fs)", status, time.perf_counter() - start_time) + + @is_fsdp_test + @require_fsdp + def test_fsdp2_sharding_structure_untied(self): + self._run_fsdp2_distributed_test( + "test_fsdp2_sharding_structure_untied", _test_fsdp2_sharding_structure_impl, False + ) + + @is_fsdp_test + @require_fsdp + def test_fsdp2_sharding_structure_tied(self): + self._run_fsdp2_distributed_test( + "test_fsdp2_sharding_structure_tied", _test_fsdp2_sharding_structure_impl, True + ) + + @is_fsdp_test + @require_fsdp + def test_fsdp2_save_load(self): + self._run_fsdp2_distributed_test("test_fsdp2_save_load", _test_fsdp2_save_load_impl) + + @is_fsdp_test + @require_fsdp + def test_fsdp2_auto_plan_vs_ddp_untied(self): + self._run_fsdp2_distributed_test( + "test_fsdp2_auto_plan_vs_ddp_untied", _test_fsdp2_plan_vs_ddp_impl, False, "auto" + ) + + @is_fsdp_test + @require_fsdp + def test_fsdp2_auto_plan_vs_ddp_tied(self): + self._run_fsdp2_distributed_test( + "test_fsdp2_auto_plan_vs_ddp_tied", _test_fsdp2_plan_vs_ddp_impl, True, "auto" + ) + + @is_fsdp_test + @require_fsdp + def test_fsdp2_manual_plan_vs_ddp_untied(self): + self._run_fsdp2_distributed_test( + "test_fsdp2_manual_plan_vs_ddp_untied", _test_fsdp2_plan_vs_ddp_impl, False, "manual" + ) + + @is_fsdp_test + @require_fsdp + def test_fsdp2_manual_plan_vs_ddp_tied(self): + self._run_fsdp2_distributed_test( + "test_fsdp2_manual_plan_vs_ddp_tied", _test_fsdp2_plan_vs_ddp_impl, True, "manual" + ) diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py index 7720205c6fef..2e60939f9ead 100644 --- a/tests/test_modeling_common.py +++ b/tests/test_modeling_common.py @@ -131,8 +131,8 @@ from torch import nn from transformers import MODEL_MAPPING + from transformers.distributed.tensor_parallel import _get_parameter_tp_plan from transformers.integrations.accelerate import compute_module_sizes - from transformers.integrations.tensor_parallel import _get_parameter_tp_plan from transformers.modeling_utils import load_state_dict from transformers.pytorch_utils import id_tensor_storage diff --git a/tests/test_tensor_parallel_mixin.py b/tests/test_tensor_parallel_mixin.py index 547bce7dacc4..d71c398c039e 100644 --- a/tests/test_tensor_parallel_mixin.py +++ b/tests/test_tensor_parallel_mixin.py @@ -16,8 +16,9 @@ from abc import ABC, abstractmethod from transformers import TorchAoConfig, set_seed -from transformers.distributed.configuration_utils import DistributedConfig -from transformers.integrations.tensor_parallel import _get_parameter_tp_plan +from transformers.distributed import DistributedConfig +from transformers.distributed.sharding_utils import _replicate_dtensor +from transformers.distributed.tensor_parallel import _get_parameter_tp_plan from transformers.testing_utils import ( is_tensor_parallel_test, is_torch_available, @@ -33,9 +34,26 @@ import torch import torch.distributed as dist import torch.multiprocessing as mp + from torch.distributed.tensor import DTensor from torch.multiprocessing.spawn import ProcessRaisedException +def _to_local(tensor): + """Extract local tensor from DTensor, or return as-is for plain tensors.""" + if hasattr(tensor, "to_local"): + # NOTE(3outeille): With Sequence Parallelism, replicated params (e.g. norm weights) get Partial + # gradients — each rank holds only its contribution from its sequence shard. + # We must all-reduce (redistribute to Replicate) before extracting, otherwise + # we'd compare an incomplete gradient against the full reference. + # In the case of real training, we will always use SP + FSDP where the last will all-reduce the + # Partial gradients for us. + + if isinstance(tensor, DTensor) and any(not p.is_replicate() for p in tensor.placements): + tensor = _replicate_dtensor(tensor) + return tensor.to_local() + return tensor + + def _find_free_port(): """Find a free port by binding a socket and releasing it.""" with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: @@ -112,50 +130,65 @@ def wrapper(*args, **kwargs): return _init_distributed_inner -def _load_tp_and_reference_models(model_path, model_class): +def _load_tp_and_reference_models(model_path, model_class, enable_sequence_parallel=False): """Load TP model and non-TP reference model for comparison. Returns: tuple: (model_tp, model_ref, device) """ - model_tp = model_class.from_pretrained(model_path, tp_plan="auto") + tp_size = dist.get_world_size() + distributed_config = DistributedConfig(tp_size=tp_size, enable_sequence_parallel=enable_sequence_parallel) + model_tp = model_class.from_pretrained( + model_path, distributed_config=distributed_config, attn_implementation="sdpa" + ) dist.barrier() device = model_tp.device - model_ref = model_class.from_pretrained(model_path) + model_ref = model_class.from_pretrained(model_path, attn_implementation="sdpa") model_ref = model_ref.to(device) return model_tp, model_ref, device +def _get_active_tp_plan(model_tp): + distributed_config = getattr(model_tp.config, "distributed_config", None) + tp_plan = getattr(distributed_config, "tp_plan", None) + return tp_plan or getattr(model_tp, "_tp_plan", None) or {} + + def _verify_tp_sharding(rank, model_tp, model_ref): """Verify TP sharding by comparing parameter shapes between TP and reference models. + For DTensor params, uses the local tensor shape (not the global DTensor shape). + Returns: list: Names of sharded parameters """ world_size = dist.get_world_size() sharded_params = [] + tp_plan = _get_active_tp_plan(model_tp) for (name, param), (_, param_full) in zip(model_tp.named_parameters(), model_ref.named_parameters()): - if param.shape != param_full.shape: + # For DTensor params, get the local shape for comparison + param_local = _to_local(param) + if param_local.shape != param_full.shape: sharded_params.append(name) if rank == 0: - print(f"[TP Test Debug] TP sharded: {name} - full: {param_full.shape} -> sharded: {param.shape}") + print(f"[TP Test Debug] TP sharded: {name} - full: {param_full.shape} -> sharded: {param_local.shape}") # Verify sharding is correct - for dim in range(param.ndim): - if param.size(dim) != param_full.size(dim): - param_plan = _get_parameter_tp_plan(name, model_tp.tp_plan, is_weight=True) - if param_plan in ("packed_colwise",): + for dim in range(param_local.ndim): + if param_local.size(dim) != param_full.size(dim): + param_plan = _get_parameter_tp_plan(name, tp_plan, is_weight=True) + if param_plan == "packed_colwise": expected_size = param_full.size(dim) // world_size - assert param.size(dim) == expected_size, ( - f"Packed weight {name} sharding incorrect: expected {expected_size}, got {param.size(dim)}" + assert param_local.size(dim) == expected_size, ( + f"Packed weight {name} sharding incorrect: expected {expected_size}, got {param_local.size(dim)}" ) else: expected_size = (param_full.size(dim) + world_size - 1) // world_size - assert param.size(dim) <= expected_size, ( - f"Weight {name} sharding incorrect: expected <= {expected_size}, got {param.size(dim)}" + assert param_local.size(dim) <= expected_size, ( + f"Weight {name} sharding incorrect: expected <= {expected_size}, got {param_local.size(dim)}" ) break @@ -166,7 +199,7 @@ def _test_tp_forward_impl(_rank, model_path, model_class, atol, rtol): """Implementation for comparing TP and non-TP model outputs.""" set_seed(0) - model_tp, model, device = _load_tp_and_reference_models(model_path, model_class) + model_tp, model, device = _load_tp_and_reference_models(model_path, model_class, enable_sequence_parallel=True) _verify_tp_sharding(_rank, model_tp, model) @@ -179,7 +212,7 @@ def _test_tp_forward_impl(_rank, model_path, model_class, atol, rtol): with torch.no_grad(): logits = model(input_ids).logits - logits_tp = model_tp(input_ids).logits + logits_tp = _to_local(model_tp(input_ids).logits) diff = (logits - logits_tp).abs() assert torch.allclose(logits, logits_tp, atol=atol, rtol=rtol), ( @@ -193,7 +226,8 @@ def _test_tp_backward_impl(rank, model_path, model_class, atol, rtol): """Implementation for comparing TP and non-TP model backward passes.""" set_seed(0) - model_tp, model, device = _load_tp_and_reference_models(model_path, model_class) + model_tp, model, device = _load_tp_and_reference_models(model_path, model_class, enable_sequence_parallel=True) + tp_plan = _get_active_tp_plan(model_tp) model_tp.train() model.train() @@ -203,32 +237,35 @@ def _test_tp_backward_impl(rank, model_path, model_class, atol, rtol): set_seed(0) labels = torch.randint(0, vocab_size, (2, 64)).to(device) - loss = model(input_ids, labels=labels).loss + loss = model(input_ids, labels=labels, use_cache=False).loss loss.backward() - loss_tp = model_tp(input_ids, labels=labels).loss + loss_tp = model_tp(input_ids, labels=labels, use_cache=False).loss loss_tp.backward() - assert torch.allclose(loss, loss_tp, atol=atol, rtol=rtol), ( + loss_tp_local = _to_local(loss_tp) + assert torch.allclose(loss, loss_tp_local, atol=atol, rtol=rtol), ( f"TP and non-TP model losses differ. " - f"Non-TP loss: {loss.item()}, TP loss: {loss_tp.item()}, " - f"Diff: {(loss - loss_tp).abs().item()}" + f"Non-TP loss: {loss.item()}, TP loss: {loss_tp_local.item()}, " + f"Diff: {(loss - loss_tp_local).abs().item()}" ) # Compare gradients for matching parameters world_size = dist.get_world_size() + + # Debug: check tied weights and parameter alignment failed_grads = {} - for (name, param), (_, param_tp) in zip(model.named_parameters(), model_tp.named_parameters()): + for (name, param), (name_tp, param_tp) in zip(model.named_parameters(), model_tp.named_parameters()): if param.grad is not None and param_tp.grad is not None: grad = param.grad - grad_tp = param_tp.grad + grad_tp = _to_local(param_tp.grad) # Slice reference gradient to match local shard if parameter is sharded if grad.shape != grad_tp.shape: for dim in range(grad.ndim): if grad.size(dim) != grad_tp.size(dim): - param_plan = _get_parameter_tp_plan(name, model_tp.tp_plan, is_weight=True) - if param_plan in ("packed_colwise",): + param_plan = _get_parameter_tp_plan(name, tp_plan, is_weight=True) + if param_plan == "packed_colwise": # interleaved slicing grad = get_packed_grad_shard(grad, world_size, rank, dim) else: @@ -239,10 +276,15 @@ def _test_tp_backward_impl(rank, model_path, model_class, atol, rtol): break if not torch.allclose(grad.cpu(), grad_tp.cpu(), atol=atol, rtol=rtol): - failed_grads[name] = (grad.cpu() - grad_tp.cpu()).abs().max().item() + max_diff = (grad.cpu() - grad_tp.cpu()).abs().max().item() + ref_abs_max = grad.cpu().abs().max().item() + tp_abs_max = grad_tp.cpu().abs().max().item() + ratio = tp_abs_max / ref_abs_max if ref_abs_max > 0 else float("inf") + failed_grads[name] = (max_diff, ref_abs_max, tp_abs_max, ratio) assert not failed_grads, f"Gradients differ for {len(failed_grads)} parameter(s):\n" + "\n".join( - f" {name}: max diff = {diff}" for name, diff in failed_grads.items() + f" {name}: max_diff={v[0]:.6f}, ref_max={v[1]:.6f}, tp_max={v[2]:.6f}, tp/ref ratio={v[3]:.4f}" + for name, v in failed_grads.items() ) dist.barrier() @@ -274,7 +316,7 @@ def _test_tp_generation_impl(_rank, model_path, model_class, atol, rtol, max_new # Compare logits/scores at each generation step scores = torch.stack(output.scores) - scores_tp = torch.stack(output_tp.scores) + scores_tp = torch.stack([_to_local(s) for s in output_tp.scores]) diff = (scores - scores_tp).abs() assert torch.allclose(scores, scores_tp, atol=atol, rtol=rtol), ( @@ -283,9 +325,10 @@ def _test_tp_generation_impl(_rank, model_path, model_class, atol, rtol, max_new ) # Compare generated token sequences - assert torch.equal(output.sequences, output_tp.sequences), ( + sequences_tp = _to_local(output_tp.sequences) + assert torch.equal(output.sequences, sequences_tp), ( f"TP and non-TP model generated different token sequences (direct load path). " - f"Non-TP: {output.sequences.tolist()} | TP: {output_tp.sequences.tolist()}" + f"Non-TP: {output.sequences.tolist()} | TP: {sequences_tp.tolist()}" ) dist.barrier() @@ -297,7 +340,11 @@ def _test_tp_generation_quantized_impl(_rank, model_path, model_class, max_new_t quantization_config = TorchAoConfig(Float8WeightOnlyConfig()) - model_tp = model_class.from_pretrained(model_path, tp_plan="auto", quantization_config=quantization_config) + model_tp = model_class.from_pretrained( + model_path, + distributed_config=DistributedConfig(tp_size=dist.get_world_size()), + quantization_config=quantization_config, + ) dist.barrier() device = model_tp.device @@ -347,9 +394,13 @@ def _test_tp_generation_quantized_impl(_rank, model_path, model_class, max_new_t def _load_ep_and_reference_models(model_path, model_class): """Load EP model and non-EP reference model for comparison.""" + tp_size = dist.get_world_size() model_ep = model_class.from_pretrained( model_path, - distributed_config=DistributedConfig(enable_expert_parallel=True), + distributed_config=DistributedConfig( + tp_size=tp_size, + enable_expert_parallel=True, + ), ) dist.barrier() @@ -426,8 +477,8 @@ class TensorParallelTesterMixin(ABC): # Configuration (can be overridden per model) # ============================================================ tensor_parallel_size: int = 2 - tensor_parallel_atol: float = 1e-5 - tensor_parallel_rtol: float = 1e-5 + tensor_parallel_atol: float = 5e-3 + tensor_parallel_rtol: float = 5e-3 @property @abstractmethod @@ -562,18 +613,7 @@ def test_tp_generation_quantized(self): if not is_torchao_available(): self.skipTest("Test requires torchao") - config = self.model_tester.get_config() - model_class = self._get_tp_model_class() - max_new_tokens = 25 - - with tempfile.TemporaryDirectory() as tmp_dir: - set_seed(42) - model = model_class(config) - model.save_pretrained(tmp_dir, save_original_format=True) - - _init_distributed(tp=self.tensor_parallel_size)(_test_tp_generation_quantized_impl)( - tmp_dir, model_class, max_new_tokens - ) + self.skipTest("Quantization is not currently supported with distributed training") @is_tensor_parallel_test def test_ep_forward(self): diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py index d35d1f6ad9b8..0253e4e50861 100644 --- a/tests/utils/test_core_model_loading.py +++ b/tests/utils/test_core_model_loading.py @@ -17,6 +17,7 @@ import torch import torch.nn as nn +from torch.distributed.tensor.placement_types import Shard from transformers import PretrainedConfig, PreTrainedModel from transformers.conversion_mapping import ( @@ -39,11 +40,13 @@ convert_and_load_state_dict_in_model, rename_source_key, revert_weight_conversion, + spawn_materialize, ) from transformers.modeling_utils import LoadStateDictConfig from transformers.utils.import_utils import is_triton_available from ..test_modeling_common import compare_state_dicts +from .test_distributed_sharding_utils import FakeMesh, _make_dtensor_shard_op class TestWeightGlobMatching(unittest.TestCase): @@ -225,6 +228,99 @@ def __init__(self, config, add_extra_moe=False, with_mlp=True): class TestConvertAndLoadStateDict(unittest.TestCase): + def test_dtensor_shard_aware_mixtral_conversion_uses_only_local_experts(self): + """Integration test: FSDP-sharded expert loading + WeightConverter. + + The problem: Mixtral has 8 experts. The checkpoint stores them separately:: + + experts.0.w1.weight (2x2) + experts.0.w3.weight (2x2) + experts.1.w1.weight (2x2) + experts.1.w3.weight (2x2) + + The model stores them packed into one tensor:: + + experts.gate_up_proj.weight (2, 4, 2) + ^ ^ ^ + | | +-- features + | +-- w1 (2) + w3 (2) concatenated + +-- num_experts + + The conversion (without FSDP) is: load all expert w1/w3 tensors, + MergeModulelist(dim=0) stacks experts, Concatenate(dim=1) joins w1+w3. + + With FSDP, Shard(0) splits the expert dim across ranks. Rank 0 owns + expert 0, rank 1 owns expert 1. So rank 0 should skip loading expert 1 + entirely -- not load it then discard it. + + What the test checks:: + + checkpoint files shard_tensor rank 0 gets + ---------------- ------------ ----------- + experts.0.w1 [[0,1],[2,3]] idx=0 -> kept [[0,1],[2,3]] + experts.1.w1 [[10,11],...] idx=1 -> None (not owned) + experts.0.w3 [[4,5],[6,7]] idx=0 -> kept [[4,5],[6,7]] + experts.1.w3 [[14,15],...] idx=1 -> None (not owned) + + WeightConverter then combines only the kept tensors:: + + MergeModulelist(dim=0): stack owned experts -> shape (1, 2, 2) each + Concatenate(dim=1): cat w1 + w3 along dim 1 + + gate_up_proj = [[[0,1],[2,3],[4,5],[6,7]]] shape (1, 4, 2) + ~~~~~~~~~~ ~~~~~~~~~~ + w1 w3 + + The key point: DtensorShardOperation.shard_tensor(tensor_idx=1) returns + None for rank 0, so the converter never even processes expert 1's data. + This saves memory during loading. + """ + shard_op = _make_dtensor_shard_op( + FakeMesh(shape=(2,), rank=0), + [Shard(0)], + param_shape=(2, 4, 2), + local_shape=(1, 4, 2), + ) + converter = WeightConverter( + ["experts.*.w1.weight", "experts.*.w3.weight"], + "experts.gate_up_proj.weight", + operations=[MergeModulelist(dim=0), Concatenate(dim=1)], + ) + + for idx, tensor in enumerate( + [ + torch.tensor([[0.0, 1.0], [2.0, 3.0]]), + torch.tensor([[10.0, 11.0], [12.0, 13.0]]), + ] + ): + converter.add_tensor( + "model.layers.0.experts.gate_up_proj.weight", + f"model.layers.0.experts.{idx}.w1.weight", + "experts.*.w1.weight", + spawn_materialize(None, tensor, device="cpu", dtype=None, sharding_op=shard_op, tensor_idx=idx), + ) + + for idx, tensor in enumerate( + [ + torch.tensor([[4.0, 5.0], [6.0, 7.0]]), + torch.tensor([[14.0, 15.0], [16.0, 17.0]]), + ] + ): + converter.add_tensor( + "model.layers.0.experts.gate_up_proj.weight", + f"model.layers.0.experts.{idx}.w3.weight", + "experts.*.w3.weight", + spawn_materialize(None, tensor, device="cpu", dtype=None, sharding_op=shard_op, tensor_idx=idx), + ) + + converted = converter.convert("model.layers.0.experts.gate_up_proj.weight") + + self.assertEqual(list(converted), ["model.layers.0.experts.gate_up_proj.weight"]) + torch.testing.assert_close( + converted["model.layers.0.experts.gate_up_proj.weight"], + torch.tensor([[[0.0, 1.0], [2.0, 3.0], [4.0, 5.0], [6.0, 7.0]]]), + ) + def test_moe_and_qkv_conversion(self): model = DummyRoot(PretrainedConfig()) @@ -509,11 +605,11 @@ def __init__(self, config): model_state = model.state_dict() self.assertFalse(torch.allclose(raw_k, expected_k)) - torch.testing.assert_close(model_state["model.layers.0.self_attn.k_proj.weight"], expected_k) - torch.testing.assert_close(model_state["model.layers.0.self_attn.v_proj.weight"], expected_v) + torch.testing.assert_close(model_state["layers.0.self_attn.k_proj.weight"], expected_k) + torch.testing.assert_close(model_state["layers.0.self_attn.v_proj.weight"], expected_v) - q_weight_key = "model.layers.0.self_attn.q_proj.weight" - scale_key = "model.layers.0.self_attn.q_proj.weight_scale_inv" + q_weight_key = "layers.0.self_attn.q_proj.weight" + scale_key = "layers.0.self_attn.q_proj.weight_scale_inv" self.assertIn(scale_key, model_state) expected_dtype = torch.float8_e4m3fn if hasattr(torch, "float8_e4m3fn") else torch.int8 self.assertEqual(model_state[q_weight_key].dtype, expected_dtype) @@ -524,11 +620,14 @@ def __init__(self, config): torch.Size((out_dim // block_size[0], in_dim // block_size[1])), ) - dequant = Fp8Dequantize(block_size=block_size) + dequant = Fp8Dequantize(quantizer) dequantized_q = dequant.convert( - [model_state[q_weight_key], model_state[scale_key]], - context={"quantization_config": quantizer.quantization_config}, - ) + { + "weight$": [model_state[q_weight_key]], + "weight_scale_inv": [model_state[scale_key]], + }, + full_layer_name=q_weight_key, + )[q_weight_key] torch.testing.assert_close(dequantized_q, expected_q, rtol=1e-2, atol=1e-2) def test_scoped_renaming_does_not_leak_to_sibling_or_parent(self): diff --git a/tests/utils/test_distributed_sharding_utils.py b/tests/utils/test_distributed_sharding_utils.py new file mode 100644 index 000000000000..677f13ca2638 --- /dev/null +++ b/tests/utils/test_distributed_sharding_utils.py @@ -0,0 +1,462 @@ +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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 os +import shutil +import tempfile +import unittest + +import torch +from torch.distributed.tensor.placement_types import Replicate, Shard, _StridedShard + +from transformers.distributed.sharding_utils import ( + DtensorShardOperation, + _find_strided_shard_placement_from_fused_params, +) + + +if torch.distributed.is_available(): + import torch.distributed as dist + import torch.distributed.checkpoint as dcp + import torch.multiprocessing as mp + from torch.distributed.device_mesh import init_device_mesh + from torch.distributed.tensor import DTensor, distribute_tensor + + from transformers.distributed.sharding_utils import ( + fuse_optimizer_state, + get_fusion_metadata, + unfuse_optimizer_state, + ) + + +class FakeMesh: + """Fake multi-dimensional device mesh for testing DtensorShardOperation.""" + + def __init__(self, shape, rank, dim_names=None): + if isinstance(shape, int): + shape = (shape,) + self.shape = tuple(shape) + self.ndim = len(self.shape) + self.mesh_dim_names = dim_names or tuple(f"dim{i}" for i in range(self.ndim)) + # Compute nD coordinate (row-major: last dim changes fastest) + self._coord = [] + r = rank + for s in reversed(self.shape): + self._coord.insert(0, r % s) + r //= s + + def get_local_rank(self): + return self._coord[0] + + def get_coordinate(self): + return tuple(self._coord) + + def size(self): + result = 1 + for s in self.shape: + result *= s + return result + + def _is_current_rank_part_of_mesh(self): + return True + + def _sym_get_coordinate(self, dim): + return self._coord[dim] + + def __getitem__(self, name): + idx = self.mesh_dim_names.index(name) + return FakeMesh( + shape=(self.shape[idx],), + rank=self._coord[idx], + dim_names=(name,), + ) + + +def _make_dtensor_shard_op(mesh, placements, param_shape, local_shape): + """Build a DtensorShardOperation without requiring a real DTensor / distributed init. + + The axis-0 ownership cache is computed by mimicking + ``compute_local_shape_and_global_offset`` for the leading dim only: + locate the mesh dim that shards param dim 0 (if any) and use its local rank. + """ + op = object.__new__(DtensorShardOperation) + op.device_mesh = mesh + op.placements = tuple(placements) + op.param_ndim = len(param_shape) + op._axis0_offset = 0 + op._axis0_local_size = local_shape[0] + for mesh_dim, p in enumerate(placements): + if hasattr(p, "dim") and (p.dim % len(param_shape)) == 0: + sub = mesh[mesh.mesh_dim_names[mesh_dim]] if mesh.ndim > 1 else mesh + op._axis0_offset = sub.get_local_rank() * local_shape[0] + break + return op + + +class TestDtensorShardOperation(unittest.TestCase): + """Unit tests for DtensorShardOperation. + + See `DtensorShardOperation` in sharding_utils.py for the placement primer + and table of checkpoint layouts. The rest of this docstring covers the + test-specific conventions you need to write new cases here. + + Running example used throughout these tests: a stack of N MoE experts, + each of shape [in, out]. The full parameter shape is [N, in, out]. + Tests are parameterized so every rank in a (fake) mesh is checked. + + The checkpoint can store this param in two layouts, and shard_tensor + behaves differently for each: + + | Layout | tensor_idx | source.shape | + |------------------------------|-------------------|----------------| + | Single stacked tensor | None | [N, in, out] | + | N separate per-expert files | 0, 1, ..., N-1 | [in, out] | + + Single-tensor case: source has the full param shape, including axis 0. + shard_tensor returns this rank's slice along every sharded dim. + + Per-piece case: shard_tensor is called once per expert. Each call's + source is just that one expert's [in, out] tensor — note source has + one fewer dim than the param, because the axis-0 index lives in + `tensor_idx`, not in source.shape. shard_tensor returns: + - None, if this rank doesn't own expert `tensor_idx` along axis 0 + (the piece is then dropped by the caller, MergeModulelist / + Concatenate), + - the inner-dim slice otherwise. After all N calls, the caller + stacks the surviving slices along axis 0 to rebuild this rank's + local [n_local, in, out]. + """ + + def test_no_shard_placements_returns_full_copy(self): + tensor = torch.arange(16).reshape(4, 4).float() + expected = { + 0: tensor, # rank 0 — no shards, full copy + 1: tensor, # rank 1 — no shards, full copy + } + for rank in range(2): + mesh = FakeMesh(shape=(2,), rank=rank) + op = _make_dtensor_shard_op(mesh, [Replicate()], param_shape=(4, 4), local_shape=(4, 4)) + torch.testing.assert_close(op.shard_tensor(tensor), expected[rank], msg=f"rank {rank}") + + def test_1D_shard(self): + tensor = torch.arange(16).reshape(4, 4).float() + expected = { + 0: tensor[:2], # rank 0 — first half + 1: tensor[2:], # rank 1 — second half + } + for rank in range(2): + mesh = FakeMesh(shape=(2,), rank=rank) + op = _make_dtensor_shard_op(mesh, [Shard(0)], param_shape=(4, 4), local_shape=(2, 4)) + torch.testing.assert_close(op.shard_tensor(tensor), expected[rank], msg=f"rank {rank}") + + def test_1D_strided_shard(self): + tensor = torch.arange(16).reshape(4, 4).float() + expected = { + 0: tensor[[0, 2]], # first piece of each group — rows {0, 2} + 1: tensor[[1, 3]], # second piece of each group — rows {1, 3} + } + for rank in range(2): + mesh = FakeMesh(shape=(2,), rank=rank) + op = _make_dtensor_shard_op( + mesh, [_StridedShard(dim=0, split_factor=2)], param_shape=(4, 4), local_shape=(2, 4) + ) + torch.testing.assert_close(op.shard_tensor(tensor), expected[rank], msg=f"rank {rank}") + + def test_2D_shard_different_dims(self): + tensor = torch.arange(64).reshape(8, 8).float() + expected = { + 0: tensor[:4, :4], # top-left + 1: tensor[:4, 4:], # top-right + 2: tensor[4:, :4], # bottom-left + 3: tensor[4:, 4:], # bottom-right + } + for rank in range(4): + mesh = FakeMesh(shape=(2, 2), rank=rank) + op = _make_dtensor_shard_op(mesh, [Shard(0), Shard(1)], param_shape=(8, 8), local_shape=(4, 4)) + torch.testing.assert_close(op.shard_tensor(tensor), expected[rank], msg=f"rank {rank}") + + def test_2D_shard_same_dim(self): + tensor = torch.arange(64).reshape(8, 8).float() + expected = { + 0: tensor[:2], # rows 0-1 + 1: tensor[2:4], # rows 2-3 + 2: tensor[4:6], # rows 4-5 + 3: tensor[6:8], # rows 6-7 + } + for rank in range(4): + mesh = FakeMesh(shape=(2, 2), rank=rank) + op = _make_dtensor_shard_op(mesh, [Shard(0), Shard(0)], param_shape=(8, 8), local_shape=(2, 8)) + torch.testing.assert_close(op.shard_tensor(tensor), expected[rank], msg=f"rank {rank}") + + def test_2D_strided_shard_same_dim(self): + tensor = torch.arange(16).reshape(4, 4).float() + expected = { + 0: tensor[[0]], # row 0 + 1: tensor[[2]], # row 2 + 2: tensor[[1]], # row 1 + 3: tensor[[3]], # row 3 + } + for rank in range(4): + mesh = FakeMesh(shape=(2, 2), rank=rank) + op = _make_dtensor_shard_op( + mesh, + [_StridedShard(dim=0, split_factor=2), Shard(0)], + param_shape=(4, 4), + local_shape=(1, 4), + ) + torch.testing.assert_close(op.shard_tensor(tensor), expected[rank], msg=f"rank {rank}") + + def test_2D_strided_shard_different_dims(self): + tensor = torch.arange(16).reshape(4, 4).float() + expected = { + 0: torch.cat([tensor[:2, 0:1], tensor[:2, 2:3]], dim=1), # top rows, cols {0, 2} + 1: torch.cat([tensor[:2, 1:2], tensor[:2, 3:4]], dim=1), # top rows, cols {1, 3} + 2: torch.cat([tensor[2:, 0:1], tensor[2:, 2:3]], dim=1), # bottom rows, cols {0, 2} + 3: torch.cat([tensor[2:, 1:2], tensor[2:, 3:4]], dim=1), # bottom rows, cols {1, 3} + } + for rank in range(4): + mesh = FakeMesh(shape=(2, 2), rank=rank) + op = _make_dtensor_shard_op( + mesh, + [Shard(0), _StridedShard(dim=1, split_factor=2)], + param_shape=(4, 4), + local_shape=(2, 2), + ) + torch.testing.assert_close(op.shard_tensor(tensor), expected[rank], msg=f"rank {rank}") + + def test_moe_1D_shard_filters_by_axis0_ownership(self): + source = torch.ones(2, 2) + expected = { + 0: { + 0: source, # first owned + 1: source, # last owned + 2: None, # first not-owned (upper boundary, exclusive) + 3: None, # not owned + }, + 1: { + 0: None, # not owned + 1: None, # last not-owned (just below offset) + 2: source, # first owned (lower boundary, inclusive) + 3: source, # last owned + }, + } + for rank in range(2): + mesh = FakeMesh(shape=(2,), rank=rank) + op = _make_dtensor_shard_op(mesh, [Shard(0)], param_shape=(4, 2, 2), local_shape=(2, 2, 2)) + for tensor_idx, exp in expected[rank].items(): + with self.subTest(rank=rank, tensor_idx=tensor_idx): + shard = op.shard_tensor(source, tensor_idx=tensor_idx) + if exp is None: + self.assertIsNone(shard) + else: + torch.testing.assert_close(shard, exp) + + def test_moe_1D_strided_shard_on_inner_dim_degrades_to_contiguous(self): + source = torch.arange(8).reshape(4, 2).float() + expected = { + 0: source[:2], # rank 0 — first half (strided silently degraded to contiguous) + 1: source[2:], # rank 1 — second half + } + for rank in range(2): + mesh = FakeMesh(shape=(2,), rank=rank) + op = _make_dtensor_shard_op( + mesh, + [_StridedShard(dim=1, split_factor=2)], + param_shape=(8, 8, 2), + local_shape=(8, 4, 2), + ) + torch.testing.assert_close(op.shard_tensor(source, tensor_idx=0), expected[rank], msg=f"rank {rank}") + + def test_moe_2D_shard_on_axis0_and_inner_dim_slices_inner(self): + source = torch.arange(8).reshape(4, 2).float() + expected = { + 0: source[:2], # owned; inner Shard(1) → source dim 0 first half + 1: source[2:], # owned; inner Shard(1) → source dim 0 second half + 2: None, # not owned (axis-0 ownership filter) + 3: None, # not owned + } + for rank in range(4): + mesh = FakeMesh(shape=(2, 2), rank=rank) + op = _make_dtensor_shard_op(mesh, [Shard(0), Shard(1)], param_shape=(4, 4, 2), local_shape=(2, 2, 2)) + shard = op.shard_tensor(source, tensor_idx=1) + if expected[rank] is None: + self.assertIsNone(shard, msg=f"rank {rank}") + else: + torch.testing.assert_close(shard, expected[rank], msg=f"rank {rank}") + + def test_moe_2D_shard_with_negative_dim_indices(self): + source = torch.arange(8).reshape(4, 2).float() + expected = { + 0: source[:, :1], # owned; inner Shard(-1) → source dim 1, first half + 1: source[:, 1:], # owned; inner Shard(-1) → source dim 1, second half + 2: None, # not owned + 3: None, # not owned + } + for rank in range(4): + mesh = FakeMesh(shape=(2, 2), rank=rank) + op = _make_dtensor_shard_op(mesh, [Shard(-3), Shard(-1)], param_shape=(4, 4, 2), local_shape=(2, 4, 1)) + shard = op.shard_tensor(source, tensor_idx=1) + if expected[rank] is None: + self.assertIsNone(shard, msg=f"rank {rank}") + else: + torch.testing.assert_close(shard, expected[rank], msg=f"rank {rank}") + + def test_strided_intervals(self): + # Direct tests for _strided_intervals(intervals, rank, world_size, split_factor). + # Keys: (input_interval, rank, world_size, split_factor) -> expected output list. + mesh = FakeMesh(shape=(2,), rank=0) + op = _make_dtensor_shard_op(mesh, [Shard(0)], param_shape=(8,), local_shape=(4,)) + expected = { + # Even (0, 8) sf=2 ws=2 -> groups (0,4) (4,8); each rank takes half of each + ((0, 8), 0, 2, 2): [(0, 2), (4, 6)], + ((0, 8), 1, 2, 2): [(2, 4), (6, 8)], + # Uneven (0, 7) sf=2 -> group 0 = (0,4), group 1 = (4,7) -> size 3 + ((0, 7), 0, 2, 2): [(0, 2), (4, 6)], # rank 0 -> half of each group + ((0, 7), 1, 2, 2): [(2, 4), (6, 7)], # rank 1's piece in group 1 is 1 elem wide + # split_factor=1 collapses to a single group -> contiguous behavior + ((0, 4), 0, 2, 1): [(0, 2)], + # split_factor=4 on size-2: groups (0,1), (1,2), (2,2), (3,2) -> last 2 empty -> skipped + ((0, 2), 0, 2, 4): [(0, 1), (1, 2)], + } + for (interval, rank, ws, sf), exp in expected.items(): + with self.subTest(interval=interval, rank=rank, ws=ws, sf=sf): + self.assertEqual(op._strided_intervals([interval], rank=rank, world_size=ws, split_factor=sf), exp) + + def test_contiguous_intervals(self): + # Direct tests for _contiguous_intervals(intervals, rank, world_size). + # Keys: (input_intervals, rank, world_size) -> expected output list. + mesh = FakeMesh(shape=(2,), rank=0) + op = _make_dtensor_shard_op(mesh, [Shard(0)], param_shape=(8,), local_shape=(4,)) + expected = { + # Single interval, even split -> rank 0 -> first 4 elems, rank 1 -> last 4 elems + (((0, 8),), 0, 2): [(0, 4)], + (((0, 8),), 1, 2): [(4, 8)], + # Uneven: size 5 / 2 ranks -> rank 0 -> first 3 elems, rank 1 -> last 2 elems + (((0, 5),), 0, 2): [(0, 3)], + (((0, 5),), 1, 2): [(3, 5)], + # Empty: size 3 / 4 ranks -> rank 3 -> nothing + (((0, 3),), 3, 4): [], + # Multi-input intervals (8 elems): rank -> takes its slice from whichever interval(s) cover it + (((0, 4), (10, 14)), 0, 2): [(0, 4)], # rank 0 -> first 4 elems = first interval entirely + (((0, 4), (10, 14)), 1, 2): [(10, 14)], # rank 1 -> last 4 elems = second interval entirely + (((0, 4), (10, 14)), 2, 4): [(10, 12)], # ws=4, rank 2 -> cuts mid-input-interval + } + for (intervals, rank, ws), exp in expected.items(): + with self.subTest(intervals=intervals, rank=rank, ws=ws): + self.assertEqual(op._contiguous_intervals(list(intervals), rank=rank, world_size=ws), exp) + + def test_slice_and_cat(self): + # Direct tests for _slice_and_cat(source, intervals, device, dtype). + tensor = torch.arange(64).reshape(8, 8).float() + mesh = FakeMesh(shape=(2,), rank=0) + op = _make_dtensor_shard_op(mesh, [Shard(0)], param_shape=(8, 8), local_shape=(4, 8)) + expected = { + # Fast path: every dim is single-interval -> one slice read, no cat + "fast_path": ([[(0, 4)], [(0, 8)]], tensor[:4, :]), + # Cat along dim 1: two disjoint col ranges -> read separately and concat + "cat_dim1": ([[(0, 4)], [(0, 2), (4, 6)]], torch.cat([tensor[:4, :2], tensor[:4, 4:6]], dim=1)), + # Cat along dim 0: two disjoint row ranges -> read separately and concat + "cat_dim0": ([[(0, 2), (4, 6)], [(0, 8)]], torch.cat([tensor[:2, :], tensor[4:6, :]], dim=0)), + } + for case, (intervals, exp) in expected.items(): + with self.subTest(case=case): + torch.testing.assert_close(op._slice_and_cat(tensor, intervals, None, None), exp) + + # Reject: two dims with disjoint ranges -> would require an outer-product of reads. + with self.assertRaises(ValueError): + op._slice_and_cat(tensor, [[(0, 2), (4, 6)], [(0, 2), (4, 6)]], None, None) + + result = op._slice_and_cat(tensor, [[(0, 4)], [(0, 8)]], None, torch.float16) + self.assertEqual(result.dtype, torch.float16) + + +class TestFindStridedShardPlacementFromFusedParams(unittest.TestCase): + def test_find_strided_shard_placement_from_fused_params(self): + expected = { + # Plain placements — no _StridedShard, nothing to do + (Replicate(),): None, + (Shard(0),): None, + (Shard(0), Shard(2)): None, + # Uncomposed _StridedShard — TP-only fused gate||up (DCP can't encode) + (_StridedShard(0, split_factor=2),): _StridedShard(0, split_factor=2), + (_StridedShard(1, split_factor=4),): _StridedShard(1, split_factor=4), + # _StridedShard on a different tensor dim than the other Shard — still uncomposed + (Shard(0), _StridedShard(1, split_factor=2)): _StridedShard(1, split_factor=2), + (_StridedShard(2, split_factor=2), Shard(0)): _StridedShard(2, split_factor=2), + # _StridedShard composed with another Shard on the SAME tensor dim — DCP-friendly + (_StridedShard(0, split_factor=2), Shard(0)): None, + (Shard(0), _StridedShard(0, split_factor=2)): None, + } + for placements, exp in expected.items(): + with self.subTest(placements=placements): + self.assertEqual(_find_strided_shard_placement_from_fused_params(placements), exp) + + +def _optimizer_state_checkpointing_e2e_worker(rank, world_size, port, ckpt_dir): + # 1. Init a 4-rank CPU process group + 2x2 (fsdp, tp) mesh. + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(port) + dist.init_process_group("gloo", rank=rank, world_size=world_size) + mesh = init_device_mesh("cpu", (2, 2), mesh_dim_names=("fsdp", "tp")) + + # 2. Build a fused DTensor with the Mixtral gate_up_proj placement: + # shape (num_experts=4, 2·intermediate=16, hidden=8), + # placements (Shard(0), _StridedShard(1, sf=2)). + full = torch.arange(4 * 16 * 8, dtype=torch.float32).reshape(4, 16, 8) + dt = distribute_tensor(full, mesh, [Shard(0), Shard(1)]) + dt = DTensor.from_local( + dt._local_tensor.clone(), + mesh, + (Shard(0), _StridedShard(1, split_factor=2)), + run_check=False, + ) + + # 3. Wrap it the way `get_optimizer_state_dict` would. + fqn = "model.layers.0.mlp.experts.gate_up_proj" + osd = { + "state": {fqn: {"exp_avg": dt, "step": torch.tensor(7.0)}}, + "param_groups": [{"lr": 1e-4}], + } + + # 4. Snapshot the rank-local buffer for later bit-exact comparison. + before = dt._local_tensor.clone() + + # 5. unfuse → DCP save → DCP load → fuse. + fusion_metadata = get_fusion_metadata(osd) + assert set(fusion_metadata) == {fqn}, f"expected one fused param, got {set(fusion_metadata)}" + unfuse_optimizer_state(osd, fusion_metadata) + dcp.save({"optimizer": osd}, checkpoint_id=ckpt_dir) + dist.barrier() + dcp.load({"optimizer": osd}, checkpoint_id=ckpt_dir) + fuse_optimizer_state(osd, fusion_metadata) + + # 6. Verify the placement is restored and the rank-local data is bit-exact. + after = osd["state"][fqn]["exp_avg"] + assert tuple(after.placements) == (Shard(0), _StridedShard(1, split_factor=2)), after.placements + assert torch.equal(after._local_tensor, before), f"rank {rank}: local data drifted after round-trip" + + dist.destroy_process_group() + + +@unittest.skipUnless(torch.distributed.is_available(), "Requires torch.distributed (gloo backend).") +class TestOptimizerStateCheckpointing(unittest.TestCase): + def test_optimizer_state_checkpointing_e2e(self): + tmp = tempfile.mkdtemp(prefix="hf_optimizer_state_checkpointing_") + try: + mp.spawn(_optimizer_state_checkpointing_e2e_worker, args=(4, 29500, tmp), nprocs=4, join=True) + finally: + shutil.rmtree(tmp, ignore_errors=True) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/utils/test_modeling_utils.py b/tests/utils/test_modeling_utils.py index fab48f9ddb8a..ba30b775588c 100644 --- a/tests/utils/test_modeling_utils.py +++ b/tests/utils/test_modeling_utils.py @@ -59,6 +59,7 @@ is_torch_available, logging, ) +from transformers.distributed import DistributedConfig from transformers.modeling_flash_attention_utils import is_flash_attn_available from transformers.models.mistral.modeling_mistral import MistralModel from transformers.testing_utils import ( @@ -431,6 +432,42 @@ def test_get_total_byte_count_does_not_require_process_group(self): self.assertIn(torch.device("cpu"), total_byte_count) self.assertGreater(total_byte_count[torch.device("cpu")], 0) + def test_model_from_pretrained_fsdp_distributes_before_loading(self): + model = GPT2LMHeadModel(GPT2Config(n_layer=1, n_head=2, n_embd=8, n_positions=8, n_ctx=8, vocab_size=32)) + + with tempfile.TemporaryDirectory() as tmp_dir: + model.save_pretrained(tmp_dir) + call_order = [] + + fake_mesh = mock.Mock() + fake_mesh.mesh_dim_names = ("fsdp",) + fake_mesh.ndim = 1 + + def fake_apply_fsdp(model, fsdp_mesh, fsdp_plan): + call_order.append("distribute") + self.assertIs(fsdp_mesh, fake_mesh) + self.assertIsNone(fsdp_plan) + return model + + def fake_load_pretrained_model(model, state_dict, checkpoint_files, load_config, expected_keys=None): + call_order.append("load") + self.assertIs(load_config.device_mesh, fake_mesh) + return mock.Mock(), None + + with ( + patch("transformers.modeling_utils.init_device_mesh", return_value=fake_mesh), + patch("transformers.distributed.utils.apply_fully_shard_data_parallel", side_effect=fake_apply_fsdp), + patch.object(GPT2LMHeadModel, "_load_pretrained_model", side_effect=fake_load_pretrained_model), + patch.object( + GPT2LMHeadModel, + "_finalize_model_loading", + side_effect=lambda model, load_config, loading_info: loading_info, + ), + ): + GPT2LMHeadModel.from_pretrained(tmp_dir, distributed_config=DistributedConfig(fsdp_size=2)) + + self.assertEqual(call_order, ["distribute", "load"]) + def test_hub_retry(self): @hub_retry(max_attempts=2) def test_func(): diff --git a/tmp.py b/tmp.py new file mode 100644 index 000000000000..05db1b840b44 --- /dev/null +++ b/tmp.py @@ -0,0 +1 @@ +"&&" diff --git a/utils/tests_fetcher.py b/utils/tests_fetcher.py index a138ef2eaacb..f09480a59875 100644 --- a/utils/tests_fetcher.py +++ b/utils/tests_fetcher.py @@ -1099,6 +1099,7 @@ def parse_commit_message(commit_message: str) -> dict[str, bool]: "tests_non_model": r"tests/[^/]*?/test_.*\.py", "tests_training_ci": r"tests/models/.*/test_modeling_.*", "tests_tensor_parallel_ci": r"(tests/models/.*/test_modeling_.*|tests/tensor_parallel(?:/test_tensor_parallel\.py)?)", + "tests_fsdp_ci": r"(tests/models/.*/test_modeling_.*|tests/test_fsdp_mixin\.py)", }