diff --git a/.circleci/create_circleci_config.py b/.circleci/create_circleci_config.py index de0646af2dbf..8e87388f622f 100644 --- a/.circleci/create_circleci_config.py +++ b/.circleci/create_circleci_config.py @@ -404,6 +404,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.) @@ -435,7 +444,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/.gitignore b/.gitignore index 903efc854eef..de7ec2b0ac2b 100644 --- a/.gitignore +++ b/.gitignore @@ -7,7 +7,7 @@ __pycache__/ # C extensions *.so - +checkpoints/ # tests and logs tests/fixtures/cached_*_text.txt logs/ diff --git a/src/transformers/__init__.py b/src/transformers/__init__.py index 71d416151f8c..29b23e1ee423 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/cli/serving/utils.py b/src/transformers/cli/serving/utils.py index 68b13b280454..222f7b67570a 100644 --- a/src/transformers/cli/serving/utils.py +++ b/src/transformers/cli/serving/utils.py @@ -491,7 +491,7 @@ def __init__( self._tokenizer = tokenizer self._loop = loop self._queue = queue - self._decode_stream = DecodeStream([], skip_special_tokens) + self._decode_stream: DecodeStream = DecodeStream([], skip_special_tokens) self._stc_id = tool_config["stc_id"] if tool_config else None self._etc_id = tool_config["etc_id"] if tool_config else None self._inside_tool_call = False @@ -523,7 +523,7 @@ def put(self, value: "torch.Tensor") -> None: is_start_or_end_token = _advance_thinking_state(self, token_id) - text = self._decode_stream.step(self._tokenizer, token_id) + text = self._decode_stream.step(self._tokenizer, token_id) # ty:ignore[unresolved-attribute] if text is None or self._inside_tool_call or token_id == self._etc_id or is_start_or_end_token: continue if self._inside_thinking: @@ -575,7 +575,7 @@ def __init__( self._loop = loop self._queue = queue self._tokenizer = tokenizer - self._decode_stream = DecodeStream([], True) + self._decode_stream: DecodeStream = DecodeStream([], True) self._stc_id = tool_config["stc_id"] if tool_config else None self._etc_id = tool_config["etc_id"] if tool_config else None self._inside_tool_call = False @@ -602,7 +602,7 @@ def put(self, output: "GenerationOutput") -> None: is_start_or_end_token = _advance_thinking_state(self, token_id) - text = self._decode_stream.step(self._tokenizer, token_id) + text = self._decode_stream.step(self._tokenizer, token_id) # ty:ignore[unresolved-attribute] if text is None or self._inside_tool_call or token_id == self._etc_id or is_start_or_end_token: continue if self._inside_thinking: diff --git a/src/transformers/configuration_utils.py b/src/transformers/configuration_utils.py index 2b6fdce06e28..4b8ec290c818 100755 --- a/src/transformers/configuration_utils.py +++ b/src/transformers/configuration_utils.py @@ -149,6 +149,16 @@ 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 when `distributed_config.enable_sequence_parallel=True` and + `enable_expert_parallel=False`. Same key/value shape as the TP plan. + - **base_model_tp_ep_plan** (`dict[str, Any]`) -- Complete plan for inference TP + expert parallel + (`enable_sequence_parallel=False`, `enable_expert_parallel=True`). + - **base_model_sp_ep_plan** (`dict[str, Any]`) -- Complete plan for training SP + expert parallel + (`enable_sequence_parallel=True`, `enable_expert_parallel=True`). + - **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. @@ -220,6 +230,10 @@ 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_tp_ep_plan: ClassVar[dict[str, Any] | None] = None + base_model_sp_ep_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 @@ -1022,6 +1036,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] @@ -1167,6 +1184,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 0d57b0b8d23c..da8bbe4e737a 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -30,21 +30,21 @@ 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 _torch_distributed_available = torch.distributed.is_available() +if _torch_distributed_available: + from torch.distributed.tensor import DTensor if TYPE_CHECKING: - from .integrations.tensor_parallel import TensorParallelLayer from .modeling_utils import LoadStateDictConfig, PreTrainedModel from .quantizers import HfQuantizer - logger = get_logger(__name__) @@ -388,7 +388,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) @@ -404,11 +404,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 @@ -610,7 +609,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) @@ -768,7 +767,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 reverse_transform.base_model_prefix = self.base_model_prefix @@ -795,6 +796,8 @@ def materialize_tensors(self) -> dict[str, list[torch.Tensor]]: # Sync loading elif callable(tensors[0]): tensors = [func() for func in tensors] + # Some may be None for some distributed setups + tensors = [tensor for tensor in tensors if tensor is not None] # Add them to the new dictionary collected_tensors[key] = tensors @@ -997,29 +1000,20 @@ 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.""" - - 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 - + """Materialize (and optionally shard) a tensor, asynchronously if a thread pool is provided. -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) @@ -1092,12 +1086,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(".") @@ -1112,27 +1106,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( @@ -1317,7 +1309,6 @@ def convert_and_load_state_dict_in_model( device_map = load_config.device_map or {"": "cpu"} hf_quantizer = load_config.hf_quantizer dtype = load_config.dtype - device_mesh = load_config.device_mesh disk_offload_folder = load_config.disk_offload_folder offload_buffers = load_config.offload_buffers dtype_plan = load_config.dtype_plan or {} @@ -1352,10 +1343,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())) @@ -1419,37 +1406,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() - ) - # Per-expert sharding (EP) needs `tensor_idx` = the expert index so the - # distributed op selects whole experts. The signal is a `MergeModulelist` - # in the chain; it isn't always `operations[0]` (e.g. an FP8 quantizer - # prepends a scale-decode op), so scan the whole chain rather than just the head. - shard_index = ( - len(mapping.collected_tensors.get(source_pattern, [])) - if isinstance(mapping, WeightConverter) - and any(isinstance(op, MergeModulelist) for op in mapping.operations) - 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) @@ -1479,14 +1452,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..8d5570b4abbb 100644 --- a/src/transformers/distributed/__init__.py +++ b/src/transformers/distributed/__init__.py @@ -19,6 +19,15 @@ _import_structure = { "configuration_utils": ["DistributedConfig"], + "fsdp": ["is_fsdp_enabled", "is_fsdp_managed_module", "verify_fsdp_plan"], + "sharding_utils": [], + "tensor_parallel": [ + "ALL_PARALLEL_STYLES", + "apply_tensor_parallel", + "select_parallel_plan", + "verify_tp_plan", + ], + "utils": [], } @@ -26,6 +35,13 @@ from .configuration_utils import ( DistributedConfig, ) + from .fsdp import is_fsdp_enabled, is_fsdp_managed_module, verify_fsdp_plan + from .tensor_parallel import ( + ALL_PARALLEL_STYLES, + apply_tensor_parallel, + select_parallel_plan, + verify_tp_plan, + ) else: import sys diff --git a/src/transformers/distributed/configuration_utils.py b/src/transformers/distributed/configuration_utils.py index 7726d9f3290d..05b63e070a6a 100644 --- a/src/transformers/distributed/configuration_utils.py +++ b/src/transformers/distributed/configuration_utils.py @@ -12,99 +12,83 @@ # 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 select the model's SP/TP and EP plan + (see ``select_parallel_plan``). Set explicitly to override. + enable_sequence_parallel (`bool`, *optional*, defaults to `False`): + Select ``base_model_sp_plan`` (dense-only) or ``base_model_sp_ep_plan`` (with EP). + enable_expert_parallel (`bool`, *optional*, defaults to `False`): + Select ``base_model_tp_ep_plan`` or ``base_model_sp_ep_plan`` when combined with SP flag. + fsdp_size (`int`, *optional*): + Number of devices for FSDP (data parallelism). If `None` and `tp_size` is set, defaults to 1. + fsdp_cpu_offload (`bool`, *optional*, defaults to `False`): + Whether to enable CPU offloading for FSDP2. + fsdp_mixed_precision (`bool`, *optional*, defaults to `False`): + Whether to enable mixed precision for FSDP2. """ + 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_cpu_offload: bool = False + fsdp_mixed_precision: bool = False + + 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() + if self.tp_size * self.fsdp_size != world_size: + raise RuntimeError( + f"tp_size ({self.tp_size}) * fsdp_size ({self.fsdp_size}) is not equal to world_size ({world_size})" + ) + # TODO (remi-or) the check above should probably happen during actual model sharding, same as the check below + # elif self.tp_size > 1 or self.fsdp_size > 1: + # issue = "not initialized" if torch.distributed.is_available() else "not available" + # raise ValueError(f"Requested {self.tp_size = }, {self.fsdp_size = } but torch.distributed is {issue}.") @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..4483a3b5d744 --- /dev/null +++ b/src/transformers/distributed/fsdp.py @@ -0,0 +1,277 @@ +# 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 + +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 + + from .configuration_utils import DistributedConfig + +if is_torch_available(): + import torch + +if is_torch_available() and is_torch_greater_or_equal("2.6"): + from torch.distributed._composable.fsdp import fully_shard + from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy + +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 _get_fsdp_policy_kwargs(distributed_config: DistributedConfig | None) -> dict[str, Any]: + """Build ``fully_shard`` policy kwargs from ``DistributedConfig`` runtime flags.""" + if distributed_config is None: + return {} + + fsdp_policy_kwargs = {} + if distributed_config.fsdp_cpu_offload: + fsdp_policy_kwargs["offload_policy"] = CPUOffloadPolicy() + if distributed_config.fsdp_mixed_precision: + fsdp_policy_kwargs["mp_policy"] = MixedPrecisionPolicy( + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + output_dtype=None, + ) + return fsdp_policy_kwargs + + +def is_norm_and_head_pair(modules: list[tuple[str, Any]]) -> bool: + """Match the canonical tail pair: one final norm + the output head (or tied embedding).""" + if len(modules) != 2: + return False + module_names = [name for name, _ in modules] + has_final_norm = any(name == "norm" or name.endswith(".norm") for name in module_names) + has_output_head = any( + name in {"lm_head", "embed_tokens"} or name.endswith((".lm_head", ".embed_tokens")) for name in module_names + ) + return has_final_norm and has_output_head + + +def _resolve_tied_embed_lm_head_plan( + fsdp_plan: dict[str, str], + *, + tie_word_embeddings: bool, +) -> dict[str, str]: + """ + Rewrite the plan so tied embed/lm_head weights are wrapped once. + Example: + {"model.embed_tokens": "free_full_weight", + "model.layers.*": "free_full_weight", + "model.norm": "keep_full_weight", + "lm_head": "keep_full_weight"} + -> + {"model.layers.*": "free_full_weight", + "model.norm": "keep_full_weight", + "model.embed_tokens": "keep_full_weight"} + """ + if not tie_word_embeddings: + return fsdp_plan + + embed_module_key = next( + (key for key in fsdp_plan if key == "embed_tokens" or key.endswith(".embed_tokens")), + None, + ) + if embed_module_key is None: + return fsdp_plan + + adapted_plan = {} + for key, sharding_strategy in fsdp_plan.items(): + if key == embed_module_key: + continue + if key == "lm_head" and sharding_strategy == "keep_full_weight": + adapted_plan[embed_module_key] = sharding_strategy + else: + adapted_plan[key] = sharding_strategy + return adapted_plan + + +def expand_fsdp_plan(model, fsdp_plan: dict[str, str]) -> list[tuple[str, nn.Module, str]]: + """Expand plan keys into ``(module_name, module, sharding_strategy)`` shard targets.""" + module_lookup = dict(model.named_modules()) + shard_targets: list[tuple[str, nn.Module, str]] = [] + + for plan_key, sharding_strategy in fsdp_plan.items(): + if plan_key in module_lookup: + shard_targets.append((plan_key, module_lookup[plan_key], sharding_strategy)) + continue + + for module_name, module in module_lookup.items(): + if replace_layer_number_by_wildcard(module_name) == plan_key: + shard_targets.append((module_name, module, sharding_strategy)) + + return shard_targets + + +def verify_fsdp_plan(module_names: list[str], fsdp_plan: dict[str, str] | None) -> None: + """ + Verify the FSDP plan of the model, log a warning if plan keys were not applied or strategies are invalid. + """ + if not fsdp_plan: + return + + name_lookup = dict.fromkeys(module_names) + unused_rules: dict[str, str] = {} + invalid_strategies: dict[str, str] = {} + + for key, strategy in fsdp_plan.items(): + if strategy not in {"free_full_weight", "keep_full_weight"}: + invalid_strategies[key] = strategy + continue + if key not in name_lookup and not any(replace_layer_number_by_wildcard(name) == key for name in name_lookup): + unused_rules[key] = strategy + + if invalid_strategies: + logger.warning(f"The following FSDP entries have unknown strategies: {invalid_strategies}") + if unused_rules: + logger.warning(f"The following FSDP rules were not applied to any module: {unused_rules}") + + +def apply_fully_shard_data_parallel(model, fsdp_mesh): + """ + Apply FSDP2 (fully_shard) to a model. + """ + if not is_torch_available(): + raise ImportError("PyTorch is required for FSDP support") + + if not is_torch_greater_or_equal("2.6"): + raise OSError("FSDP2 requires torch>=2.6") + + fsdp_plan = dict(getattr(model, "_fsdp_plan", None) or {}) + if not fsdp_plan: + raise ValueError( + f"{type(model).__name__} does not have a FSDP2 plan declared. Set " + "`base_model_fsdp_plan` on the config and `_fsdp_plan` on the head class." + ) + + distributed_config = getattr(model.config, "distributed_config", None) + fsdp_policy_kwargs = _get_fsdp_policy_kwargs(distributed_config) + tie_word_embeddings = getattr(model.config, "tie_word_embeddings", False) + + adapted_fsdp_plan = _resolve_tied_embed_lm_head_plan(fsdp_plan, tie_word_embeddings=tie_word_embeddings) + shard_targets = expand_fsdp_plan(model, adapted_fsdp_plan) + + reshard_targets = [] + no_reshard_targets = [] + for module_name, module, sharding_strategy in shard_targets: + if sharding_strategy == "keep_full_weight": + no_reshard_targets.append((module_name, module)) + else: + reshard_targets.append((module_name, module)) + + for module_name, module in reshard_targets: + fully_shard(module, mesh=fsdp_mesh, reshard_after_forward=True, **fsdp_policy_kwargs) + logger.debug(f"Applied fully_shard to {module_name} (reshard=True)") + + # Optimization: when the keep buffer is exactly the (final_norm, lm_head/embed) + # tail pair, bundle them into one fully_shard so that we dont need to do all-gather during backward pass. + if is_norm_and_head_pair(no_reshard_targets): + module_names = [name for name, _ in no_reshard_targets] + modules = [module for _, module in no_reshard_targets] + fully_shard(modules, mesh=fsdp_mesh, reshard_after_forward=False, **fsdp_policy_kwargs) + logger.debug(f"Grouped tail {module_names} (reshard=False)") + else: + for module_name, module in no_reshard_targets: + fully_shard(module, mesh=fsdp_mesh, reshard_after_forward=False, **fsdp_policy_kwargs) + logger.debug(f"Applied fully_shard to {module_name} (reshard=False)") + + # Apply FSDP2 to the root module + fully_shard(model, mesh=fsdp_mesh, **fsdp_policy_kwargs) + + logger.info(f"FSDP2 applied to model via _fsdp_plan: {len(fsdp_plan)} entries") + + # Used by generation code to detect FSDP and enable synced_gpus. + model._is_fsdp_managed_module = True + + if tie_word_embeddings and hasattr(model, "tie_weights"): + 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..03606a2768d5 --- /dev/null +++ b/src/transformers/distributed/sharding_utils.py @@ -0,0 +1,495 @@ +# 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 + + # torch < 2.10 names as an underscore before `local_shard_size_and_offset`: alias it the non-underscored version + if not hasattr(Shard, "local_shard_size_and_offset") and hasattr(Shard, "_local_shard_size_and_offset"): + Shard.local_shard_size_and_offset = Shard._local_shard_size_and_offset + + +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..60a8e4a85c70 --- /dev/null +++ b/src/transformers/distributed/tensor_parallel.py @@ -0,0 +1,811 @@ +# 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 functools +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.placement_types import _StridedShard + + from .sharding_utils import _find_strided_shard_placement_from_fused_params + + # 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 + """ + # Lookup order (most -> least specific): + # 1. exact param key -> layers.1.mlp.experts.gate_up_proj + if parameter_name in tp_plan: + return tp_plan[parameter_name] + generic_param_name = replace_layer_number_by_wildcard(parameter_name) + # 2. wildcard param key -> layers.*.mlp.experts.gate_up_proj + if generic_param_name in tp_plan: + return tp_plan[generic_param_name] + # 3. parent module (indexed) -> layers.1.mlp.experts + elif is_weight and "." in parameter_name and (module_name := parameter_name.rsplit(".", 1)[0]) in tp_plan: + return tp_plan[module_name] + # 4. parent module (wildcard) -> layers.*.mlp.experts + 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 + + +class TensorParallelLayer: + def shard_param(self, module, param, mesh): + """Wrap ONE parameter as a DTensor placeholder. Default: no-op.""" + pass + + 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 install_forward(self, module, mesh): + """Install pre / around / post transforms by replacing module.forward.""" + 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 + + +def _matching_weight_plan_key(generic_key: str, weight_plan: dict[str, str]) -> str | None: + """Return the weight_plan entry that covers a wildcarded parameter name, or None.""" + if generic_key in weight_plan: + return generic_key + # To avoid stripping .gate_up_proj for MoE + if generic_key.endswith(".weight"): + param_name = generic_key.removesuffix(".weight") + elif generic_key.endswith(".bias"): + param_name = generic_key.removesuffix(".bias") + else: + param_name = generic_key + if param_name != generic_key and param_name in weight_plan: + return param_name + if "." in param_name and (parent := param_name.rsplit(".", 1)[0]) in weight_plan: + return parent + return None + + +def verify_tp_plan(expected_keys: list[str], tp_plan: dict[str, str] | None) -> 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, …) that override shard_param are checked. Comm-only entries + (activation, ep_router, moe_experts_allreduce, …) are excluded. + """ + if tp_plan is None: + return + + weight_plan = { + k: v + for k, v in tp_plan.items() + if v in ALL_PARALLEL_STYLES and type(ALL_PARALLEL_STYLES[v]).shard_param is not TensorParallelLayer.shard_param + } + + 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: + if (matched_key := _matching_weight_plan_key(key, weight_plan)) is not None: + unused_rules.pop(matched_key, 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 ColwiseParallel(TensorParallelLayer): + """Column-wise: weight & bias → Shard(0) (Embedding: Shard(1)); input replicated, output Shard(-1).""" + + def __init__(self, *, input_layouts=None, output_layouts=None, use_local_output: bool = True): + self.input_layouts = input_layouts or Replicate() + self.output_layouts = output_layouts if output_layouts is not None else Shard(-1) + self.use_local_output = use_local_output + + def shard_param(self, module, param, mesh): + meta = module._parameters.get(param) + if meta is None: + return + placement = Shard(1) if isinstance(module, torch.nn.Embedding) else Shard(0) + module._parameters[param] = torch.nn.Parameter( + distribute_tensor(meta, mesh, [placement], src_data_rank=None), + requires_grad=meta.requires_grad, + ) + + def transform_inputs_pre_forward(self, module, args, kwargs, mesh): + x = args[0] + if not isinstance(x, DTensor): + x = DTensor.from_local(x, mesh, [self.input_layouts], run_check=False) + if x.placements != (Replicate(),): + x = x.redistribute(placements=[Replicate()], async_op=True) + return (x,) + args[1:], kwargs # stay DTensor into F.linear + + def transform_output_post_forward(self, module, output, mesh): + if not isinstance(output, DTensor): + return output + if output.placements != (self.output_layouts,): + output = output.redistribute(placements=[self.output_layouts], async_op=True) + return output.to_local() if self.use_local_output else output + + def __repr__(self) -> str: + return ( + f"{self.__class__.__name__}(input_layouts={self.input_layouts}, " + f"output_layouts={self.output_layouts}, use_local_output={self.use_local_output})" + ) + + +class RowwiseParallel(TensorParallelLayer): + """Row-wise: weight → Shard(1), bias → Replicate (Embedding: weight → Shard(0)). + + Linear input is sharded on the last dim; Embedding input is replicated. The module + forward produces a Partial output which the boundary redistribute reduces to + output_layouts (Replicate → allreduce, Shard(1) → reduce-scatter). + """ + + def __init__(self, *, input_layouts=None, output_layouts=None, use_local_output: bool = True): + self.input_layouts = input_layouts or Shard(-1) + self.output_layouts = output_layouts or Replicate() + self.use_local_output = use_local_output + + def shard_param(self, module, param, mesh): + meta = module._parameters.get(param) + if meta is None: + return + if isinstance(module, torch.nn.Embedding): + placement = Shard(0) + else: + # bias is replicated (added after the row-reduce); weight shards on input dim + placement = Replicate() if param == "bias" else Shard(1) + module._parameters[param] = torch.nn.Parameter( + distribute_tensor(meta, mesh, [placement], src_data_rank=None), + requires_grad=meta.requires_grad, + ) + + def transform_inputs_pre_forward(self, module, args, kwargs, mesh): + # Embedding runtime sharding needs a replicated input; Linear needs Shard(-1). + desired = Replicate() if isinstance(module, torch.nn.Embedding) else Shard(-1) + x = args[0] + if not isinstance(x, DTensor): + x = DTensor.from_local(x, mesh, [self.input_layouts], run_check=False) + if x.placements != (desired,): + x = x.redistribute(placements=[desired], async_op=True) + return (x,) + args[1:], kwargs + + def transform_output_post_forward(self, module, output, mesh): + if not isinstance(output, DTensor): + return output + if output.placements != (self.output_layouts,): + output = output.redistribute(placements=[self.output_layouts], async_op=True) + return output.to_local() if self.use_local_output else output + + def __repr__(self) -> str: + return ( + f"{self.__class__.__name__}(input_layouts={self.input_layouts}, " + f"output_layouts={self.output_layouts}, use_local_output={self.use_local_output})" + ) + + +class SequenceParallel(TensorParallelLayer): + def __init__(self, *, sequence_dim: int = 1, use_local_output: bool = True): + self.sequence_dim = sequence_dim + self.use_local_output = use_local_output + + def install_forward(self, module, mesh): + # Replicate the module's params (LayerNorm/RMSNorm ones-init → from_local is safe). + for p_name, p in list(module.named_parameters(recurse=False)): + module.register_parameter( + p_name, torch.nn.Parameter(DTensor.from_local(p, mesh, [Replicate()], run_check=False)) + ) + return super().install_forward(module, mesh) + + def transform_inputs_pre_forward(self, module, args, kwargs, mesh): + seq = Shard(self.sequence_dim) + x = args[0] + if not isinstance(x, DTensor): + x = DTensor.from_local(x, mesh, [seq], run_check=False) + elif x.placements != (seq,): + x = x.redistribute(placements=[seq], async_op=True) + return (x,) + args[1:], kwargs + + def transform_output_post_forward(self, module, output, mesh): + if isinstance(output, DTensor): + return output.to_local() if self.use_local_output else output + return output + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(sequence_dim={self.sequence_dim}, use_local_output={self.use_local_output})" + + +class PrepareModuleInputOutput(TensorParallelLayer): + """Allgather input (Shard(1) → Replicate) + local split output (Replicate → Shard(1)).""" + + def __init__(self, use_local_output=True): + 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() + + +class PrepareModuleInput(TensorParallelLayer): + """Allgather a module input (default input_layout → desired_layout, then to-local).""" + + def __init__(self, *, input_kwarg=None, input_layout=None, desired_layout=None, use_local_output=True): + self.input_kwarg = input_kwarg + self.input_layout = input_layout or Shard(1) + self.desired_layout = desired_layout or Replicate() + self.use_local_output = use_local_output + + def _prepare(self, x, mesh): + if not isinstance(x, DTensor): + x = DTensor.from_local(x, mesh, [self.input_layout], run_check=False) + if x.placements != (self.desired_layout,): + x = x.redistribute(placements=[self.desired_layout]) + return x.to_local() if self.use_local_output else x + + def transform_inputs_pre_forward(self, module, args, kwargs, mesh): + if self.input_kwarg is not None: + if kwargs.get(self.input_kwarg) is not None: + kwargs = {**kwargs, self.input_kwarg: self._prepare(kwargs[self.input_kwarg], mesh)} + return args, kwargs + return (self._prepare(args[0], mesh),) + args[1:], kwargs + + def __repr__(self) -> str: + return ( + f"{self.__class__.__name__}(input_kwarg={self.input_kwarg!r}, input_layout={self.input_layout}, " + f"desired_layout={self.desired_layout}, use_local_output={self.use_local_output})" + ) + + +# ============================================================================= +# MoE / packed-linear local-param swap (grouped_mm needs plain tensors) +# ============================================================================= + + +def _accumulate_local_param_grad(original_param: DTensor, local_grad: torch.Tensor) -> torch.Tensor: + """Stitch a local grad back onto the original DTensor parameter. + + Used when the forward swap detaches the local leaf (``_StridedShard`` params) because + DTensor backward redistribute does not support that placement as a source. + """ + tensor_meta = original_param._spec.tensor_meta + detached_grad = local_grad.detach() + with torch.no_grad(): + if original_param.grad is None: + original_param.grad = DTensor.from_local( + detached_grad, + original_param.device_mesh, + original_param.placements, + run_check=False, + shape=tensor_meta.shape, + stride=tensor_meta.stride, + ) + elif isinstance(original_param.grad, DTensor): + original_param.grad._local_tensor.add_(detached_grad) + else: + original_param.grad.add_(detached_grad) + return local_grad + + +class PackedColwiseParallel(TensorParallelLayer): + """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, + ): + self.input_layouts = (input_layouts or Replicate(),) + self.use_local_output = use_local_output + self.split_factor = split_factor + + def shard_param(self, module, param, mesh): + if not isinstance(module, torch.nn.Linear): + raise NotImplementedError("PackedColwiseParallel currently only supports nn.Linear!") + meta = module._parameters.get(param) + if meta is None: + return + # Wrap as a DTensor placeholder. Runs on meta — distribute_tensor builds metadata only. + placement = _StridedShard(dim=0, split_factor=self.split_factor) + module._parameters[param] = torch.nn.Parameter( + distribute_tensor(meta, mesh, [placement], src_data_rank=None), + requires_grad=meta.requires_grad, + ) + + 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 + + @contextlib.contextmanager + def context_around_forward(self, module): + # grouped_mm etc needs plain tensors, so swap the params + to_swap_params = [ + (name, param) for name, param in module.named_parameters(recurse=False) if isinstance(param, DTensor) + ] + for name, param in to_swap_params: + del module._parameters[name] + if _find_strided_shard_placement_from_fused_params(param.placements) is not None: + local = torch.nn.Parameter(param._local_tensor.detach(), requires_grad=param.requires_grad) + if param.requires_grad: + local.register_hook(functools.partial(_accumulate_local_param_grad, param)) + else: + local = param.to_local() + setattr(module, name, local) + try: + yield + finally: + for name, param in to_swap_params: + # restore the original DTensor params + delattr(module, name) + module.register_parameter(name, param) + + 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 __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})" + ) + + +class MoEParamShard(TensorParallelLayer): + """Param-only TP/EP style for MoE expert weights: shard the named 3D params as DTensor + placeholders; the matching forward comm is a separate moe_experts_allreduce entry. + + shards_expert_dim (EP grouped_gemm) shards dim 0 and updates module.num_experts to the + per-rank local count, so the experts forward and ep_router sentinel agree. + """ + + def __init__(self, placement, *, shards_expert_dim: bool = False): + self.placement = placement + self.shards_expert_dim = shards_expert_dim + + def shard_param(self, module, param, mesh): + # Wrap one expert param as a DTensor placeholder. Runs on meta — + # distribute_tensor builds metadata only, no collective. + meta = module._parameters.get(param) + if meta is None: + return + if self.shards_expert_dim: + # dim 0 is the expert dimension; record the per-rank local count so the + # experts forward and the EP router agree on local expert ids/sentinel. + module.num_experts = meta.shape[0] // mesh.size() + module._parameters[param] = torch.nn.Parameter( + distribute_tensor(meta, mesh, [self.placement], src_data_rank=None), + requires_grad=meta.requires_grad, + ) + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(placement={self.placement}, shards_expert_dim={self.shards_expert_dim})" + + +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(TensorParallelLayer): + def __init__(self, output_layouts=None): + self.output_layouts = output_layouts or Replicate() + + def transform_inputs_pre_forward(self, module, args, kwargs, mesh, *, is_expert_parallel=False): + hidden_states, top_k_index, top_k_weights = args + tp_group = mesh.get_group() if mesh.ndim == 1 else mesh.get_group("tp") + if not isinstance(hidden_states, DTensor): + hidden_states = DTensor.from_local(hidden_states, mesh, [Replicate()], run_check=False) + hidden_states = hidden_states.to_local() + hidden_states = _AllReduceBackward.apply(hidden_states, tp_group) + + if isinstance(top_k_weights, DTensor): + top_k_weights = top_k_weights.to_local() + # Under TP the router runs replicated, so its routing-weight gradient is partial + # across ranks and needs an all-reduce-sum in backward. Under EP the router output + # is already sliced per-rank by `ep_router` (non-local scores zeroed), so each + # rank's routing weights are independent — skip the all-reduce to avoid + # double-counting the gradient. + if not is_expert_parallel: + top_k_weights = _AllReduceBackward.apply(top_k_weights, tp_group) + + return (hidden_states, top_k_index, top_k_weights), kwargs + + def install_forward(self, module, mesh, *, is_expert_parallel=False): + """Install pre / around / post transforms; ``is_expert_parallel`` is baked into the closure.""" + original_forward = module.forward + + def tp_forward(*args, **kwargs): + args, kwargs = self.transform_inputs_pre_forward( + module, args, kwargs, mesh, is_expert_parallel=is_expert_parallel + ) + 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 + + @contextlib.contextmanager + def context_around_forward(self, module): + # grouped_mm experts forward needs plain tensors, so swap the params + to_swap_params = [ + (name, param) for name, param in module.named_parameters(recurse=False) if isinstance(param, DTensor) + ] + for name, param in to_swap_params: + del module._parameters[name] + # Happens only when train in TP only (experts dont use grouped_gemm but rely on moe_tp_gate_up_colwise) + if _find_strided_shard_placement_from_fused_params(param.placements) is not None: + local = torch.nn.Parameter(param._local_tensor.detach(), requires_grad=param.requires_grad) + if param.requires_grad: + local.register_hook(functools.partial(_accumulate_local_param_grad, param)) + else: + local = param.to_local() + setattr(module, name, local) + try: + yield + finally: + # restore the original DTensor params + for name, param in to_swap_params: + delattr(module, name) + module.register_parameter(name, param) + + 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 MoeIdentityParallel(TensorParallelLayer): + """ + TP class for zero/identity experts in MoE layers. + + Under TP, the parent `MoEExpertsParallel` all-reduces the expert module + output by summation. 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 transform_inputs_pre_forward(self, module, args, kwargs, mesh): + input_tensor = args[0] + return (input_tensor / mesh.size(), *args[1:]), kwargs + + +class EpRouterParallel(TensorParallelLayer): + """Expert-parallel router: forward-only slicing of router outputs to local experts. + + Ported from the original RouterParallel (#39501). The gate runs replicated on + every rank and emits global (router_logits, router_scores, router_indices). Under + EP each rank owns num_experts // ep_size experts, so this post-forward hook zeroes + the scores of non-local experts and remaps the surviving indices into the local range, + using num_local_experts as the sentinel for dropped (non-local) slots. The + downstream experts forward masks that sentinel and moe_experts_allreduce sums the + partial per-rank results. + + No parameter sharding — the gate weight stays replicated. + + Example: 128 experts, EP=8 → num_local_experts=16. Rank 0 (owns experts 0-15) keeps + only indices in [0, 16), remaps them via fmod, and sets everything else to the + sentinel 16; its scores are zeroed everywhere except the surviving local slots. + """ + + def transform_output_post_forward(self, module, output, mesh): + ep_rank, ep_size = mesh.get_local_rank(), mesh.size() + num_experts = getattr(module, "num_experts", None) + if num_experts is None: + num_experts = getattr(getattr(module, "config", None), "num_experts", None) + if num_experts is None: + raise AttributeError( + f"Router module {type(module).__name__} is missing `num_experts` and `config.num_experts`" + ) + if num_experts % ep_size != 0: + raise ValueError(f"num_experts must be divisible by ep_size: {num_experts} % {ep_size} != 0") + num_local_experts = num_experts // ep_size + + router_logits, router_scores, router_indices = output + 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) + # `-1 % 1 == 0`, so fmod only remaps correctly when there is more than one local expert. + 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 + + +class ParallelInterface(GeneralInterface): + """Registry of named TP styles. Configs and modeling files reference these by string name. + + Styles are split into the three parallelism groups they belong to (TENSOR_/SEQUENCE_/ + EXPERT_PARALLEL_STYLES); _global_mapping is the union of all three. Keeping them separate + lets other code reason about a group as a whole — e.g. SEQUENCE_PARALLEL_STYLES is reused to + validate that a user-provided tp_plan actually shards the sequence dimension when + enable_sequence_parallel=True. + + Adding a new entry under one of the three groups 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, so the dicts live behind a + torch-availability guard and are empty when torch is unavailable (this module stays + importable; styles are only ever looked up inside apply_tensor_parallel, which needs torch). + """ + + if is_torch_available() and is_torch_greater_or_equal("2.5") and _torch_distributed_available: + TENSOR_PARALLEL_STYLES = { + # 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()), + # MoE intra-expert weight sharding — param-level (no forward hook). The matching forward + # comm is declared separately via "moe_experts_allreduce" on the experts module. gate_up + # is packed (gate||up) along dim -2 (Qwen3/Mixtral: [E, 2*inter, hidden]). + "moe_tp_gate_up_colwise": MoEParamShard(_StridedShard(dim=-2, split_factor=2)), # packed gate/up + "moe_tp_down_rowwise": MoEParamShard(Shard(-1)), # down_proj input dim + # Row-parallel + "rowwise_allreduce": RowwiseParallel(input_layouts=Shard(-1), output_layouts=Replicate()), + # Vocab / embedding (rowwise sharding on vocab dim) + "vocab_allreduce": RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()), + } + SEQUENCE_PARALLEL_STYLES = { + "rowwise_reduce_scatter": RowwiseParallel(input_layouts=Shard(-1), output_layouts=Shard(1)), + "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 (allgather seq-sharded input → replicate, to local). + # use_local_output=True: our modeling code expects plain tensors downstream. + "module_allgather": PrepareModuleInput(input_layout=Shard(1), desired_layout=Replicate()), + "module_allgather_hidden_states": PrepareModuleInput( + input_kwarg="hidden_states", input_layout=Shard(1), desired_layout=Replicate() + ), + "module_allgather_split": PrepareModuleInputOutput(), + } + EXPERT_PARALLEL_STYLES = { + "grouped_gemm": MoEParamShard(Shard(0), shards_expert_dim=True), # expert dim + # Experts' forward comm (no weight sharding — that is declared per-param via "grouped_gemm" + # / "moe_tp_*"). Installs the forward hook that all-reduces partial per-rank expert outputs. + # Reused by dense TP-MoE as well as EP. + "moe_experts_allreduce": MoEExpertsParallel(output_layouts=Replicate()), + "moe_identity_expert": MoeIdentityParallel(), + # EP router — forward-only slicing of router outputs to local experts. + "ep_router": EpRouterParallel(), + } + else: + TENSOR_PARALLEL_STYLES = SEQUENCE_PARALLEL_STYLES = EXPERT_PARALLEL_STYLES = {} + + _global_mapping = {**TENSOR_PARALLEL_STYLES, **SEQUENCE_PARALLEL_STYLES, **EXPERT_PARALLEL_STYLES} + + +ALL_PARALLEL_STYLES: ParallelInterface = ParallelInterface() + + +def select_parallel_plan(model) -> dict[str, str]: + """ + Select the parallel plan to apply from explicit combo plans on the model. + """ + distributed_config = model.config.distributed_config + user_tp_plan = distributed_config.tp_plan + if user_tp_plan is not None: + return dict(user_tp_plan) + + sp = distributed_config.enable_sequence_parallel + ep = distributed_config.enable_expert_parallel + if sp and ep: + plan_attr = "_sp_ep_plan" + elif sp: + plan_attr = "_sp_plan" + elif ep: + plan_attr = "_tp_ep_plan" + else: + plan_attr = "_tp_plan" + config_attr = f"base_model{plan_attr}" + + plan = getattr(model, plan_attr, None) or {} + if not plan: + raise ValueError( + f"Model {model.config.model_type!r} has no {config_attr} (model.{plan_attr} is empty) but " + f"DistributedConfig requested enable_sequence_parallel={sp}, enable_expert_parallel={ep}. " + f"Add {config_attr} to the model config, set DistributedConfig.tp_plan explicitly, " + f"or change the flags." + ) + return dict(plan) + + +def apply_tensor_parallel(model, tp_mesh): + """Apply tensor parallelism by looking each style up by string name. + + A single named_modules() walk shards each module's direct params (param-keyed plan + entries → shard_param) and installs its forward comm (module-keyed entries → + install_forward). Plus the cross-cutting setup: tie-weights handling, SP position_ids + injection, and loss_parallel activation. + """ + enable_sp = bool(getattr(model.config.distributed_config, "enable_sequence_parallel", False)) + enable_ep = bool(getattr(model.config.distributed_config, "enable_expert_parallel", False)) + + # case when enable_sequence_parallel=True and tp_plan!=None (but tp_plan has no sequence parallel styles keys) raise an error + user_tp_plan = model.config.distributed_config.tp_plan + if ( + enable_sp + and user_tp_plan is not None # user provided tp_plan manually + and not any(style in ParallelInterface.SEQUENCE_PARALLEL_STYLES for style in user_tp_plan.values()) + ): + raise ValueError( + "enable_sequence_parallel=True but the provided tp_plan shards no sequence dimension " + f"(none of {sorted(ParallelInterface.SEQUENCE_PARALLEL_STYLES)}). Provide a sequence-parallel plan " + "(e.g. 'vocab_reduce_scatter' on the embedding) or set enable_sequence_parallel=False. " + f"tp_plan: {user_tp_plan}" + ) + + model.tp_plan = dict(select_parallel_plan(model)) + logger.info(f"TP plan has been resolved: {model.tp_plan}") + + # 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 model.tp_plan) + if not tied_source_in_plan: + model.tp_plan.pop("lm_head", None) + + for name, module in model.named_modules(): + # Shard each parameter directly on the module using the specified style. + for p_name, _ in list(module.named_parameters(recurse=False)): + full = f"{name}.{p_name}" if name else p_name + style_name = _get_parameter_tp_plan(parameter_name=full, tp_plan=model.tp_plan, is_weight=True) + if style_name is not None and style_name in ALL_PARALLEL_STYLES: + ALL_PARALLEL_STYLES[style_name].shard_param(module, p_name, tp_mesh) + # Install forward hooks for modules as needed by the plan. + style_name = _get_parameter_tp_plan(parameter_name=name, tp_plan=model.tp_plan, is_weight=False) + if style_name is not None and style_name in ALL_PARALLEL_STYLES: + if style_name == "moe_experts_allreduce": + ALL_PARALLEL_STYLES[style_name].install_forward(module, tp_mesh, is_expert_parallel=enable_ep) + else: + ALL_PARALLEL_STYLES[style_name].install_forward(module, 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 model.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..4146058a7e84 --- /dev/null +++ b/src/transformers/distributed/utils.py @@ -0,0 +1,299 @@ +# 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.state_dict import ( + get_model_state_dict, + get_optimizer_state_dict, + set_optimizer_state_dict, + ) + from torch.distributed.tensor import DTensor + + if is_torch_greater_or_equal("2.7"): + from torch.distributed.checkpoint.hf_storage import HuggingFaceStorageWriter + + +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`.") + + if distributed_config.fsdp_size > 1 and not is_torch_greater_or_equal("2.6"): + raise OSError("FSDP2 requires `torch>=2.6`.") + + 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 the N-dimensional device mesh + mesh = torch.distributed.init_device_mesh(device_type, tuple(dims), mesh_dim_names=tuple(names)) + # If N > 1, create a flattened sub-mesh so all-reduces across the world mesh ae done in one collective + 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) + 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) + 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. + """ + if not is_torch_greater_or_equal("2.7"): + raise OSError("Distributed checkpoint saving requires `torch>=2.7`.") + + 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 93588ae77b8d..7405d1142e4e 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 ( @@ -118,6 +118,18 @@ logger = logging.get_logger(__name__) + +def _get_dist_world_size() -> int: + if ( + dist.is_available() + and hasattr(dist, "is_initialized") + and dist.is_initialized() + and hasattr(dist, "get_world_size") + ): + return dist.get_world_size() + return 1 + + if is_accelerate_available(): from accelerate.hooks import AlignDevicesHook, add_hook_to_module @@ -2211,7 +2223,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 = _get_dist_world_size() generation_mode_kwargs["synced_gpus"] = ( (is_deepspeed_zero3_enabled() or is_fsdp_managed_module(self)) and world_size > 1 if synced_gpus is None @@ -2658,7 +2670,8 @@ 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 + if hasattr(dist, "all_reduce") and hasattr(dist, "ReduceOp"): + 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 2ecc31ae54cf..27181748defc 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"], "gemma_quant": [ "QuantizedEmbedding", "QuantizedLinear", @@ -167,11 +166,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() @@ -216,7 +210,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 .gemma_quant import ( QuantizedEmbedding, QuantizedLinear, @@ -319,12 +312,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 8fc01b7f09e3..7bf8ccef724b 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 f6a2e533585e..65b79f19c37b 100644 --- a/src/transformers/integrations/moe.py +++ b/src/transformers/integrations/moe.py @@ -371,7 +371,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. @@ -446,6 +446,12 @@ def grouped_mm_experts_forward( proj_out = _grouped_linear( selected_hidden_states_g, selected_weights, offsets, bias=selected_biases, is_transposed=self.is_transposed ) # (S, 2 * intermediate_dim) or (S, intermediate_dim) depending on whether we have gating + # Mid-mask (bwd path). The up grouped_mm leaves sentinel rows of `proj_out` uninit. This + # matters in BACKWARD: the masked_fill_ backward zeros the gradient at sentinel rows (select, + # so even NaN → 0) before the non-linearity bwd's garbage can reach grad(gate_up_proj[_bias]). + # Forward needs no masking here (the down grouped_mm skips these rows) — the value-zeroing is + # just a harmless side effect. + proj_out.masked_fill_(sentinel_mask, 0.0) # Apply gating or activation if self.has_gate: diff --git a/src/transformers/integrations/mxfp4.py b/src/transformers/integrations/mxfp4.py index 100a02e5b250..0086644cf18d 100644 --- a/src/transformers/integrations/mxfp4.py +++ b/src/transformers/integrations/mxfp4.py @@ -488,28 +488,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) @@ -534,24 +521,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/modeling_utils.py b/src/transformers/modeling_utils.py index 2522d5101a75..99f58106ecda 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, verify_fsdp_plan +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 @@ -147,8 +152,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 @@ -186,6 +189,8 @@ 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 + fsdp_plan: dict[str, str] | None = None weights_only: bool = True weight_mapping: list[WeightConverter | WeightRenaming] | None = None disable_mmap: bool | None = None @@ -209,6 +214,12 @@ def _get_torch_distributed_world_size() -> int: return torch.distributed.get_world_size() +def _get_torch_distributed_rank() -> int: + if not _is_torch_distributed_initialized() or not hasattr(torch.distributed, "get_rank"): + return 0 + return torch.distributed.get_rank() + + def is_local_dist_rank_0(): return _is_torch_distributed_initialized() and int(os.environ.get("LOCAL_RANK", "-1")) == 0 @@ -1249,6 +1260,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 @@ -1395,13 +1410,34 @@ 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_tp_ep_plan = getattr(self, "_tp_ep_plan", None) or {} + cls_sp_ep_plan = getattr(self, "_sp_ep_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._tp_ep_plan = dict(cls_tp_ep_plan) + self._sp_ep_plan = dict(cls_sp_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._ep_plan = self.config.base_model_ep_plan.copy() if self.config.base_model_ep_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._tp_ep_plan = ( + self.config.base_model_tp_ep_plan.copy() if self.config.base_model_tp_ep_plan is not None else {} + ) + self._sp_ep_plan = ( + self.config.base_model_sp_ep_plan.copy() if self.config.base_model_sp_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` @@ -1419,12 +1455,18 @@ def post_init(self): # This works because the way the `__init__` and `post_init` are called on all submodules is depth-first in the graph for name, module in self.named_children(): # Parallel plans - if plan := getattr(module, "_ep_plan", None): - 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, "_tp_ep_plan", None): + self._tp_ep_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()}) + if plan := getattr(module, "_sp_ep_plan", None): + self._sp_ep_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()}) @@ -1454,13 +1496,12 @@ def post_init(self): @property def tp_plan(self) -> dict[str, str]: - """ - The full tp plan for the model's modules - """ - if hasattr(self.config, "distributed_config") and self.config.distributed_config.enable_expert_parallel: - return self._ep_plan return self._tp_plan + @property + def fsdp_plan(self) -> dict[str, str]: + return self._fsdp_plan + @property def pp_plan(self) -> dict[str, tuple[str, str]]: return self._pp_plan @@ -1473,14 +1514,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()] @@ -2085,6 +2118,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__"): @@ -2104,6 +2139,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__"): @@ -3344,6 +3381,7 @@ def save_pretrained( token: str | bool | None = None, save_peft_format: bool = True, save_original_format: bool = True, + distributed_checkpoint: bool = False, **kwargs, ): """ @@ -3409,7 +3447,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." ) @@ -3419,9 +3457,13 @@ def save_pretrained( return os.makedirs(save_directory, exist_ok=True) + save_on_this_rank = is_main_process + if _is_torch_distributed_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 _get_torch_distributed_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) @@ -3446,46 +3488,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) - active_adapter = self.active_adapters() + 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 - 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] + active_adapter = self.active_adapters() - current_peft_config = self.peft_config[active_adapter] - current_peft_config.save_pretrained(save_directory) + 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) + finally: + if distributed_config is not None: + model_to_save.config.distributed_config = distributed_config + + if distributed_checkpoint: + if _is_torch_distributed_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 + # 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 @@ -3511,10 +3579,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) @@ -3556,54 +3620,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) @@ -3619,6 +3684,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 [] @@ -4015,13 +4087,23 @@ 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_cpu_offload`, and + `fsdp_mixed_precision`. 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 uses the model's ``base_model_fsdp_plan`` with FSDP2 + (`fully_shard`). `tp_plan` can override the tensor parallel layout (e.g. + `{"model.layers.*.self_attn.q_proj": "colwise"}`). + + Examples: + - TP-only: `DistributedConfig(tp_size=4)` + - FSDP-only: `DistributedConfig(fsdp_size=4)` + - FSDP with mixed precision: `DistributedConfig(fsdp_size=4, fsdp_mixed_precision=True)` + - 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 @@ -4115,8 +4197,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) @@ -4125,8 +4205,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"]: @@ -4157,17 +4246,24 @@ 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) + # When using any distributed setup, we simply set the device_map here to the current rank so that we correctly + # load the params on the right device + if device_mesh.device_type == "cpu": + device_map = {"": torch.device("cpu")} + else: + device_map = {"": torch.device(device_mesh.device_type, int(os.environ["LOCAL_RANK"]))} + 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`.") @@ -4180,7 +4276,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: @@ -4317,17 +4412,15 @@ 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: + if distributed_config is not None: + model = distribute_model(model, distributed_config, device_mesh) + elif device_map is not None: + # Expand device_map if it was passed as a `str`, i.e. `device_map="auto"` device_map = _get_device_map(model, device_map, max_memory, hf_quantizer) - # Finalize model weight initialization load_config = LoadStateDictConfig( pretrained_model_name_or_path=pretrained_model_name_or_path, ignore_mismatched_sizes=ignore_mismatched_sizes, @@ -4339,6 +4432,8 @@ def from_pretrained( dtype_plan=dtype_plan, hf_quantizer=hf_quantizer, device_mesh=device_mesh, + tp_plan=model.tp_plan, + fsdp_plan=model.fsdp_plan, weights_only=weights_only, weight_mapping=weight_conversions, use_safetensors=use_safetensors, @@ -4404,10 +4499,12 @@ def _load_pretrained_model( } # Model's definition arriving here is final (TP hooks added, quantized layers replaces) - expected_keys = list(model.state_dict().keys()) if expected_keys is None else expected_keys + expected_tp_plan_keys = list(model.state_dict().keys()) if expected_keys is None else expected_keys + expected_fsdp_plan_keys = [name for name, _ in model.named_modules()] if logger.level >= logging.WARNING: - verify_tp_plan(expected_keys, getattr(model, "_tp_plan", None)) + verify_tp_plan(expected_tp_plan_keys, load_config.tp_plan) + verify_fsdp_plan(expected_fsdp_plan_keys, load_config.fsdp_plan) # This offload index if for params explicitly on the "disk" in the device_map disk_offload_index = None @@ -4425,7 +4522,7 @@ def _load_pretrained_model( # Warmup cuda to load the weights much faster on devices if load_config.device_map is not None and not is_hqq_or_quark: - expanded_device_map = expand_device_map(load_config.device_map, expected_keys) + expanded_device_map = expand_device_map(load_config.device_map, expected_tp_plan_keys) caching_allocator_warmup(model, expanded_device_map, load_config.hf_quantizer) error_msgs = [] @@ -4484,7 +4581,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, ) @@ -4636,8 +4733,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): @@ -4756,15 +4853,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(): diff --git a/src/transformers/models/afmoe/modeling_afmoe.py b/src/transformers/models/afmoe/modeling_afmoe.py index 421119b33deb..cf4a595bdcda 100644 --- a/src/transformers/models/afmoe/modeling_afmoe.py +++ b/src/transformers/models/afmoe/modeling_afmoe.py @@ -615,8 +615,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..b055b222e38d 100644 --- a/src/transformers/models/afmoe/modular_afmoe.py +++ b/src/transformers/models/afmoe/modular_afmoe.py @@ -395,7 +395,11 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..85bb6b27fd23 100644 --- a/src/transformers/models/apertus/modeling_apertus.py +++ b/src/transformers/models/apertus/modeling_apertus.py @@ -422,8 +422,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..5aaee35657a2 100644 --- a/src/transformers/models/arcee/modeling_arcee.py +++ b/src/transformers/models/arcee/modeling_arcee.py @@ -424,8 +424,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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 6ac87ac8d3b9..479f4e3d80d9 100644 --- a/src/transformers/models/aria/modeling_aria.py +++ b/src/transformers/models/aria/modeling_aria.py @@ -759,8 +759,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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 022d536699cc..58a49d034d01 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 330a1a7d4beb..a9cecb8842a3 100644 --- a/src/transformers/models/audioflamingo3/modeling_audioflamingo3.py +++ b/src/transformers/models/audioflamingo3/modeling_audioflamingo3.py @@ -447,6 +447,7 @@ def forward(self, audio_features): ) class AudioFlamingo3Model(AudioFlamingo3PreTrainedModel): _tp_plan = None + _sp_plan = None _pp_plan = None _keep_in_fp32_modules_strict = None diff --git a/src/transformers/models/audioflamingo3/modular_audioflamingo3.py b/src/transformers/models/audioflamingo3/modular_audioflamingo3.py index e312a87394fb..def99468fdc8 100644 --- a/src/transformers/models/audioflamingo3/modular_audioflamingo3.py +++ b/src/transformers/models/audioflamingo3/modular_audioflamingo3.py @@ -182,6 +182,7 @@ def __init__(self, config: AudioFlamingo3Config): ) class AudioFlamingo3Model(VoxtralModel): _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 95e7696e7d24..77a49ebe7d00 100644 --- a/src/transformers/models/bamba/modeling_bamba.py +++ b/src/transformers/models/bamba/modeling_bamba.py @@ -1070,8 +1070,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..b51bd8fdabf1 100644 --- a/src/transformers/models/bitnet/modeling_bitnet.py +++ b/src/transformers/models/bitnet/modeling_bitnet.py @@ -423,7 +423,11 @@ def forward( class BitNetForCausalLM(BitNetPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = None + _sp_plan = None + _tp_ep_plan = None + _sp_ep_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..9bb7991e6d72 100644 --- a/src/transformers/models/bitnet/modular_bitnet.py +++ b/src/transformers/models/bitnet/modular_bitnet.py @@ -110,6 +110,9 @@ class BitNetModel(LlamaModel): class BitNetForCausalLM(LlamaForCausalLM): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = None + _sp_plan = None + _tp_ep_plan = None + _sp_ep_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 ba9b68e3e586..6d4a6de09588 100644 --- a/src/transformers/models/cohere/modeling_cohere.py +++ b/src/transformers/models/cohere/modeling_cohere.py @@ -452,8 +452,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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 471797ee8e8f..e78e1adaae6b 100644 --- a/src/transformers/models/cohere2/modeling_cohere2.py +++ b/src/transformers/models/cohere2/modeling_cohere2.py @@ -430,8 +430,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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 b40a5056232b..ecd41c3e86d4 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/cohere2_moe/configuration_cohere2_moe.py b/src/transformers/models/cohere2_moe/configuration_cohere2_moe.py index 9c1abfc06cd7..76050d934197 100644 --- a/src/transformers/models/cohere2_moe/configuration_cohere2_moe.py +++ b/src/transformers/models/cohere2_moe/configuration_cohere2_moe.py @@ -74,16 +74,38 @@ class Cohere2MoeConfig(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", + # Depending on layers, `mlp` can be a MoE or an MLP layer + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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", + } + # 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", + # Depending on layers, `mlp` can be a MoE or an MLP layer - so they are updated by `_update_sp_plan` to avoid collisions + "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", + } vocab_size: int = 256000 hidden_size: int = 8192 @@ -142,7 +164,6 @@ def __post_init__(self, **kwargs): for i in range(self.num_hidden_layers - first_k_dense_replace) ] self.layer_types = prefix_layers + rest_layers - self.validate_layer_type() if self.mlp_layer_types is None: @@ -151,6 +172,30 @@ def __post_init__(self, **kwargs): ] super().__post_init__(**kwargs) + self._update_sp_plan() + + def _update_sp_plan(self): + """Depending on layers, `mlp` can be a MoE or an MLP layer - so we update the plan dynamically here to avoid collisions""" + self.base_model_sp_plan = self.base_model_sp_plan.copy() + for i, mlp_type in enumerate(self.mlp_layer_types): + # This is a MLP layer + if mlp_type == "dense": + self.base_model_sp_plan.update( + { + f"layers.{i}.mlp": "module_allgather", + f"layers.{i}.mlp.gate_proj": "colwise", + f"layers.{i}.mlp.up_proj": "colwise", + f"layers.{i}.mlp.down_proj": "rowwise_reduce_scatter", + } + ) + # This is a MoE layer + else: + self.base_model_sp_plan.update( + { + f"layers.{i}.mlp": "module_allgather_split", + f"layers.{i}.mlp.experts": "moe_experts_allreduce", + } + ) __all__ = ["Cohere2MoeConfig"] diff --git a/src/transformers/models/cohere2_moe/modeling_cohere2_moe.py b/src/transformers/models/cohere2_moe/modeling_cohere2_moe.py index 5e2a37c43d1d..20b4abef6c37 100644 --- a/src/transformers/models/cohere2_moe/modeling_cohere2_moe.py +++ b/src/transformers/models/cohere2_moe/modeling_cohere2_moe.py @@ -580,8 +580,12 @@ def forward( @auto_docstring class Cohere2MoeForCausalLM(Cohere2MoePreTrainedModel, 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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/csm/modeling_csm.py b/src/transformers/models/csm/modeling_csm.py index dd0a1f0aa749..97f5bf6425aa 100644 --- a/src/transformers/models/csm/modeling_csm.py +++ b/src/transformers/models/csm/modeling_csm.py @@ -549,7 +549,11 @@ def forward(self, hidden_states, codebook_indices=None): class CsmDepthDecoderForCausalLM(CsmPreTrainedModel, GenerationMixin): _tied_weights_keys = None _tp_plan = None + _sp_plan = None + _tp_ep_plan = None + _sp_ep_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..4fa7c150285b 100644 --- a/src/transformers/models/csm/modular_csm.py +++ b/src/transformers/models/csm/modular_csm.py @@ -274,6 +274,9 @@ def forward(self, hidden_states, codebook_indices=None): class CsmDepthDecoderForCausalLM(LlamaForCausalLM, GenerationMixin): _tied_weights_keys = None _tp_plan = None + _sp_plan = None + _tp_ep_plan = None + _sp_ep_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..fd4e9176aef8 100644 --- a/src/transformers/models/cwm/modeling_cwm.py +++ b/src/transformers/models/cwm/modeling_cwm.py @@ -426,8 +426,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..a425a78ff82d 100644 --- a/src/transformers/models/dbrx/modeling_dbrx.py +++ b/src/transformers/models/dbrx/modeling_dbrx.py @@ -642,7 +642,11 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..ab7b3998b085 100644 --- a/src/transformers/models/dbrx/modular_dbrx.py +++ b/src/transformers/models/dbrx/modular_dbrx.py @@ -430,7 +430,11 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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_ocr2/configuration_deepseek_ocr2.py b/src/transformers/models/deepseek_ocr2/configuration_deepseek_ocr2.py index 2d9c601c1f2a..34e666e7a260 100644 --- a/src/transformers/models/deepseek_ocr2/configuration_deepseek_ocr2.py +++ b/src/transformers/models/deepseek_ocr2/configuration_deepseek_ocr2.py @@ -88,7 +88,6 @@ class DeepseekOcr2VisionEncoderConfig(PreTrainedConfig): model_type = "deepseek_ocr2_encoder" keys_to_ignore_at_inference = ["past_key_values"] - # Default tensor parallel plan for base model `DeepseekOcr2VisionEncoder` base_model_tp_plan = { "layers.*.self_attn.q_proj": "colwise", diff --git a/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py b/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py index cddbd6f2151e..a9ebd9471616 100644 --- a/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py +++ b/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py @@ -533,6 +533,18 @@ class DeepseekOcr2VisionEncoderConfig(Qwen2Config): ```""" base_config_key = "encoder_config" + # Default tensor parallel plan for base model `DeepseekOcr2VisionEncoder` + 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.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", + } + base_model_sp_plan = AttributeError() + base_model_fsdp_plan = AttributeError() @auto_docstring(checkpoint="deepseek-community/DeepSeek-OCR-2") @@ -600,6 +612,8 @@ class DeepseekOcr2TextConfig(DeepseekV2Config): "layers.*.mlp.up_proj": "colwise", "layers.*.mlp.down_proj": "rowwise", } + base_model_sp_plan = AttributeError() + base_model_fsdp_plan = AttributeError() # Remove unused attributes inherited from DeepseekV2Config first_k_dense_replace = AttributeError() diff --git a/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py b/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py index 1b8005f8efef..f0093fc52a3f 100644 --- a/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py +++ b/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py @@ -54,18 +54,43 @@ 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.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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 +98,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..1d4995293a32 100644 --- a/src/transformers/models/deepseek_v2/modeling_deepseek_v2.py +++ b/src/transformers/models/deepseek_v2/modeling_deepseek_v2.py @@ -541,8 +541,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..9458009391cb 100644 --- a/src/transformers/models/deepseek_v2/modular_deepseek_v2.py +++ b/src/transformers/models/deepseek_v2/modular_deepseek_v2.py @@ -69,18 +69,43 @@ 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.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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..a94ef5660e80 100644 --- a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py @@ -49,21 +49,27 @@ 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.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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 +117,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 9ca55dad5580..876a2ccd5a36 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -635,8 +635,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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_v32/modular_deepseek_v32.py b/src/transformers/models/deepseek_v32/modular_deepseek_v32.py index 77be4aa9c943..11d12a756eb8 100644 --- a/src/transformers/models/deepseek_v32/modular_deepseek_v32.py +++ b/src/transformers/models/deepseek_v32/modular_deepseek_v32.py @@ -98,6 +98,7 @@ class DeepseekV32Config(Glm4MoeLiteConfig, RotaryEmbeddingConfigMixin): "layers.*.mlp.up_proj": "colwise", "layers.*.mlp.down_proj": "rowwise", } + base_model_fsdp_plan = AttributeError() attribute_map = {"num_local_experts": "num_experts"} @@ -369,7 +370,11 @@ class DeepseekV32Model(DeepseekV3Model): class DeepseekV32ForCausalLM(DeepseekV3ForCausalLM): - pass + _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_ep_plan = AttributeError() + _sp_plan = AttributeError() + _sp_ep_plan = AttributeError() + _fsdp_plan = AttributeError() __all__ = [ diff --git a/src/transformers/models/deepseek_v4/configuration_deepseek_v4.py b/src/transformers/models/deepseek_v4/configuration_deepseek_v4.py index 78e7f11291b1..d9498c1bc984 100644 --- a/src/transformers/models/deepseek_v4/configuration_deepseek_v4.py +++ b/src/transformers/models/deepseek_v4/configuration_deepseek_v4.py @@ -130,12 +130,18 @@ 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", "layers.*.self_attn.compressor.indexer.q_b_proj": "colwise", "layers.*.self_attn.compressor.indexer.scorer.weights_proj": "colwise", "layers.*.self_attn.compressor.indexer.scorer": "all_reduce", } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 129280 hidden_size: int = 4096 moe_intermediate_size: int = 2048 diff --git a/src/transformers/models/deepseek_v4/modeling_deepseek_v4.py b/src/transformers/models/deepseek_v4/modeling_deepseek_v4.py index b0fc580df777..dcbbb9f0ed47 100644 --- a/src/transformers/models/deepseek_v4/modeling_deepseek_v4.py +++ b/src/transformers/models/deepseek_v4/modeling_deepseek_v4.py @@ -1423,8 +1423,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..51f4642d3366 100644 --- a/src/transformers/models/diffllama/modeling_diffllama.py +++ b/src/transformers/models/diffllama/modeling_diffllama.py @@ -660,8 +660,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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/diffusion_gemma/modular_diffusion_gemma.py b/src/transformers/models/diffusion_gemma/modular_diffusion_gemma.py index 331bc5187f7d..963eab4dd7c8 100644 --- a/src/transformers/models/diffusion_gemma/modular_diffusion_gemma.py +++ b/src/transformers/models/diffusion_gemma/modular_diffusion_gemma.py @@ -91,11 +91,37 @@ class DiffusionGemmaTextConfig(Gemma4TextConfig): model_type = "diffusion_gemma_text" final_logit_softcapping = 30.0 + 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.*.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", + } + 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.*.router": "ep_router", + "layers.*.experts.gate_up_proj": "grouped_gemm", + "layers.*.experts.down_proj": "grouped_gemm", + "layers.*.experts": "moe_tp_experts", + } + base_model_tp_ep_plan = AttributeError() base_model_pp_plan = { "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = AttributeError() enable_moe_block = AttributeError() attention_k_eq_v = AttributeError() 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..885321f8d197 100644 --- a/src/transformers/models/doge/modeling_doge.py +++ b/src/transformers/models/doge/modeling_doge.py @@ -716,8 +716,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..565fd7b4bd42 100644 --- a/src/transformers/models/dots1/configuration_dots1.py +++ b/src/transformers/models/dots1/configuration_dots1.py @@ -51,18 +51,16 @@ 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.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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 +68,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..8708a661c2fe 100644 --- a/src/transformers/models/dots1/modeling_dots1.py +++ b/src/transformers/models/dots1/modeling_dots1.py @@ -569,8 +569,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..74ae84c5c6b1 100644 --- a/src/transformers/models/dots1/modular_dots1.py +++ b/src/transformers/models/dots1/modular_dots1.py @@ -65,18 +65,16 @@ 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.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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 +82,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 6f43af5f8350..4a48b9e60dbd 100644 --- a/src/transformers/models/emu3/modeling_emu3.py +++ b/src/transformers/models/emu3/modeling_emu3.py @@ -1275,8 +1275,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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 b910f1cbd437..55691646b590 100644 --- a/src/transformers/models/ernie4_5/modeling_ernie4_5.py +++ b/src/transformers/models/ernie4_5/modeling_ernie4_5.py @@ -421,8 +421,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..5983c9d0fd15 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,16 @@ 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.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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 +83,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 663cb83ae32d..53c04e6bb86b 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 @@ -653,8 +653,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..e54aff4408b9 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,19 @@ 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.text_moe.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.text_moe.experts.down_proj": "moe_tp_down_rowwise", + "layers.*.mlp.text_moe.experts": "moe_experts_allreduce", + "layers.*.mlp.vision_moe.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.vision_moe.experts.down_proj": "moe_tp_down_rowwise", + "layers.*.mlp.vision_moe.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"]), @@ -100,6 +106,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 c970757313a3..660d6277eacd 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 @@ -100,9 +100,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 @@ -140,13 +140,19 @@ 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.text_moe.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.text_moe.experts.down_proj": "moe_tp_down_rowwise", + "layers.*.mlp.text_moe.experts": "moe_experts_allreduce", + "layers.*.mlp.vision_moe.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.vision_moe.experts.down_proj": "moe_tp_down_rowwise", + "layers.*.mlp.vision_moe.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", } 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..01642bf2be06 100644 --- a/src/transformers/models/eurobert/modeling_eurobert.py +++ b/src/transformers/models/eurobert/modeling_eurobert.py @@ -408,7 +408,11 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..c58246231e40 100644 --- a/src/transformers/models/eurobert/modular_eurobert.py +++ b/src/transformers/models/eurobert/modular_eurobert.py @@ -141,7 +141,11 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..4516901f11e3 100644 --- a/src/transformers/models/exaone4/modeling_exaone4.py +++ b/src/transformers/models/exaone4/modeling_exaone4.py @@ -440,8 +440,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..df164851667d 100644 --- a/src/transformers/models/exaone_moe/configuration_exaone_moe.py +++ b/src/transformers/models/exaone_moe/configuration_exaone_moe.py @@ -64,24 +64,36 @@ class ExaoneMoeConfig(PreTrainedConfig): model_type = "exaone_moe" keys_to_ignore_at_inference = ["past_key_values"] - # Default tensor parallel plan for base model `LlamaModel` + 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.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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_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..7f1fcd5fc177 100644 --- a/src/transformers/models/exaone_moe/modeling_exaone_moe.py +++ b/src/transformers/models/exaone_moe/modeling_exaone_moe.py @@ -563,8 +563,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..3e5dfda698ee 100644 --- a/src/transformers/models/exaone_moe/modular_exaone_moe.py +++ b/src/transformers/models/exaone_moe/modular_exaone_moe.py @@ -80,6 +80,30 @@ class ExaoneMoeConfig(Exaone4Config): >>> configuration = model.config ```""" + base_model_sp_plan = None + + 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_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + + 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..823adc24ab84 100644 --- a/src/transformers/models/falcon_h1/modeling_falcon_h1.py +++ b/src/transformers/models/falcon_h1/modeling_falcon_h1.py @@ -1166,8 +1166,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..a1dc3273f571 100644 --- a/src/transformers/models/flex_olmo/configuration_flex_olmo.py +++ b/src/transformers/models/flex_olmo/configuration_flex_olmo.py @@ -50,13 +50,29 @@ 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 + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "layers.*.mlp.experts": "moe_experts_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_experts_allreduce", + } + base_model_tp_ep_plan = { + "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": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -64,6 +80,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 aadd3c0fa1aa..6f3f8a4af84f 100644 --- a/src/transformers/models/flex_olmo/modeling_flex_olmo.py +++ b/src/transformers/models/flex_olmo/modeling_flex_olmo.py @@ -603,8 +603,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..ea11df43a532 100644 --- a/src/transformers/models/flex_olmo/modular_flex_olmo.py +++ b/src/transformers/models/flex_olmo/modular_flex_olmo.py @@ -60,13 +60,29 @@ 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 + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "layers.*.mlp.experts": "moe_experts_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_experts_allreduce", + } + base_model_tp_ep_plan = { + "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": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -74,6 +90,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..ad2b3cdcd751 100644 --- a/src/transformers/models/gemma/modeling_gemma.py +++ b/src/transformers/models/gemma/modeling_gemma.py @@ -450,8 +450,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..9a88238fd852 100644 --- a/src/transformers/models/gemma2/modeling_gemma2.py +++ b/src/transformers/models/gemma2/modeling_gemma2.py @@ -476,8 +476,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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 227c6a2b51d8..b786a390c7a9 100644 --- a/src/transformers/models/gemma3/modeling_gemma3.py +++ b/src/transformers/models/gemma3/modeling_gemma3.py @@ -585,8 +585,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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 99ee6c8b225f..13482907b5e8 100644 --- a/src/transformers/models/gemma3n/modeling_gemma3n.py +++ b/src/transformers/models/gemma3n/modeling_gemma3n.py @@ -1827,8 +1827,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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 10bdca384bd7..d39a7e779be5 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 66d1548fe453..a415cc1fc74b 100644 --- a/src/transformers/models/gemma4/configuration_gemma4.py +++ b/src/transformers/models/gemma4/configuration_gemma4.py @@ -125,28 +125,35 @@ 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) + # EP-only: shard routed experts; keep attention and dense MLP replicated. + "layers.*.router": "ep_router", + "layers.*.experts.gate_up_proj": "grouped_gemm", + "layers.*.experts.down_proj": "grouped_gemm", + "layers.*.experts": "moe_experts_allreduce", + } + base_model_tp_ep_plan = { + "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", "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 +161,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 @@ -242,12 +255,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 73887286dc22..92fdd704560e 100644 --- a/src/transformers/models/gemma4/modeling_gemma4.py +++ b/src/transformers/models/gemma4/modeling_gemma4.py @@ -1818,8 +1818,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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 88ecf590eb4f..177094e3cd47 100644 --- a/src/transformers/models/gemma4_assistant/modeling_gemma4_assistant.py +++ b/src/transformers/models/gemma4_assistant/modeling_gemma4_assistant.py @@ -110,6 +110,8 @@ 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"} + _tp_ep_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/gemma4_unified/modular_gemma4_unified.py b/src/transformers/models/gemma4_unified/modular_gemma4_unified.py index 17d2cd107c20..d3aa2283c59b 100644 --- a/src/transformers/models/gemma4_unified/modular_gemma4_unified.py +++ b/src/transformers/models/gemma4_unified/modular_gemma4_unified.py @@ -456,6 +456,8 @@ class Gemma4UnifiedTextConfig(Gemma4TextConfig): "layers.*.mlp.down_proj": "rowwise", } base_model_ep_plan = None # No MoE + base_model_tp_ep_plan = AttributeError() + base_model_fsdp_plan = AttributeError() sliding_window: int = 1024 use_bidirectional_attention: Literal["all", "vision"] | None = "vision" @@ -757,6 +759,12 @@ def forward( class Gemma4UnifiedForCausalLM(Gemma4ForCausalLM): + _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_ep_plan = AttributeError() + _sp_plan = AttributeError() + _sp_ep_plan = AttributeError() + _fsdp_plan = AttributeError() + def forward( self, input_ids: torch.LongTensor | None = None, diff --git a/src/transformers/models/gemma4_unified_assistant/modular_gemma4_unified_assistant.py b/src/transformers/models/gemma4_unified_assistant/modular_gemma4_unified_assistant.py index 9b9e619c2a53..c5638a637a6f 100644 --- a/src/transformers/models/gemma4_unified_assistant/modular_gemma4_unified_assistant.py +++ b/src/transformers/models/gemma4_unified_assistant/modular_gemma4_unified_assistant.py @@ -16,6 +16,9 @@ class Gemma4UnifiedAssistantForCausalLM(Gemma4AssistantForCausalLM): + _tp_ep_plan = AttributeError() + _fsdp_plan = AttributeError() + def forward(**super_kwargs): r""" shared_kv_states (`dict[str, torch.Tensor` of shape `(batch_size, 1, q_len, kv_len)`, *optional*): 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 58b8517a117d..3c17a1159bc6 100644 --- a/src/transformers/models/glm/modeling_glm.py +++ b/src/transformers/models/glm/modeling_glm.py @@ -438,8 +438,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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 9f4c5474c538..5f9883211880 100644 --- a/src/transformers/models/glm4/modeling_glm4.py +++ b/src/transformers/models/glm4/modeling_glm4.py @@ -443,8 +443,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..9856c3ae9dd0 100644 --- a/src/transformers/models/glm4_moe/configuration_glm4_moe.py +++ b/src/transformers/models/glm4_moe/configuration_glm4_moe.py @@ -57,22 +57,29 @@ 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.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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 53e1abf53c02..d312f1bd50bc 100644 --- a/src/transformers/models/glm4_moe/modeling_glm4_moe.py +++ b/src/transformers/models/glm4_moe/modeling_glm4_moe.py @@ -576,8 +576,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..1e7266ec206d 100644 --- a/src/transformers/models/glm4_moe/modular_glm4_moe.py +++ b/src/transformers/models/glm4_moe/modular_glm4_moe.py @@ -70,22 +70,29 @@ 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.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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..5ab913299ea1 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,27 @@ 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.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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 e0d7b8fddb14..3002cbd66fbb 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 @@ -652,8 +652,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..44a369cc8205 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,27 @@ 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.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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 ba9f0c31fe64..d3d38123e1ea 100644 --- a/src/transformers/models/glm4v/modular_glm4v.py +++ b/src/transformers/models/glm4v/modular_glm4v.py @@ -135,15 +135,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..67399f1868bc 100644 --- a/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py +++ b/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py @@ -56,16 +56,29 @@ 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.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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_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 +111,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..a6515ccd88c8 100644 --- a/src/transformers/models/glm4v_moe/modular_glm4v_moe.py +++ b/src/transformers/models/glm4v_moe/modular_glm4v_moe.py @@ -88,16 +88,29 @@ 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.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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_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 c5df188ba99c..03a68e1292ef 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 5ecf328bd170..0b2806cd0c65 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 @@ -64,25 +64,23 @@ class GlmMoeDsaConfig(PreTrainedConfig, RotaryEmbeddingConfigMixin): 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.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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"]), } - attribute_map = { "num_local_experts": "n_routed_experts", } @@ -127,6 +125,12 @@ class GlmMoeDsaConfig(PreTrainedConfig, RotaryEmbeddingConfigMixin): head_dim: int = 64 first_k_dense_replace: int = 3 layer_types: list[str] | None = None + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } # `"full"` runs the indexer, `"shared"` reuses the previous full layer's index mask. indexer_types: list[str] | None = None 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 34d7bf139c36..dbf6fc56b55a 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 @@ -781,8 +781,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"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _sp_plan = {"lm_head": "colwise_loss_parallel"} + _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 a4e91beb4633..a618a6d4b6f9 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 @@ -78,6 +78,31 @@ class GlmMoeDsaConfig(DeepseekV32Config): >>> configuration = model.config ```""" + base_model_tp_plan = { + "layers.*.self_attn.q_b_proj": "colwise", + "layers.*.self_attn.kv_b_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "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", } @@ -333,7 +358,9 @@ def forward( class GlmMoeDsaForCausalLM(DeepseekV32ForCausalLM): - pass + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} + _fsdp_plan = {"lm_head": "keep_full_weight"} __all__ = [ 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 06f389cb8a72..6dc305c289f0 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 3530887574f6..d4e95c84503c 100644 --- a/src/transformers/models/glmasr/modeling_glmasr.py +++ b/src/transformers/models/glmasr/modeling_glmasr.py @@ -371,6 +371,7 @@ class GlmAsrModelOutputWithPast(BaseModelOutputWithPast): ) class GlmAsrModel(GlmAsrPreTrainedModel): _tp_plan = None + _sp_plan = None _pp_plan = None _keep_in_fp32_modules_strict = None 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..09e98780ca9d 100755 --- a/src/transformers/models/gpt_neox/modeling_gpt_neox.py +++ b/src/transformers/models/gpt_neox/modeling_gpt_neox.py @@ -394,7 +394,8 @@ 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"} + _tp_ep_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..2d8d31deb650 100644 --- a/src/transformers/models/gpt_neox/modular_gpt_neox.py +++ b/src/transformers/models/gpt_neox/modular_gpt_neox.py @@ -341,7 +341,8 @@ 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"} + _tp_ep_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 8ba86ffc5802..d21f8b45c69b 100644 --- a/src/transformers/models/gpt_oss/modeling_gpt_oss.py +++ b/src/transformers/models/gpt_oss/modeling_gpt_oss.py @@ -592,8 +592,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..e9fad65eabd5 100644 --- a/src/transformers/models/granite/modeling_granite.py +++ b/src/transformers/models/granite/modeling_granite.py @@ -445,8 +445,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..dae4348efbcf 100644 --- a/src/transformers/models/granitemoe/modeling_granitemoe.py +++ b/src/transformers/models/granitemoe/modeling_granitemoe.py @@ -626,8 +626,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..fe927faa79d0 100644 --- a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py +++ b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py @@ -1307,8 +1307,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..b6f2884fec46 100644 --- a/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py +++ b/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py @@ -695,8 +695,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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 7d22c79095f8..6c5a91b0725e 100644 --- a/src/transformers/models/helium/modeling_helium.py +++ b/src/transformers/models/helium/modeling_helium.py @@ -421,8 +421,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..24d757afd551 100644 --- a/src/transformers/models/hrm_text/modeling_hrm_text.py +++ b/src/transformers/models/hrm_text/modeling_hrm_text.py @@ -555,8 +555,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..88e79b690295 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,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..7eea7af92237 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,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..7335ba8d7a78 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.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "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": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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..c335daf46af1 100644 --- a/src/transformers/models/hy_v3/modeling_hy_v3.py +++ b/src/transformers/models/hy_v3/modeling_hy_v3.py @@ -544,8 +544,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..c109e8973aa5 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.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "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": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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..a5420a95e290 100644 --- a/src/transformers/models/hyperclovax/modeling_hyperclovax.py +++ b/src/transformers/models/hyperclovax/modeling_hyperclovax.py @@ -452,8 +452,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..b74f8acca655 100644 --- a/src/transformers/models/jais2/modeling_jais2.py +++ b/src/transformers/models/jais2/modeling_jais2.py @@ -397,8 +397,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..ae303a3aec07 100755 --- a/src/transformers/models/jamba/modeling_jamba.py +++ b/src/transformers/models/jamba/modeling_jamba.py @@ -843,8 +843,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..22dd63b3a71d 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,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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 6dae15bd1970..96eef65c3b1a 100644 --- a/src/transformers/models/laguna/configuration_laguna.py +++ b/src/transformers/models/laguna/configuration_laguna.py @@ -65,18 +65,44 @@ 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.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.down_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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", + } + # Expert-only EP plan: shards MoE experts along the expert axis; gate stays replicated. + 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_experts_allreduce", + } + base_model_tp_ep_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.g_proj": "colwise", + "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_allreduce", + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "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_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -84,6 +110,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 828745c348c3..b6a63f11d8bd 100644 --- a/src/transformers/models/laguna/modeling_laguna.py +++ b/src/transformers/models/laguna/modeling_laguna.py @@ -678,8 +678,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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 659cefb8a9ce..039abdc91e18 100644 --- a/src/transformers/models/laguna/modular_laguna.py +++ b/src/transformers/models/laguna/modular_laguna.py @@ -83,18 +83,37 @@ 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.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.down_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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_tp_ep_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.g_proj": "colwise", + "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_allreduce", + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "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_allreduce", } vocab_size: int = 100352 diff --git a/src/transformers/models/lfm2/modeling_lfm2.py b/src/transformers/models/lfm2/modeling_lfm2.py index b03df7583544..5e75dbed8def 100644 --- a/src/transformers/models/lfm2/modeling_lfm2.py +++ b/src/transformers/models/lfm2/modeling_lfm2.py @@ -559,8 +559,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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 159e82f95147..2e629902da46 100644 --- a/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py +++ b/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py @@ -649,8 +649,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..7cd78fce37f5 100644 --- a/src/transformers/models/llama/modeling_llama.py +++ b/src/transformers/models/llama/modeling_llama.py @@ -428,8 +428,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..2612a236d0eb 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,41 @@ 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.experts": "moe_experts_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", + "layers.*.feed_forward.router": "ep_router", + } + base_model_tp_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_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_allreduce", + "layers.*.feed_forward.gate_proj": "colwise", + "layers.*.feed_forward.up_proj": "colwise", + "layers.*.feed_forward.down_proj": "rowwise_allreduce", + "layers.*.feed_forward.experts.gate_up_proj": "grouped_gemm", + "layers.*.feed_forward.experts.down_proj": "grouped_gemm", + "layers.*.feed_forward.experts": "moe_experts_allreduce", "layers.*.feed_forward.router": "ep_router", } @@ -179,7 +194,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 3ee103f54dc3..ab0e5f3e42f0 100644 --- a/src/transformers/models/llama4/modeling_llama4.py +++ b/src/transformers/models/llama4/modeling_llama4.py @@ -589,7 +589,11 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..45fa53fc824d 100644 --- a/src/transformers/models/longcat_flash/configuration_longcat_flash.py +++ b/src/transformers/models/longcat_flash/configuration_longcat_flash.py @@ -55,16 +55,15 @@ 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.*.self_attn.*.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts.identity_expert": "moe_identity_expert", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts": "moe_experts_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 +72,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 4a287fde6e02..b4fe7f92c4f7 100644 --- a/src/transformers/models/longcat_flash/modeling_longcat_flash.py +++ b/src/transformers/models/longcat_flash/modeling_longcat_flash.py @@ -644,8 +644,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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/mellum/configuration_mellum.py b/src/transformers/models/mellum/configuration_mellum.py index 49deff0e2e69..4da68e0cf56c 100644 --- a/src/transformers/models/mellum/configuration_mellum.py +++ b/src/transformers/models/mellum/configuration_mellum.py @@ -54,15 +54,28 @@ class MellumConfig(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.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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", + } + + # Per-layer MLP entries are filled in by `_update_sp_plan` (dense vs sparse). + 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", + "norm": "activation", } # Expert-only EP plan: only shards MoE experts, not attention. # Attention is left unsharded — FSDP2 handles attention weight distribution. @@ -71,7 +84,43 @@ class MellumConfig(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", + } + # Inference TP + EP (per-layer overrides applied in `_update_parallel_plans`). + base_model_tp_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_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + # Training SP + EP (per-layer overrides applied in `_update_parallel_plans`). + base_model_sp_ep_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.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "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_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -79,6 +128,12 @@ class MellumConfig(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 = 98304 hidden_size: int = 2304 intermediate_size: int = 7168 @@ -120,11 +175,88 @@ def __post_init__(self, **kwargs): "sliding_attention": {"rope_type": "default", "rope_theta": 10000.0}, } + self._update_parallel_plans() + super().__post_init__( **kwargs, ignore_keys_at_rope_validation={"sliding_attention", "full_attention"}, ) + def _is_moe_layer(self, layer_idx: int) -> bool: + return self.mlp_layer_types[layer_idx] == "sparse" + + def _update_parallel_plans(self): + self._update_sp_plan() + self._update_sp_ep_plan() + self._update_tp_ep_plan() + + def _update_sp_plan(self): + """Set per-layer SP entries depending on whether the MLP is dense or sparse.""" + self.base_model_sp_plan = self.base_model_sp_plan.copy() + for i, mlp_type in enumerate(self.mlp_layer_types): + if mlp_type == "dense": + self.base_model_sp_plan.update( + { + f"layers.{i}.mlp": "module_allgather", + f"layers.{i}.mlp.gate_proj": "colwise", + f"layers.{i}.mlp.up_proj": "colwise", + f"layers.{i}.mlp.down_proj": "rowwise_reduce_scatter", + } + ) + else: + self.base_model_sp_plan.update( + { + f"layers.{i}.mlp": "module_allgather_split", + f"layers.{i}.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + f"layers.{i}.mlp.experts.down_proj": "moe_tp_down_rowwise", + f"layers.{i}.mlp.experts": "moe_experts_allreduce", + } + ) + + def _update_sp_ep_plan(self): + self.base_model_sp_ep_plan = self.base_model_sp_ep_plan.copy() + for i, mlp_type in enumerate(self.mlp_layer_types): + if mlp_type == "dense": + self.base_model_sp_ep_plan.update( + { + f"layers.{i}.mlp": "module_allgather", + f"layers.{i}.mlp.gate_proj": "colwise", + f"layers.{i}.mlp.up_proj": "colwise", + f"layers.{i}.mlp.down_proj": "rowwise_reduce_scatter", + } + ) + else: + self.base_model_sp_ep_plan.update( + { + f"layers.{i}.mlp": "module_allgather_split", + f"layers.{i}.mlp.gate": "ep_router", + f"layers.{i}.mlp.experts.gate_up_proj": "grouped_gemm", + f"layers.{i}.mlp.experts.down_proj": "grouped_gemm", + f"layers.{i}.mlp.experts": "moe_experts_allreduce", + } + ) + + def _update_tp_ep_plan(self): + self.base_model_tp_ep_plan = self.base_model_tp_ep_plan.copy() + for i, mlp_type in enumerate(self.mlp_layer_types): + if mlp_type == "dense": + self.base_model_tp_ep_plan.update( + { + f"layers.{i}.mlp.gate_proj": "colwise", + f"layers.{i}.mlp.up_proj": "colwise", + f"layers.{i}.mlp.down_proj": "rowwise_allreduce", + } + ) + else: + self.base_model_tp_ep_plan.update( + { + f"layers.{i}.mlp.gate": "ep_router", + f"layers.{i}.mlp.experts.gate_up_proj": "grouped_gemm", + f"layers.{i}.mlp.experts.down_proj": "grouped_gemm", + f"layers.{i}.mlp.experts": "moe_experts_allreduce", + } + ) + def convert_rope_params_to_dict(self, **kwargs): # No need to handle BC for new models, because they have no old-format `rope_scaling` return kwargs diff --git a/src/transformers/models/mellum/modeling_mellum.py b/src/transformers/models/mellum/modeling_mellum.py index bc206be988ec..db04cfa0ac9a 100644 --- a/src/transformers/models/mellum/modeling_mellum.py +++ b/src/transformers/models/mellum/modeling_mellum.py @@ -633,8 +633,12 @@ def load_balancing_loss_func( @auto_docstring class MellumForCausalLM(MellumPreTrainedModel, 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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/mellum/modular_mellum.py b/src/transformers/models/mellum/modular_mellum.py index 889b4790088c..da08ca93e687 100644 --- a/src/transformers/models/mellum/modular_mellum.py +++ b/src/transformers/models/mellum/modular_mellum.py @@ -54,6 +54,57 @@ class MellumConfig(Qwen3MoeConfig): model_type = "mellum" + # Per-layer MLP entries are filled in by `_update_sp_plan` (dense vs sparse). + 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", + "norm": "activation", + } + # Inference TP + EP (per-layer overrides applied in `_update_parallel_plans`). + base_model_tp_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_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + # Training SP + EP (per-layer overrides applied in `_update_parallel_plans`). + base_model_sp_ep_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.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "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", + } + vocab_size: int = 98304 hidden_size: int = 2304 intermediate_size: int = 7168 @@ -84,12 +135,89 @@ def __post_init__(self, **kwargs): "sliding_attention": {"rope_type": "default", "rope_theta": 10000.0}, } + self._update_parallel_plans() + PreTrainedConfig.__post_init__( self, **kwargs, ignore_keys_at_rope_validation={"sliding_attention", "full_attention"}, ) + def _is_moe_layer(self, layer_idx: int) -> bool: + return self.mlp_layer_types[layer_idx] == "sparse" + + def _update_parallel_plans(self): + self._update_sp_plan() + self._update_sp_ep_plan() + self._update_tp_ep_plan() + + def _update_sp_plan(self): + """Set per-layer SP entries depending on whether the MLP is dense or sparse.""" + self.base_model_sp_plan = self.base_model_sp_plan.copy() + for i, mlp_type in enumerate(self.mlp_layer_types): + if mlp_type == "dense": + self.base_model_sp_plan.update( + { + f"layers.{i}.mlp": "module_allgather", + f"layers.{i}.mlp.gate_proj": "colwise", + f"layers.{i}.mlp.up_proj": "colwise", + f"layers.{i}.mlp.down_proj": "rowwise_reduce_scatter", + } + ) + else: + self.base_model_sp_plan.update( + { + f"layers.{i}.mlp": "module_allgather_split", + f"layers.{i}.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + f"layers.{i}.mlp.experts.down_proj": "moe_tp_down_rowwise", + f"layers.{i}.mlp.experts": "moe_experts_allreduce", + } + ) + + def _update_sp_ep_plan(self): + self.base_model_sp_ep_plan = self.base_model_sp_ep_plan.copy() + for i, mlp_type in enumerate(self.mlp_layer_types): + if mlp_type == "dense": + self.base_model_sp_ep_plan.update( + { + f"layers.{i}.mlp": "module_allgather", + f"layers.{i}.mlp.gate_proj": "colwise", + f"layers.{i}.mlp.up_proj": "colwise", + f"layers.{i}.mlp.down_proj": "rowwise_reduce_scatter", + } + ) + else: + self.base_model_sp_ep_plan.update( + { + f"layers.{i}.mlp": "module_allgather_split", + f"layers.{i}.mlp.gate": "ep_router", + f"layers.{i}.mlp.experts.gate_up_proj": "grouped_gemm", + f"layers.{i}.mlp.experts.down_proj": "grouped_gemm", + f"layers.{i}.mlp.experts": "moe_experts_allreduce", + } + ) + + def _update_tp_ep_plan(self): + self.base_model_tp_ep_plan = self.base_model_tp_ep_plan.copy() + for i, mlp_type in enumerate(self.mlp_layer_types): + if mlp_type == "dense": + self.base_model_tp_ep_plan.update( + { + f"layers.{i}.mlp.gate_proj": "colwise", + f"layers.{i}.mlp.up_proj": "colwise", + f"layers.{i}.mlp.down_proj": "rowwise_allreduce", + } + ) + else: + self.base_model_tp_ep_plan.update( + { + f"layers.{i}.mlp.gate": "ep_router", + f"layers.{i}.mlp.experts.gate_up_proj": "grouped_gemm", + f"layers.{i}.mlp.experts.down_proj": "grouped_gemm", + f"layers.{i}.mlp.experts": "moe_experts_allreduce", + } + ) + def convert_rope_params_to_dict(self, **kwargs): # No need to handle BC for new models, because they have no old-format `rope_scaling` return kwargs diff --git a/src/transformers/models/minimax/configuration_minimax.py b/src/transformers/models/minimax/configuration_minimax.py index 9a3a0023725e..8ff7d9a6d3db 100644 --- a/src/transformers/models/minimax/configuration_minimax.py +++ b/src/transformers/models/minimax/configuration_minimax.py @@ -65,16 +65,23 @@ 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.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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..0f6392c9976a 100644 --- a/src/transformers/models/minimax/modeling_minimax.py +++ b/src/transformers/models/minimax/modeling_minimax.py @@ -789,8 +789,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..6f9ec2655a58 100644 --- a/src/transformers/models/minimax/modular_minimax.py +++ b/src/transformers/models/minimax/modular_minimax.py @@ -91,16 +91,23 @@ 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.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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..9adab09b43cf 100644 --- a/src/transformers/models/minimax_m2/configuration_minimax_m2.py +++ b/src/transformers/models/minimax_m2/configuration_minimax_m2.py @@ -48,19 +48,41 @@ 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.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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 08b4f8117f60..b14efdb3ceb5 100644 --- a/src/transformers/models/minimax_m2/modeling_minimax_m2.py +++ b/src/transformers/models/minimax_m2/modeling_minimax_m2.py @@ -587,8 +587,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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 e71774ef59fd..b5917b630669 100644 --- a/src/transformers/models/minimax_m2/modular_minimax_m2.py +++ b/src/transformers/models/minimax_m2/modular_minimax_m2.py @@ -67,19 +67,41 @@ 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.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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_m3_vl/modular_minimax_m3_vl.py b/src/transformers/models/minimax_m3_vl/modular_minimax_m3_vl.py index 87d1589a501b..10384c82f6c3 100644 --- a/src/transformers/models/minimax_m3_vl/modular_minimax_m3_vl.py +++ b/src/transformers/models/minimax_m3_vl/modular_minimax_m3_vl.py @@ -93,6 +93,17 @@ class MiniMaxM3VLTextConfig(MiniMaxM2Config): model_type = "minimax_m3_vl_text" base_config_key = "text_config" + 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", + } + base_model_sp_plan = AttributeError() + base_model_fsdp_plan = AttributeError() base_model_ep_plan = { "layers.*.mlp.gate": "ep_router", "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", @@ -101,6 +112,14 @@ class MiniMaxM3VLTextConfig(MiniMaxM2Config): "layers.*.mlp.experts.down_proj_scale_inv": "grouped_gemm", "layers.*.mlp.experts": "moe_tp_experts", } + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), + } + attribute_map = { + "num_experts": "num_local_experts", + } hidden_size: int = 6144 intermediate_size: int = 3072 @@ -718,6 +737,11 @@ def forward( class MiniMaxM3VLForCausalLM(MiniMaxM2ForCausalLM): config: MiniMaxM3VLTextConfig _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_ep_plan = AttributeError() + _sp_plan = AttributeError() + _sp_ep_plan = AttributeError() + _fsdp_plan = AttributeError() def __init__(self, config: MiniMaxM3VLTextConfig): super().__init__(config) 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..0e1d344bf25f 100644 --- a/src/transformers/models/ministral/modeling_ministral.py +++ b/src/transformers/models/ministral/modeling_ministral.py @@ -429,8 +429,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..8b7dfec23fbe 100644 --- a/src/transformers/models/ministral3/modeling_ministral3.py +++ b/src/transformers/models/ministral3/modeling_ministral3.py @@ -412,8 +412,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..5fa28011ccfa 100644 --- a/src/transformers/models/mistral/modeling_mistral.py +++ b/src/transformers/models/mistral/modeling_mistral.py @@ -401,8 +401,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..9979792faffe 100644 --- a/src/transformers/models/mistral4/configuration_mistral4.py +++ b/src/transformers/models/mistral4/configuration_mistral4.py @@ -47,21 +47,27 @@ 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.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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 23bfe1091b69..10eecfa312ea 100644 --- a/src/transformers/models/mistral4/modeling_mistral4.py +++ b/src/transformers/models/mistral4/modeling_mistral4.py @@ -641,8 +641,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..717e672afe00 100644 --- a/src/transformers/models/mixtral/configuration_mixtral.py +++ b/src/transformers/models/mixtral/configuration_mixtral.py @@ -42,20 +42,80 @@ 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.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "layers.*.mlp.experts": "moe_experts_allreduce", + "norm": "activation", + } + # Expert-only EP plan: shards MoE experts along the expert axis; gate stays replicated. + 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_experts_allreduce", + } + # Inference TP + EP (complete plan — no runtime merge). + base_model_tp_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_allreduce", + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + # Training SP + EP (complete plan — no runtime merge). + base_model_sp_ep_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.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "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..b0947d3f4271 100644 --- a/src/transformers/models/mixtral/modeling_mixtral.py +++ b/src/transformers/models/mixtral/modeling_mixtral.py @@ -580,8 +580,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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 6d08da835303..a262d6c77112 100644 --- a/src/transformers/models/musicflamingo/modeling_musicflamingo.py +++ b/src/transformers/models/musicflamingo/modeling_musicflamingo.py @@ -213,6 +213,7 @@ def apply_rotary_time_emb(hidden_states, cos, sin): ) class MusicFlamingoModel(MusicFlamingoPreTrainedModel): _tp_plan = None + _sp_plan = None _pp_plan = None _keep_in_fp32_modules_strict = None 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..48011dac3736 100644 --- a/src/transformers/models/nanochat/modeling_nanochat.py +++ b/src/transformers/models/nanochat/modeling_nanochat.py @@ -432,8 +432,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..6c6f79ae5750 100644 --- a/src/transformers/models/nanochat/modular_nanochat.py +++ b/src/transformers/models/nanochat/modular_nanochat.py @@ -198,7 +198,11 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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 b1e250db58a5..df5a149eaf91 100644 --- a/src/transformers/models/olmo/modeling_olmo.py +++ b/src/transformers/models/olmo/modeling_olmo.py @@ -426,8 +426,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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 d8def37b37d8..b03d73d36929 100644 --- a/src/transformers/models/olmo2/modeling_olmo2.py +++ b/src/transformers/models/olmo2/modeling_olmo2.py @@ -430,8 +430,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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 d167ebe15951..4d765e7d56bd 100644 --- a/src/transformers/models/olmo3/modeling_olmo3.py +++ b/src/transformers/models/olmo3/modeling_olmo3.py @@ -434,8 +434,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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 18a8cff6fe5d..250c2f65b425 100644 --- a/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py +++ b/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py @@ -1046,8 +1046,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..01cc17a429d7 100644 --- a/src/transformers/models/olmoe/configuration_olmoe.py +++ b/src/transformers/models/olmoe/configuration_olmoe.py @@ -46,13 +46,67 @@ 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.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "layers.*.mlp.experts": "moe_experts_allreduce", + "norm": "activation", + } + # Expert-only EP plan: shards MoE experts along the expert axis; gate stays replicated. + 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_experts_allreduce", + } + base_model_tp_ep_plan = { + "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": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + base_model_sp_ep_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.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "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..d25bdd6f5282 100644 --- a/src/transformers/models/olmoe/modeling_olmoe.py +++ b/src/transformers/models/olmoe/modeling_olmoe.py @@ -604,8 +604,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..12317fbfc3cd 100644 --- a/src/transformers/models/phi/modeling_phi.py +++ b/src/transformers/models/phi/modeling_phi.py @@ -406,8 +406,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..7ea8ccfcb241 100644 --- a/src/transformers/models/phi3/modeling_phi3.py +++ b/src/transformers/models/phi3/modeling_phi3.py @@ -432,8 +432,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..99003339fed4 100644 --- a/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py +++ b/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py @@ -1596,8 +1596,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..753b2535c1e4 100644 --- a/src/transformers/models/phimoe/modeling_phimoe.py +++ b/src/transformers/models/phimoe/modeling_phimoe.py @@ -772,8 +772,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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 575620c870fe..ff3055f47d29 100644 --- a/src/transformers/models/pi0/modeling_pi0.py +++ b/src/transformers/models/pi0/modeling_pi0.py @@ -219,7 +219,8 @@ 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"} + _tp_ep_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 2bc69237d13c..968844bfb7a4 100644 --- a/src/transformers/models/pi0/modular_pi0.py +++ b/src/transformers/models/pi0/modular_pi0.py @@ -479,7 +479,8 @@ 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"} + _tp_ep_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..ded1903f822f 100644 --- a/src/transformers/models/qwen2/modeling_qwen2.py +++ b/src/transformers/models/qwen2/modeling_qwen2.py @@ -416,8 +416,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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 97587d4611fc..a9170e0b0a80 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..b000ad0942f3 100644 --- a/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py +++ b/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py @@ -57,10 +57,33 @@ 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.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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", + } + # Expert-only EP plan: shards MoE experts along the expert axis; gate stays replicated. + 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_experts_allreduce", + } + base_model_tp_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_allreduce", + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -68,6 +91,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..3303d94c4848 100644 --- a/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py +++ b/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py @@ -617,8 +617,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..84d55ae84b47 100644 --- a/src/transformers/models/qwen2_moe/modular_qwen2_moe.py +++ b/src/transformers/models/qwen2_moe/modular_qwen2_moe.py @@ -230,7 +230,11 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..9e573bccbd7f 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; runtime flags (cpu_offload, mixed_precision) + # live on DistributedConfig. 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..fe3de31fcac7 100644 --- a/src/transformers/models/qwen3/modeling_qwen3.py +++ b/src/transformers/models/qwen3/modeling_qwen3.py @@ -441,8 +441,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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 13551c377991..7929156952f6 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 21a8032e34f4..d113c4c99dd9 100644 --- a/src/transformers/models/qwen3_5/modeling_qwen3_5.py +++ b/src/transformers/models/qwen3_5/modeling_qwen3_5.py @@ -1614,8 +1614,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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 6a49a4e8f602..4dd857bc553e 100644 --- a/src/transformers/models/qwen3_5/modular_qwen3_5.py +++ b/src/transformers/models/qwen3_5/modular_qwen3_5.py @@ -97,12 +97,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 4a32828987ed..ee125e10a6e9 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,13 @@ 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.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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 +74,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 8249194c7f61..87cc297c4e94 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 @@ -1800,8 +1800,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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.*"] @@ -1906,7 +1910,12 @@ class Qwen3_5MoeForConditionalGeneration(Qwen3_5MoePreTrainedModel, GenerationMi _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"} # Reference: fix gemma3 grad acc #37208 accepts_loss_kwargs = False - _tp_plan = {"lm_head": "colwise_gather_output"} + config: Qwen3_5MoeConfig + _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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 4c6d8259e74e..361ffe18fc57 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 @@ -90,15 +90,13 @@ 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.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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"} @@ -252,7 +250,12 @@ def __init__(self, config): class Qwen3_5MoeForConditionalGeneration(Qwen3VLMoeForConditionalGeneration): - _tp_plan = {"lm_head": "colwise_gather_output"} + config: Qwen3_5MoeConfig + _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..21893966542c 100644 --- a/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py +++ b/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py @@ -57,15 +57,31 @@ 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.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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 +90,43 @@ 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", + } + # Inference TP + EP (per-layer overrides applied in `_update_parallel_plans`). + base_model_tp_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_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + # Training SP + EP (per-layer overrides applied in `_update_parallel_plans`). + base_model_sp_ep_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.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "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_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -82,6 +134,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 @@ -115,6 +173,86 @@ def __post_init__(self, **kwargs): self.sliding_window = self.sliding_window if self.use_sliding_window else None self.mlp_only_layers = [] if self.mlp_only_layers is None else self.mlp_only_layers super().__post_init__(**kwargs) + self._update_parallel_plans() + + def _is_moe_layer(self, layer_idx: int) -> bool: + return (layer_idx not in self.mlp_only_layers) and ( + self.num_experts > 0 and (layer_idx + 1) % self.decoder_sparse_step == 0 + ) + + def _update_parallel_plans(self): + self._update_sp_plan() + self._update_sp_ep_plan() + self._update_tp_ep_plan() + + def _update_sp_plan(self): + self.base_model_sp_plan = self.base_model_sp_plan.copy() + + for i in range(self.num_hidden_layers): + if self._is_moe_layer(i): + self.base_model_sp_plan.update( + { + f"layers.{i}.mlp": "module_allgather_split", + f"layers.{i}.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + f"layers.{i}.mlp.experts.down_proj": "moe_tp_down_rowwise", + f"layers.{i}.mlp.experts": "moe_experts_allreduce", + } + ) + else: + self.base_model_sp_plan.update( + { + f"layers.{i}.mlp": "module_allgather", + f"layers.{i}.mlp.gate_proj": "colwise", + f"layers.{i}.mlp.up_proj": "colwise", + f"layers.{i}.mlp.down_proj": "rowwise_reduce_scatter", + } + ) + + def _update_sp_ep_plan(self): + self.base_model_sp_ep_plan = self.base_model_sp_ep_plan.copy() + + for i in range(self.num_hidden_layers): + if self._is_moe_layer(i): + self.base_model_sp_ep_plan.update( + { + f"layers.{i}.mlp": "module_allgather_split", + f"layers.{i}.mlp.gate": "ep_router", + f"layers.{i}.mlp.experts.gate_up_proj": "grouped_gemm", + f"layers.{i}.mlp.experts.down_proj": "grouped_gemm", + f"layers.{i}.mlp.experts": "moe_experts_allreduce", + } + ) + else: + self.base_model_sp_ep_plan.update( + { + f"layers.{i}.mlp": "module_allgather", + f"layers.{i}.mlp.gate_proj": "colwise", + f"layers.{i}.mlp.up_proj": "colwise", + f"layers.{i}.mlp.down_proj": "rowwise_reduce_scatter", + } + ) + + def _update_tp_ep_plan(self): + self.base_model_tp_ep_plan = self.base_model_tp_ep_plan.copy() + + for i in range(self.num_hidden_layers): + if self._is_moe_layer(i): + self.base_model_tp_ep_plan.update( + { + f"layers.{i}.mlp.gate": "ep_router", + f"layers.{i}.mlp.experts.gate_up_proj": "grouped_gemm", + f"layers.{i}.mlp.experts.down_proj": "grouped_gemm", + f"layers.{i}.mlp.experts": "moe_experts_allreduce", + } + ) + else: + self.base_model_tp_ep_plan.update( + { + f"layers.{i}.mlp.gate_proj": "colwise", + f"layers.{i}.mlp.up_proj": "colwise", + f"layers.{i}.mlp.down_proj": "rowwise_allreduce", + } + ) __all__ = ["Qwen3MoeConfig"] diff --git a/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py b/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py index ddf84fc575b7..4b71fa5bbc5b 100644 --- a/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py +++ b/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py @@ -609,8 +609,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..f09c965d1093 100644 --- a/src/transformers/models/qwen3_next/configuration_qwen3_next.py +++ b/src/transformers/models/qwen3_next/configuration_qwen3_next.py @@ -62,18 +62,16 @@ 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.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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 +79,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 b4cafa6bb089..5ef52f8227ba 100644 --- a/src/transformers/models/qwen3_next/modeling_qwen3_next.py +++ b/src/transformers/models/qwen3_next/modeling_qwen3_next.py @@ -1087,8 +1087,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..d52c179d58ff 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; runtime flags (cpu_offload, mixed_precision) + # live on DistributedConfig. 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,73 @@ 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.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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", } 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", + } + # Inference TP + EP (per-layer overrides applied in `_update_parallel_plans`). + base_model_tp_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_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + # Training SP + EP (per-layer overrides applied in `_update_parallel_plans`). + base_model_sp_ep_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.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "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_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -379,6 +469,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 @@ -409,8 +505,89 @@ class Qwen3OmniMoeTalkerTextConfig(PreTrainedConfig): def __post_init__(self, **kwargs): self.sliding_window = self.sliding_window + self.sliding_window = self.sliding_window if self.use_sliding_window else None self.mlp_only_layers = [] if self.mlp_only_layers is None else self.mlp_only_layers super().__post_init__(**kwargs) + self._update_parallel_plans() + + def _is_moe_layer(self, layer_idx: int) -> bool: + return (layer_idx not in self.mlp_only_layers) and ( + self.num_experts > 0 and (layer_idx + 1) % self.decoder_sparse_step == 0 + ) + + def _update_parallel_plans(self): + self._update_sp_plan() + self._update_sp_ep_plan() + self._update_tp_ep_plan() + + def _update_sp_plan(self): + self.base_model_sp_plan = self.base_model_sp_plan.copy() + + for i in range(self.num_hidden_layers): + if self._is_moe_layer(i): + self.base_model_sp_plan.update( + { + f"layers.{i}.mlp": "module_allgather_split", + f"layers.{i}.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + f"layers.{i}.mlp.experts.down_proj": "moe_tp_down_rowwise", + f"layers.{i}.mlp.experts": "moe_experts_allreduce", + } + ) + else: + self.base_model_sp_plan.update( + { + f"layers.{i}.mlp": "module_allgather", + f"layers.{i}.mlp.gate_proj": "colwise", + f"layers.{i}.mlp.up_proj": "colwise", + f"layers.{i}.mlp.down_proj": "rowwise_reduce_scatter", + } + ) + + def _update_sp_ep_plan(self): + self.base_model_sp_ep_plan = self.base_model_sp_ep_plan.copy() + + for i in range(self.num_hidden_layers): + if self._is_moe_layer(i): + self.base_model_sp_ep_plan.update( + { + f"layers.{i}.mlp": "module_allgather_split", + f"layers.{i}.mlp.gate": "ep_router", + f"layers.{i}.mlp.experts.gate_up_proj": "grouped_gemm", + f"layers.{i}.mlp.experts.down_proj": "grouped_gemm", + f"layers.{i}.mlp.experts": "moe_experts_allreduce", + } + ) + else: + self.base_model_sp_ep_plan.update( + { + f"layers.{i}.mlp": "module_allgather", + f"layers.{i}.mlp.gate_proj": "colwise", + f"layers.{i}.mlp.up_proj": "colwise", + f"layers.{i}.mlp.down_proj": "rowwise_reduce_scatter", + } + ) + + def _update_tp_ep_plan(self): + self.base_model_tp_ep_plan = self.base_model_tp_ep_plan.copy() + + for i in range(self.num_hidden_layers): + if self._is_moe_layer(i): + self.base_model_tp_ep_plan.update( + { + f"layers.{i}.mlp.gate": "ep_router", + f"layers.{i}.mlp.experts.gate_up_proj": "grouped_gemm", + f"layers.{i}.mlp.experts.down_proj": "grouped_gemm", + f"layers.{i}.mlp.experts": "moe_experts_allreduce", + } + ) + else: + self.base_model_tp_ep_plan.update( + { + f"layers.{i}.mlp.gate_proj": "colwise", + f"layers.{i}.mlp.up_proj": "colwise", + f"layers.{i}.mlp.down_proj": "rowwise_allreduce", + } + ) @auto_docstring(checkpoint="Qwen/Qwen3-30B-A3B-Base") 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 b4076e96fae1..a6099aa8df41 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 @@ -2632,8 +2632,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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 = { @@ -3016,8 +3020,12 @@ 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"} + _tp_ep_plan = {"codec_head": "colwise_allgather"} + _sp_ep_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 a4c9d884a173..5d1ab7fe5a62 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,43 @@ 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", + } + # Inference TP + EP (per-layer overrides applied in `_update_parallel_plans`). + base_model_tp_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_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + # Training SP + EP (per-layer overrides applied in `_update_parallel_plans`). + base_model_sp_ep_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.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "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", } vocab_size: int = 3072 @@ -408,8 +449,8 @@ class Qwen3OmniMoeTalkerTextConfig(Qwen3MoeConfig): use_sliding_window = AttributeError() def __post_init__(self, **kwargs): - super().__post_init__(**kwargs) self.sliding_window = self.sliding_window + super().__post_init__(**kwargs) @auto_docstring(checkpoint="Qwen/Qwen3-30B-A3B-Base") @@ -1659,8 +1700,12 @@ 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"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"codec_head": "colwise_allgather"} + _sp_ep_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_vl_moe/configuration_qwen3_vl_moe.py b/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py index aaa6b2e52ef8..00eb23290288 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,70 @@ 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", + } + # Inference TP + EP (per-layer overrides applied in `_update_parallel_plans`). + base_model_tp_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_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + # Training SP + EP (per-layer overrides applied in `_update_parallel_plans`). + base_model_sp_ep_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.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "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_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -76,6 +130,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 @@ -112,8 +172,89 @@ def __post_init__(self, **kwargs): self.head_dim = self.head_dim or self.hidden_size // self.num_attention_heads self.sliding_window = None + self.sliding_window = self.sliding_window if self.use_sliding_window else None self.mlp_only_layers = [] if self.mlp_only_layers is None else self.mlp_only_layers super().__post_init__(**kwargs) + self._update_parallel_plans() + + def _is_moe_layer(self, layer_idx: int) -> bool: + return (layer_idx not in self.mlp_only_layers) and ( + self.num_experts > 0 and (layer_idx + 1) % self.decoder_sparse_step == 0 + ) + + def _update_parallel_plans(self): + self._update_sp_plan() + self._update_sp_ep_plan() + self._update_tp_ep_plan() + + def _update_sp_plan(self): + self.base_model_sp_plan = self.base_model_sp_plan.copy() + + for i in range(self.num_hidden_layers): + if self._is_moe_layer(i): + self.base_model_sp_plan.update( + { + f"layers.{i}.mlp": "module_allgather_split", + f"layers.{i}.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + f"layers.{i}.mlp.experts.down_proj": "moe_tp_down_rowwise", + f"layers.{i}.mlp.experts": "moe_experts_allreduce", + } + ) + else: + self.base_model_sp_plan.update( + { + f"layers.{i}.mlp": "module_allgather", + f"layers.{i}.mlp.gate_proj": "colwise", + f"layers.{i}.mlp.up_proj": "colwise", + f"layers.{i}.mlp.down_proj": "rowwise_reduce_scatter", + } + ) + + def _update_sp_ep_plan(self): + self.base_model_sp_ep_plan = self.base_model_sp_ep_plan.copy() + + for i in range(self.num_hidden_layers): + if self._is_moe_layer(i): + self.base_model_sp_ep_plan.update( + { + f"layers.{i}.mlp": "module_allgather_split", + f"layers.{i}.mlp.gate": "ep_router", + f"layers.{i}.mlp.experts.gate_up_proj": "grouped_gemm", + f"layers.{i}.mlp.experts.down_proj": "grouped_gemm", + f"layers.{i}.mlp.experts": "moe_experts_allreduce", + } + ) + else: + self.base_model_sp_ep_plan.update( + { + f"layers.{i}.mlp": "module_allgather", + f"layers.{i}.mlp.gate_proj": "colwise", + f"layers.{i}.mlp.up_proj": "colwise", + f"layers.{i}.mlp.down_proj": "rowwise_reduce_scatter", + } + ) + + def _update_tp_ep_plan(self): + self.base_model_tp_ep_plan = self.base_model_tp_ep_plan.copy() + + for i in range(self.num_hidden_layers): + if self._is_moe_layer(i): + self.base_model_tp_ep_plan.update( + { + f"layers.{i}.mlp.gate": "ep_router", + f"layers.{i}.mlp.experts.gate_up_proj": "grouped_gemm", + f"layers.{i}.mlp.experts.down_proj": "grouped_gemm", + f"layers.{i}.mlp.experts": "moe_experts_allreduce", + } + ) + else: + self.base_model_tp_ep_plan.update( + { + f"layers.{i}.mlp.gate_proj": "colwise", + f"layers.{i}.mlp.up_proj": "colwise", + f"layers.{i}.mlp.down_proj": "rowwise_allreduce", + } + ) @auto_docstring(checkpoint="Qwen/Qwen3-VL-30B-A3B-Instruct") 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..91dbf4b3241d 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,52 @@ 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", + } + # Inference TP + EP (per-layer overrides applied in `_update_parallel_plans`). + base_model_tp_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_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + # Training SP + EP (per-layer overrides applied in `_update_parallel_plans`). + base_model_sp_ep_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.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "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_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -125,8 +161,8 @@ def __post_init__(self, **kwargs): self.num_key_value_heads = self.num_attention_heads self.head_dim = self.head_dim or self.hidden_size // self.num_attention_heads - super().__post_init__(**kwargs) self.sliding_window = None + super().__post_init__(**kwargs) @auto_docstring(checkpoint="Qwen/Qwen3-VL-30B-A3B-Instruct") 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..e2f667e03a36 100644 --- a/src/transformers/models/seed_oss/modeling_seed_oss.py +++ b/src/transformers/models/seed_oss/modeling_seed_oss.py @@ -429,8 +429,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..a8077ced3274 100644 --- a/src/transformers/models/smollm3/modeling_smollm3.py +++ b/src/transformers/models/smollm3/modeling_smollm3.py @@ -445,8 +445,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..aefa196808fe 100644 --- a/src/transformers/models/solar_open/configuration_solar_open.py +++ b/src/transformers/models/solar_open/configuration_solar_open.py @@ -41,16 +41,23 @@ 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.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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 +88,21 @@ 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.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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..3ba4782d3439 100644 --- a/src/transformers/models/solar_open/modeling_solar_open.py +++ b/src/transformers/models/solar_open/modeling_solar_open.py @@ -552,8 +552,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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..51cf4c0e0b39 100644 --- a/src/transformers/models/solar_open/modular_solar_open.py +++ b/src/transformers/models/solar_open/modular_solar_open.py @@ -47,10 +47,31 @@ 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.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "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..ab9ed58cfdfd 100644 --- a/src/transformers/models/starcoder2/modeling_starcoder2.py +++ b/src/transformers/models/starcoder2/modeling_starcoder2.py @@ -409,8 +409,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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 074877e17f21..c8aed75038a2 100644 --- a/src/transformers/models/t5gemma/modeling_t5gemma.py +++ b/src/transformers/models/t5gemma/modeling_t5gemma.py @@ -946,7 +946,8 @@ 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"} + _tp_ep_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 553cfdb41c2a..43fed4ea81a0 100644 --- a/src/transformers/models/t5gemma/modular_t5gemma.py +++ b/src/transformers/models/t5gemma/modular_t5gemma.py @@ -85,6 +85,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() @@ -785,7 +791,8 @@ 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"} + _tp_ep_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 2eba9554039d..587cd8768420 100644 --- a/src/transformers/models/t5gemma2/modeling_t5gemma2.py +++ b/src/transformers/models/t5gemma2/modeling_t5gemma2.py @@ -1183,7 +1183,8 @@ 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"} + _tp_ep_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 8b0ba064e16a..10f741252cd0 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,8 @@ 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"} + _tp_ep_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..72a1b88beaae 100644 --- a/src/transformers/models/vaultgemma/modeling_vaultgemma.py +++ b/src/transformers/models/vaultgemma/modeling_vaultgemma.py @@ -467,8 +467,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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/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 3660946a1a4d..d5c669a08495 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/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 f293235f5cbb..f2e2211337a8 100644 --- a/src/transformers/models/youtu/modeling_youtu.py +++ b/src/transformers/models/youtu/modeling_youtu.py @@ -534,8 +534,12 @@ 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"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_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 48b342e00559..b19c10307029 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. @@ -4216,6 +4252,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.""" @@ -4255,8 +4330,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 @@ -4266,7 +4342,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) @@ -4284,7 +4360,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() @@ -4295,26 +4371,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 672cb7151dfa..f19453b7bc96 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 @@ -2415,9 +2415,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/models/gpt_oss/test_modeling_gpt_oss.py b/tests/models/gpt_oss/test_modeling_gpt_oss.py index 5b3c13f8c25a..c7c0d1d93e93 100644 --- a/tests/models/gpt_oss/test_modeling_gpt_oss.py +++ b/tests/models/gpt_oss/test_modeling_gpt_oss.py @@ -69,6 +69,7 @@ class GptOssModelTest(CausalLMModelTest, unittest.TestCase): _is_stateful = True model_split_percents = [0.5, 0.6] model_tester_class = GptOssModelTester + tensor_parallel_atol = 7e-3 @require_kernels @require_torch_accelerator diff --git a/tests/models/hrm_text/test_modeling_hrm_text.py b/tests/models/hrm_text/test_modeling_hrm_text.py index 97644dede00d..441a632699b1 100644 --- a/tests/models/hrm_text/test_modeling_hrm_text.py +++ b/tests/models/hrm_text/test_modeling_hrm_text.py @@ -73,6 +73,14 @@ def test_tp_forward(self): def test_tp_backward(self): pass + @unittest.skip(reason="Higher tols (likely due to different recursion and grad patterns). FIXME later") + def test_sp_forward(self): + pass + + @unittest.skip(reason="Higher tols (likely due to different recursion and grad patterns). FIXME later") + def test_sp_backward(self): + pass + @unittest.skip(reason="Higher tols (likely due to different recursion and grad patterns). FIXME later") def test_tp_generation(self): pass diff --git a/tests/models/qwen3_moe/test_modeling_qwen3_moe.py b/tests/models/qwen3_moe/test_modeling_qwen3_moe.py index 13fafa6969ae..78cbe6b08ef7 100644 --- a/tests/models/qwen3_moe/test_modeling_qwen3_moe.py +++ b/tests/models/qwen3_moe/test_modeling_qwen3_moe.py @@ -45,6 +45,15 @@ class Qwen3MoeModelTester(CausalLMModelTester): if is_torch_available(): base_model_class = Qwen3MoeModel + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # Make the tiny config build alternating dense/sparse layers so the SP plan is exercised + # against both a dense Qwen3MoeMLP and a sparse Qwen3MoeSparseMoeBlock. decoder_sparse_step=2 + # makes even layers dense and odd layers MoE; pin >=2 layers so a sparse layer actually + # exists (don't rely on the shared tester num_hidden_layers default staying >= 2). + self.num_hidden_layers = max(self.num_hidden_layers, 2) + self.decoder_sparse_step = 2 + @require_torch class Qwen3MoeModelTest(CausalLMModelTest, unittest.TestCase): 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..d6b2b0ab2bd5 --- /dev/null +++ b/tests/test_distributed_config.py @@ -0,0 +1,97 @@ +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_cpu_offload is False + assert dc.fsdp_mixed_precision is False + + 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_cpu_offload is False + assert dc.fsdp_mixed_precision is False + 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_cpu_offload is False + assert dc.fsdp_mixed_precision is False + + 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_cpu_offload": False, + "fsdp_mixed_precision": False, + } + + 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..280ea724f63a --- /dev/null +++ b/tests/test_fsdp_mixin.py @@ -0,0 +1,710 @@ +# 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 _resolve_tied_embed_lm_head_plan, expand_fsdp_plan + 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 _fsdp_target_module_names(model) -> set[str]: + """Return module names that should receive a ``fully_shard`` call (excluding root).""" + tie_word_embeddings = getattr(model.config, "tie_word_embeddings", False) + adapted_fsdp_plan = _resolve_tied_embed_lm_head_plan(model._fsdp_plan, tie_word_embeddings=tie_word_embeddings) + return {name for name, _, _ in expand_fsdp_plan(model, adapted_fsdp_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).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_cpu_offload=False, + fsdp_mixed_precision=False, +): + # -- 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_cpu_offload=fsdp_cpu_offload, + fsdp_mixed_precision=fsdp_mixed_precision, + ) + pre_ckpt_model = AutoModelForCausalLM.from_pretrained( + init_model_dir, torch_dtype=dtype, distributed_config=distributed_config + ) + 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 + ) + 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) + 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) + 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 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 + + init_tmpdir, init_tmpdir_obj = _save_init_pretrained(rank, config, torch.float32) + try: + _set_determinism(SEED) + model = AutoModelForCausalLM.from_pretrained( + init_tmpdir, distributed_config=DistributedConfig(fsdp_size=dist.get_world_size()) + ) + dist.barrier() + finally: + if rank == 0 and init_tmpdir_obj is not None: + init_tmpdir_obj.cleanup() + + # 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 = {""} | _fsdp_target_module_names(model) + + 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, policy_options=None, dtype=None +): + """Validate DDP-vs-FSDP2 trace matching using the model's declared FSDP plan.""" + 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 + + fsdp_cpu_offload = "cpu_offload" in policy_options + fsdp_mixed_precision = "mixed_precision" in policy_options + test_label = f"FSDP2{'+policies' if policy_options else ''}" + + 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_cpu_offload=fsdp_cpu_offload, + fsdp_mixed_precision=fsdp_mixed_precision, + ) + 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_plan_vs_ddp_untied(self): + self._run_fsdp2_distributed_test("test_fsdp2_plan_vs_ddp_untied", _test_fsdp2_plan_vs_ddp_impl, False) + + @is_fsdp_test + @require_fsdp + def test_fsdp2_plan_vs_ddp_tied(self): + self._run_fsdp2_distributed_test("test_fsdp2_plan_vs_ddp_tied", _test_fsdp2_plan_vs_ddp_impl, True) diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py index 660f5fe0e6d1..3735ad06211a 100644 --- a/tests/test_modeling_common.py +++ b/tests/test_modeling_common.py @@ -133,8 +133,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..f35fa305ef44 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: @@ -44,31 +62,6 @@ def _find_free_port(): return s.getsockname()[1] -def get_packed_grad_shard(grad, world_size, rank, dim): - """Get the correct shard of a packed gradient (matching get_packed_weights interleaved logic). - - Packed weights like gate_up_proj are sharded with interleaving: - Original: [G0 G1 G2 G3 | U0 U1 U2 U3] (gate | up) - Rank 0: [G0 G1 | U0 U1] - Rank 1: [G2 G3 | U2 U3] - """ - total_size = grad.shape[dim] - # Packed weights have 2 blocks (gate and up) - block_size = total_size // 2 - shard_block_size = block_size // world_size - - # Build interleaved indices - indices = [] - for block_idx in range(2): # gate block, then up block - block_offset = block_idx * block_size - start = block_offset + rank * shard_block_size - stop = block_offset + (rank + 1) * shard_block_size - indices.extend(range(start, stop)) - - # Select along the sharded dimension - return grad.index_select(dim, torch.tensor(indices, device=grad.device)) - - def _global_wrapper(rank, func, tp, port, backend, func_args, func_kwargs): """Wrapper to set up distributed environment and run the test function.""" @@ -112,152 +105,164 @@ def wrapper(*args, **kwargs): return _init_distributed_inner -def _load_tp_and_reference_models(model_path, model_class): - """Load TP model and non-TP reference model for comparison. +def _load_distributed_and_reference_models(model_path, model_class, mode: str): + """Load a distributed model and an unsharded reference model for comparison.""" + tp_size = dist.get_world_size() + if mode in ("ep", "tp_ep"): + distributed_config = DistributedConfig(tp_size=tp_size, enable_expert_parallel=True) + elif mode in ("sp", "sp_ep"): + distributed_config = DistributedConfig( + tp_size=tp_size, enable_sequence_parallel=True, enable_expert_parallel=(mode == "sp_ep") + ) + else: + distributed_config = DistributedConfig(tp_size=tp_size, enable_sequence_parallel=False) - Returns: - tuple: (model_tp, model_ref, device) - """ - model_tp = model_class.from_pretrained(model_path, tp_plan="auto") + from_pretrained_kwargs = {"distributed_config": distributed_config} + if mode not in ("ep", "tp_ep"): + from_pretrained_kwargs["attn_implementation"] = "sdpa" + + model_dist = model_class.from_pretrained(model_path, **from_pretrained_kwargs) dist.barrier() - device = model_tp.device - model_ref = model_class.from_pretrained(model_path) + device = model_dist.device + if mode in ("ep", "tp_ep"): + model_ref = model_class.from_pretrained(model_path) + else: + model_ref = model_class.from_pretrained(model_path, attn_implementation="sdpa") model_ref = model_ref.to(device) - return model_tp, model_ref, device + return model_dist, model_ref, device -def _verify_tp_sharding(rank, model_tp, model_ref): - """Verify TP sharding by comparing parameter shapes between TP and reference models. - - Returns: - list: Names of sharded parameters - """ - world_size = dist.get_world_size() - sharded_params = [] - - for (name, param), (_, param_full) in zip(model_tp.named_parameters(), model_ref.named_parameters()): - if param.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}") - - # 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",): - 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)}" - ) - 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)}" - ) - break - - return sharded_params - - -def _test_tp_forward_impl(_rank, model_path, model_class, atol, rtol): - """Implementation for comparing TP and non-TP model outputs.""" +def _test_forward_impl(_rank, model_path, model_class, atol, rtol, mode: str): + """Compare distributed and reference model forward outputs for the given mode.""" + assert mode in ("tp", "sp", "ep", "tp_ep", "sp_ep") set_seed(0) - model_tp, model, device = _load_tp_and_reference_models(model_path, model_class) + model_dist, model_ref, device = _load_distributed_and_reference_models(model_path, model_class, mode) - _verify_tp_sharding(_rank, model_tp, model) - - model_tp.eval() - model.eval() + model_dist.eval() + model_ref.eval() - vocab_size = model.config.vocab_size + vocab_size = model_ref.config.vocab_size set_seed(0) input_ids = torch.randint(0, vocab_size, (2, 64)).to(device) with torch.no_grad(): - logits = model(input_ids).logits - logits_tp = model_tp(input_ids).logits + logits_ref = model_ref(input_ids).logits + logits_dist = model_dist(input_ids).logits + if mode in ("tp", "sp", "sp_ep"): + logits_dist = _to_local(logits_dist) - diff = (logits - logits_tp).abs() - assert torch.allclose(logits, logits_tp, atol=atol, rtol=rtol), ( - f"TP and non-TP model outputs differ. Max diff: {diff.max().item()} | Min diff: {diff.min().item()}" + diff = (logits_ref - logits_dist).abs() + assert torch.allclose(logits_ref, logits_dist, atol=atol, rtol=rtol), ( + f"{mode.upper()} and reference model outputs differ. Max diff: {diff.max().item()} | Min diff: {diff.min().item()}" ) dist.barrier() -def _test_tp_backward_impl(rank, model_path, model_class, atol, rtol): - """Implementation for comparing TP and non-TP model backward passes.""" +def _get_packed_grad_shard(grad, world_size, rank, dim): + """Get the correct shard of a packed gradient (matching get_packed_weights interleaved logic). + + Packed weights (packed_colwise, moe_tp_gate_up_colwise) are sharded with interleaving: + Original: [G0 G1 G2 G3 | U0 U1 U2 U3] (gate | up) + Rank 0: [G0 G1 | U0 U1] + Rank 1: [G2 G3 | U2 U3] + """ + total_size = grad.shape[dim] + # Packed weights have 2 blocks (gate and up) + block_size = total_size // 2 + shard_block_size = block_size // world_size + + # Build interleaved indices + indices = [] + for block_idx in range(2): # gate block, then up block + block_offset = block_idx * block_size + start = block_offset + rank * shard_block_size + stop = block_offset + (rank + 1) * shard_block_size + indices.extend(range(start, stop)) + + # Select along the sharded dimension + return grad.index_select(dim, torch.tensor(indices, device=grad.device)) + + +def _test_backward_impl(rank, model_path, model_class, atol, rtol, mode: str): + """Compare distributed and reference model backward passes for the given mode.""" + assert mode in ("tp", "sp", "ep", "tp_ep", "sp_ep") set_seed(0) - model_tp, model, device = _load_tp_and_reference_models(model_path, model_class) - model_tp.train() - model.train() + model_dist, model_ref, device = _load_distributed_and_reference_models(model_path, model_class, mode) + applied_plan = getattr(model_dist, "tp_plan", None) or getattr(model_dist, f"_{mode}_plan", None) or {} + model_dist.train() + model_ref.train() - vocab_size = model.config.vocab_size + vocab_size = model_ref.config.vocab_size set_seed(0) input_ids = torch.randint(0, vocab_size, (2, 64)).to(device) set_seed(0) labels = torch.randint(0, vocab_size, (2, 64)).to(device) - loss = model(input_ids, labels=labels).loss - loss.backward() + loss_ref = model_ref(input_ids, labels=labels, use_cache=False).loss + loss_ref.backward() + + loss_dist = model_dist(input_ids, labels=labels, use_cache=False).loss + loss_dist.backward() - loss_tp = model_tp(input_ids, labels=labels).loss - loss_tp.backward() + if mode in ("tp", "sp", "sp_ep"): + loss_dist_value = _to_local(loss_dist) + else: + loss_dist_value = loss_dist - assert torch.allclose(loss, loss_tp, 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()}" + assert torch.allclose(loss_ref, loss_dist_value, atol=atol, rtol=rtol), ( + f"{mode.upper()} and reference model losses differ. " + f"Reference loss: {loss_ref.item()}, {mode.upper()} loss: {loss_dist_value.item()}, " + f"Diff: {(loss_ref - loss_dist_value).abs().item()}" ) - # Compare gradients for matching parameters + # Compare parameter gradients world_size = dist.get_world_size() - failed_grads = {} - for (name, param), (_, 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 - - # 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",): - # interleaved slicing - grad = get_packed_grad_shard(grad, world_size, rank, dim) - else: - # regular slicing - shard_size = grad_tp.size(dim) - start = rank * shard_size - grad = grad.narrow(dim, start, shard_size) - 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() - - 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() - ) + failed_grads = [] + for (name, param), (_, param_dist) in zip(model_ref.named_parameters(), model_dist.named_parameters()): + if param.grad is None or param_dist.grad is None: + continue + + grad = param.grad + grad_dist = _to_local(param_dist.grad) + + if grad.shape != grad_dist.shape: + for dim in range(grad.ndim): + if grad.size(dim) != grad_dist.size(dim): + param_plan = _get_parameter_tp_plan(name, applied_plan, is_weight=True) + if param_plan in ("packed_colwise", "moe_tp_gate_up_colwise"): + grad = _get_packed_grad_shard(grad, world_size, rank, dim) + else: + shard_size = grad_dist.size(dim) + start = rank * shard_size + grad = grad.narrow(dim, start, shard_size) + break + + try: + torch.testing.assert_close(grad, grad_dist, atol=atol, rtol=rtol) + except AssertionError as e: + failed_grads.append(f"{name}: {e}") + + assert not failed_grads, "Gradients differ:\n" + "\n".join(failed_grads) dist.barrier() -def _test_tp_generation_impl(_rank, model_path, model_class, atol, rtol, max_new_tokens): - """Implementation for comparing TP and non-TP model generation outputs (direct load path).""" +def _test_generation_impl(_rank, model_path, model_class, atol, rtol, mode: str, max_new_tokens): + """Compare distributed and reference model generation outputs for the given mode.""" + assert mode in ("tp", "ep", "tp_ep") set_seed(0) - model_tp, model, device = _load_tp_and_reference_models(model_path, model_class) - model_tp.eval() - model.eval() + model_dist, model_ref, device = _load_distributed_and_reference_models(model_path, model_class, mode) + model_dist.eval() + model_ref.eval() set_seed(0) - vocab_size = model.config.vocab_size + vocab_size = model_ref.config.vocab_size input_ids = torch.randint(0, vocab_size, (1, 10)).to(device) generation_kwargs = { "max_new_tokens": max_new_tokens, @@ -269,35 +274,48 @@ def _test_tp_generation_impl(_rank, model_path, model_class, atol, rtol, max_new } with torch.no_grad(): - output = model.generate(input_ids, **generation_kwargs) - output_tp = model_tp.generate(input_ids, **generation_kwargs) + output = model_ref.generate(input_ids, **generation_kwargs) + output_dist = model_dist.generate(input_ids, **generation_kwargs) # Compare logits/scores at each generation step scores = torch.stack(output.scores) - scores_tp = torch.stack(output_tp.scores) - - diff = (scores - scores_tp).abs() - assert torch.allclose(scores, scores_tp, atol=atol, rtol=rtol), ( - f"TP and non-TP model generation logits differ (direct load path). " + if mode in ("tp", "sp"): + scores_dist = torch.stack([_to_local(s) for s in output_dist.scores]) + else: + scores_dist = torch.stack(output_dist.scores) + + diff = (scores - scores_dist).abs() + assert torch.allclose(scores, scores_dist, atol=atol, rtol=rtol), ( + f"{mode.upper()} and reference model generation logits differ. " f"Max diff: {diff.max().item()} | Mean diff: {diff.mean().item()}" ) # Compare generated token sequences - assert torch.equal(output.sequences, output_tp.sequences), ( - f"TP and non-TP model generated different token sequences (direct load path). " - f"Non-TP: {output.sequences.tolist()} | TP: {output_tp.sequences.tolist()}" + sequences_dist = _to_local(output_dist.sequences) if mode in ("tp", "sp") else output_dist.sequences + assert torch.equal(output.sequences, sequences_dist), ( + f"{mode.upper()} and reference model generated different token sequences. " + f"Reference: {output.sequences.tolist()} | {mode.upper()}: {sequences_dist.tolist()}" ) dist.barrier() -def _test_tp_generation_quantized_impl(_rank, model_path, model_class, max_new_tokens): +def _test_generation_quantized_impl(_rank, model_path, model_class, max_new_tokens, mode: str = "tp"): """Implementation for comparing TP+quantized and non-TP quantized generation (sequence equality).""" + assert mode in ("tp", "tp_ep") set_seed(0) quantization_config = TorchAoConfig(Float8WeightOnlyConfig()) - model_tp = model_class.from_pretrained(model_path, tp_plan="auto", quantization_config=quantization_config) + distributed_config = DistributedConfig( + tp_size=dist.get_world_size(), + enable_expert_parallel=(mode == "tp_ep"), + ) + model_tp = model_class.from_pretrained( + model_path, + distributed_config=distributed_config, + quantization_config=quantization_config, + ) dist.barrier() device = model_tp.device @@ -345,72 +363,6 @@ def _test_tp_generation_quantized_impl(_rank, model_path, model_class, max_new_t dist.barrier() -def _load_ep_and_reference_models(model_path, model_class): - """Load EP model and non-EP reference model for comparison.""" - model_ep = model_class.from_pretrained( - model_path, - distributed_config=DistributedConfig(enable_expert_parallel=True), - ) - dist.barrier() - - device = model_ep.device - model_ref = model_class.from_pretrained(model_path) - model_ref = model_ref.to(device) - - return model_ep, model_ref, device - - -def _test_ep_forward_impl(_rank, model_path, model_class, atol, rtol): - """Implementation for comparing EP and non-EP model outputs.""" - set_seed(0) - - model_ep, model_ref, device = _load_ep_and_reference_models(model_path, model_class) - - model_ep.eval() - model_ref.eval() - - vocab_size = model_ref.config.vocab_size - input_ids = torch.randint(0, vocab_size, (2, 64)).to(device) - - with torch.no_grad(): - logits_ref = model_ref(input_ids).logits - logits_ep = model_ep(input_ids).logits - - diff = (logits_ref - logits_ep).abs() - assert torch.allclose(logits_ref, logits_ep, atol=atol, rtol=rtol), ( - f"EP and non-EP model outputs differ. Max diff: {diff.max().item()} | Min diff: {diff.min().item()}" - ) - - dist.barrier() - - -def _test_ep_backward_impl(_rank, model_path, model_class, atol, rtol): - """Implementation for comparing EP and non-EP model backward passes.""" - set_seed(0) - - model_ep, model_ref, device = _load_ep_and_reference_models(model_path, model_class) - model_ep.train() - model_ref.train() - - vocab_size = model_ref.config.vocab_size - input_ids = torch.randint(0, vocab_size, (2, 64)).to(device) - labels = torch.randint(0, vocab_size, (2, 64)).to(device) - - loss_ref = model_ref(input_ids, labels=labels).loss - loss_ref.backward() - - loss_ep = model_ep(input_ids, labels=labels).loss - loss_ep.backward() - - assert torch.allclose(loss_ref, loss_ep, atol=atol, rtol=rtol), ( - f"EP and non-EP model losses differ. " - f"Non-EP loss: {loss_ref.item()}, EP loss: {loss_ep.item()}, " - f"Diff: {(loss_ref - loss_ep).abs().item()}" - ) - - dist.barrier() - - class TensorParallelTesterMixin(ABC): """ Mixin for tensor parallel tests. Add to model test classes alongside ModelTesterMixin. @@ -420,14 +372,21 @@ class TensorParallelTesterMixin(ABC): - causal_lm_class, base_model_class, etc. This mixin adds tensor parallel-specific tests using that infrastructure. + + Distributed test modes: + - "tp": _tp_plan, enable_sequence_parallel=False — inference-style TP. + - "sp": _sp_plan with per-layer MLP entries — training-style sequence parallel. + - "ep": _ep_plan, enable_expert_parallel=True — expert parallel on MoE weights (legacy EP-only path). + - "tp_ep": _tp_plan ∪ _ep_plan — inference TP dense + EP MoE. + - "sp_ep": _sp_plan ∪ _ep_plan — training SP dense + EP MoE. """ # ============================================================ # 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 @@ -448,14 +407,49 @@ def _has_tp_plan(self) -> bool: config = self.model_tester.get_config() return hasattr(config, "base_model_tp_plan") and config.base_model_tp_plan is not None - def _get_tp_model_class(self): + def _has_sp_plan(self) -> bool: + config = self.model_tester.get_config() + return getattr(config, "base_model_sp_plan", None) and config.base_model_sp_plan is not None + + def _get_model_class(self): """Get the model class to use for TP tests (prefers *ForCausalLM).""" if hasattr(self.model_tester, "causal_lm_class") and self.model_tester.causal_lm_class is not None: return self.model_tester.causal_lm_class return self.all_model_classes[0] - def _skip_if_not_supported(self): - """Check and skip test if TP is not supported for this model/environment.""" + def _assert_mixed_mlp_layers(self, model, config): + """If the config interleaves dense and sparse/MoE layers, fail fast unless the tiny test + model actually built both kinds. + """ + layer_types = getattr(config, "mlp_layer_types", None) + if layer_types is not None: + # MoE marker is 'sparse' or 'moe' depending on the model; dense is always 'dense'. + intends_mixed = "dense" in layer_types and any(t != "dense" for t in layer_types) + else: + sparse_step = getattr(config, "decoder_sparse_step", 1) or 1 + intends_mixed = sparse_step > 1 or bool(getattr(config, "mlp_only_layers", None)) + if not intends_mixed: + return + + kinds = { + "moe" if hasattr(module, "experts") else "dense" + for name, module in model.named_modules() + if name.endswith(".mlp") + } + # Empty means this model's feed-forward isn't named `.mlp` (e.g. `block_sparse_moe`); the + # structural check doesn't apply, so skip rather than spuriously fail its TP test. + if not kinds: + return + self.assertEqual( + kinds, + {"dense", "moe"}, + f"{type(model).__name__}: config interleaves dense/MoE mlp layers but the tiny test config " + f"built only {kinds}. Adjust the model tester (num_hidden_layers / decoder_sparse_step / " + f"mlp_layer_types) so both kinds exist.", + ) + + def _skip_if_mode_not_supported(self, mode: str) -> None: + """Check and skip test if the mode is not supported for this model/environment.""" if not is_torch_greater_or_equal("2.9"): self.skipTest("Tensor parallel tests require torch >= 2.9") @@ -471,60 +465,101 @@ def _skip_if_not_supported(self): if not hasattr(self.model_tester, "causal_lm_class") or self.model_tester.causal_lm_class is None: self.skipTest("Model tester does not have causal_lm_class (not using CausalLMModelTester)") - if not self._has_tp_plan(): + if mode == "tp" and not self._has_tp_plan(): self.skipTest("Model does not have a tensor parallel plan (base_model_tp_plan)") + elif mode == "sp" and not self._has_sp_plan(): + self.skipTest("Model does not have a sequence parallel plan (base_model_sp_plan)") + elif mode == "ep" and not self._has_ep_plan(): + self.skipTest("Model does not have an expert parallel plan (base_model_ep_plan)") + elif mode == "tp_ep": + if not self._has_tp_plan(): + self.skipTest("Model does not have a tensor parallel plan (base_model_tp_plan)") + if not self._has_ep_plan(): + self.skipTest("Model does not have an expert parallel plan (base_model_ep_plan)") + elif mode == "sp_ep": + if not self._has_sp_plan(): + self.skipTest("Model does not have a sequence parallel plan (base_model_sp_plan)") + if not self._has_ep_plan(): + self.skipTest("Model does not have an expert parallel plan (base_model_ep_plan)") # # Skip encoder-decoder models (TP not supported) # if getattr(self, "is_encoder_decoder", False): # self.skipTest("TP tests not supported for encoder-decoder models") - # # Skip VLM models for now - # config = self.model_tester.get_config() - # if hasattr(config, "vision_config") and config.vision_config is not None: - # self.skipTest("VLM models are not yet supported in TP tests") + @is_tensor_parallel_test + def test_tp_forward(self): + """Tensor parallel forward (_tp_plan, no sequence parallel).""" + self._skip_if_mode_not_supported("tp") + config = self.model_tester.get_config() + model_class = self._get_model_class() + atol = self.tensor_parallel_atol + rtol = self.tensor_parallel_rtol - def _skip_if_ep_not_supported(self): - """Check and skip test if EP is not supported for this model/environment.""" - if not is_torch_greater_or_equal("2.9"): - self.skipTest("Expert parallel tests require torch >= 2.9") + with tempfile.TemporaryDirectory() as tmp_dir: + set_seed(42) + model = model_class(config) + self._assert_mixed_mlp_layers(model, config) + model.save_pretrained(tmp_dir, save_original_format=True) + _init_distributed(tp=self.tensor_parallel_size)(_test_forward_impl)(tmp_dir, model_class, atol, rtol, "tp") - if torch.cuda.is_available() or torch.xpu.is_available(): - self.skipTest("Expert parallel mixin tests are CPU-only and should not run on GPU or XPU machines") + @is_tensor_parallel_test + def test_tp_backward(self): + """Tensor parallel backward (_tp_plan, no sequence parallel).""" + self._skip_if_mode_not_supported("tp") + config = self.model_tester.get_config() + model_class = self._get_model_class() + atol = self.tensor_parallel_atol + rtol = self.tensor_parallel_rtol - if os.cpu_count() < self.tensor_parallel_size: - self.skipTest( - f"Expert parallel tests require at least {self.tensor_parallel_size} CPUs, " - f"but only {os.cpu_count()} available" + with tempfile.TemporaryDirectory() as tmp_dir: + set_seed(42) + model = model_class(config) + self._assert_mixed_mlp_layers(model, config) + model.save_pretrained(tmp_dir, save_original_format=True) + _init_distributed(tp=self.tensor_parallel_size)(_test_backward_impl)( + tmp_dir, model_class, atol, rtol, "tp" ) - if not hasattr(self.model_tester, "causal_lm_class") or self.model_tester.causal_lm_class is None: - self.skipTest("Model tester does not have causal_lm_class (not using CausalLMModelTester)") + @is_tensor_parallel_test + def test_sp_forward(self): + """Sequence-parallel forward (_sp_plan including per-layer MLP).""" + self._skip_if_mode_not_supported("sp") + config = self.model_tester.get_config() + model_class = self._get_model_class() + atol = self.tensor_parallel_atol + rtol = self.tensor_parallel_rtol - if not self._has_ep_plan(): - self.skipTest("Model does not have an expert parallel plan (base_model_ep_plan)") + with tempfile.TemporaryDirectory() as tmp_dir: + set_seed(42) + model = model_class(config) + self._assert_mixed_mlp_layers(model, config) + model.save_pretrained(tmp_dir, save_original_format=True) + _init_distributed(tp=self.tensor_parallel_size)(_test_forward_impl)(tmp_dir, model_class, atol, rtol, "sp") @is_tensor_parallel_test - def test_tp_forward(self): - self._skip_if_not_supported() - + def test_sp_backward(self): + """Sequence-parallel backward (_sp_plan including per-layer MLP).""" + self._skip_if_mode_not_supported("sp") config = self.model_tester.get_config() - model_class = self._get_tp_model_class() + model_class = self._get_model_class() atol = self.tensor_parallel_atol rtol = self.tensor_parallel_rtol with tempfile.TemporaryDirectory() as tmp_dir: set_seed(42) model = model_class(config) + self._assert_mixed_mlp_layers(model, config) model.save_pretrained(tmp_dir, save_original_format=True) - - _init_distributed(tp=self.tensor_parallel_size)(_test_tp_forward_impl)(tmp_dir, model_class, atol, rtol) + _init_distributed(tp=self.tensor_parallel_size)(_test_backward_impl)( + tmp_dir, model_class, atol, rtol, "sp" + ) @is_tensor_parallel_test - def test_tp_backward(self): - self._skip_if_not_supported() - + def test_ep_forward(self): + """Expert-parallel forward (_ep_plan).""" + self._skip_if_mode_not_supported("ep") config = self.model_tester.get_config() - model_class = self._get_tp_model_class() + model_class = self._get_model_class() atol = self.tensor_parallel_atol rtol = self.tensor_parallel_rtol @@ -532,71 +567,146 @@ def test_tp_backward(self): set_seed(42) model = model_class(config) model.save_pretrained(tmp_dir, save_original_format=True) + _init_distributed(tp=self.tensor_parallel_size)(_test_forward_impl)(tmp_dir, model_class, atol, rtol, "ep") - _init_distributed(tp=self.tensor_parallel_size)(_test_tp_backward_impl)(tmp_dir, model_class, atol, rtol) + @is_tensor_parallel_test + def test_ep_backward(self): + """Expert-parallel backward (_ep_plan).""" + self._skip_if_mode_not_supported("ep") + config = self.model_tester.get_config() + model_class = self._get_model_class() + atol = self.tensor_parallel_atol + rtol = self.tensor_parallel_rtol + + 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_backward_impl)( + tmp_dir, model_class, atol, rtol, "ep" + ) @is_tensor_parallel_test def test_tp_generation(self): - # Test TP generation: unfused checkpoint → conversion mapping (if needed) → TP sharding → model → generate - self._skip_if_not_supported() - + """Tensor parallel generation (_tp_plan).""" + self._skip_if_mode_not_supported("tp") config = self.model_tester.get_config() + model_class = self._get_model_class() + atol = self.tensor_parallel_atol + rtol = self.tensor_parallel_rtol - model_class = self._get_tp_model_class() + 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_generation_impl)( + tmp_dir, model_class, atol, rtol, "tp", 25 + ) + + @is_tensor_parallel_test + def test_ep_generation(self): + """Expert-parallel generation (_ep_plan).""" + self._skip_if_mode_not_supported("ep") + config = self.model_tester.get_config() + model_class = self._get_model_class() atol = self.tensor_parallel_atol rtol = self.tensor_parallel_rtol - 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_impl)( - tmp_dir, model_class, atol, rtol, max_new_tokens + _init_distributed(tp=self.tensor_parallel_size)(_test_generation_impl)( + tmp_dir, model_class, atol, rtol, "ep", 25 ) @is_tensor_parallel_test def test_tp_generation_quantized(self): - self._skip_if_not_supported() + self._skip_if_mode_not_supported("tp") if not is_torchao_available(): self.skipTest("Test requires torchao") + self.skipTest("Quantization is not currently supported with distributed training (dtensor)") + + @is_tensor_parallel_test + def test_tp_ep_forward(self): + """Inference TP + EP forward (_tp_plan ∪ _ep_plan).""" + self._skip_if_mode_not_supported("tp_ep") config = self.model_tester.get_config() - model_class = self._get_tp_model_class() - max_new_tokens = 25 + model_class = self._get_model_class() + atol = self.tensor_parallel_atol + rtol = self.tensor_parallel_rtol with tempfile.TemporaryDirectory() as tmp_dir: set_seed(42) model = model_class(config) + self._assert_mixed_mlp_layers(model, 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 + _init_distributed(tp=self.tensor_parallel_size)(_test_forward_impl)( + tmp_dir, model_class, atol, rtol, "tp_ep" ) @is_tensor_parallel_test - def test_ep_forward(self): - self._skip_if_ep_not_supported() - + def test_tp_ep_backward(self): + """Inference TP + EP backward (_tp_plan ∪ _ep_plan).""" + self._skip_if_mode_not_supported("tp_ep") config = self.model_tester.get_config() - model_class = self._get_tp_model_class() + model_class = self._get_model_class() atol = self.tensor_parallel_atol rtol = self.tensor_parallel_rtol with tempfile.TemporaryDirectory() as tmp_dir: set_seed(42) model = model_class(config) + self._assert_mixed_mlp_layers(model, config) model.save_pretrained(tmp_dir, save_original_format=True) + _init_distributed(tp=self.tensor_parallel_size)(_test_backward_impl)( + tmp_dir, model_class, atol, rtol, "tp_ep" + ) - _init_distributed(tp=self.tensor_parallel_size)(_test_ep_forward_impl)(tmp_dir, model_class, atol, rtol) + @is_tensor_parallel_test + def test_sp_ep_forward(self): + """Training SP + EP forward (_sp_plan ∪ _ep_plan).""" + self._skip_if_mode_not_supported("sp_ep") + config = self.model_tester.get_config() + model_class = self._get_model_class() + atol = self.tensor_parallel_atol + rtol = self.tensor_parallel_rtol + + with tempfile.TemporaryDirectory() as tmp_dir: + set_seed(42) + model = model_class(config) + self._assert_mixed_mlp_layers(model, config) + model.save_pretrained(tmp_dir, save_original_format=True) + _init_distributed(tp=self.tensor_parallel_size)(_test_forward_impl)( + tmp_dir, model_class, atol, rtol, "sp_ep" + ) @is_tensor_parallel_test - def test_ep_backward(self): - self._skip_if_ep_not_supported() + def test_sp_ep_backward(self): + """Training SP + EP backward (_sp_plan ∪ _ep_plan).""" + self._skip_if_mode_not_supported("sp_ep") + config = self.model_tester.get_config() + model_class = self._get_model_class() + atol = self.tensor_parallel_atol + rtol = self.tensor_parallel_rtol + + with tempfile.TemporaryDirectory() as tmp_dir: + set_seed(42) + model = model_class(config) + self._assert_mixed_mlp_layers(model, config) + model.save_pretrained(tmp_dir, save_original_format=True) + _init_distributed(tp=self.tensor_parallel_size)(_test_backward_impl)( + tmp_dir, model_class, atol, rtol, "sp_ep" + ) + @is_tensor_parallel_test + def test_tp_ep_generation(self): + """Inference TP + EP generation (_tp_plan ∪ _ep_plan).""" + self._skip_if_mode_not_supported("tp_ep") config = self.model_tester.get_config() - model_class = self._get_tp_model_class() + model_class = self._get_model_class() atol = self.tensor_parallel_atol rtol = self.tensor_parallel_rtol @@ -604,5 +714,15 @@ def test_ep_backward(self): set_seed(42) model = model_class(config) model.save_pretrained(tmp_dir, save_original_format=True) + _init_distributed(tp=self.tensor_parallel_size)(_test_generation_impl)( + tmp_dir, model_class, atol, rtol, "tp_ep", 25 + ) + + @is_tensor_parallel_test + def test_tp_ep_generation_quantized(self): + self._skip_if_mode_not_supported("tp_ep") + + if not is_torchao_available(): + self.skipTest("Test requires torchao") - _init_distributed(tp=self.tensor_parallel_size)(_test_ep_backward_impl)(tmp_dir, model_class, atol, rtol) + self.skipTest("Quantization is not currently supported with distributed training (dtensor)") diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py index 9ac9810c9808..91e24988b81b 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 4331d0987bc7..a83bfdf20a16 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 ( @@ -436,6 +437,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 + fake_mesh.device_type = "cpu" + + def fake_apply_fsdp(model, fsdp_mesh): + call_order.append("distribute") + self.assertIs(fsdp_mesh, fake_mesh) + 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/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)", }