diff --git a/.circleci/create_circleci_config.py b/.circleci/create_circleci_config.py index de0646af2dbf..d004a4c37cf0 100644 --- a/.circleci/create_circleci_config.py +++ b/.circleci/create_circleci_config.py @@ -313,7 +313,7 @@ def job_name(self): torch_job = CircleCIJob( "torch", docker_image=[{"image": "huggingface/transformers-torch-light"}], - marker="not generate", + marker="not (generate or is_training_test or is_tensor_parallel_test or is_fsdp_test)", parallelism=6, ) @@ -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 ."], + 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/conftest.py b/conftest.py index 3970d3ffee9b..899c38acebd3 100644 --- a/conftest.py +++ b/conftest.py @@ -97,6 +97,7 @@ def pytest_configure(config): ) config.addinivalue_line("markers", "training_ci: mark test for training CI validation") config.addinivalue_line("markers", "tensor_parallel_ci: mark test for tensor parallel CI validation") + config.addinivalue_line("markers", "fsdp_ci: mark test for FSDP CI validation") os.environ["DISABLE_SAFETENSORS_CONVERSION"] = "true" register_network_debug_plugin(config) diff --git a/docs/source/en/expert_parallelism.md b/docs/source/en/expert_parallelism.md index 61d298bb5cef..63256d518309 100644 --- a/docs/source/en/expert_parallelism.md +++ b/docs/source/en/expert_parallelism.md @@ -22,11 +22,16 @@ rendered properly in your Markdown viewer. Enable expert parallelism with the [`DistributedConfig`] class and the `enable_expert_parallel` argument. ```py +import os + import torch from transformers import AutoModelForCausalLM, AutoTokenizer from transformers.distributed.configuration_utils import DistributedConfig -distributed_config = DistributedConfig(enable_expert_parallel=True) +distributed_config = DistributedConfig( + tp_size=int(os.environ["WORLD_SIZE"]), + enable_expert_parallel=True, +) model = AutoModelForCausalLM.from_pretrained( "openai/gpt-oss-120b", diff --git a/docs/source/en/model_doc/minimax_m3_vl.md b/docs/source/en/model_doc/minimax_m3_vl.md index 5a73737b7e67..56a07cfab807 100644 --- a/docs/source/en/model_doc/minimax_m3_vl.md +++ b/docs/source/en/model_doc/minimax_m3_vl.md @@ -165,8 +165,10 @@ model = AutoModelForImageTextToText.from_pretrained( dtype=torch.bfloat16, # Dequantize the native MXFP8 weights to bf16 at load (the speed win); needs even TP/EP sharding. quantization_config=FineGrainedFP8Config(dequantize=True), - tp_plan="auto", - distributed_config=DistributedConfig(enable_expert_parallel=True), + distributed_config=DistributedConfig( + tp_size=int(os.environ["WORLD_SIZE"]), + enable_expert_parallel=True, + ), attn_implementation="kernels-staging/msa@v0", # MSA block-sparse attention kernel ) model.eval() diff --git a/docs/source/en/pr_checks.md b/docs/source/en/pr_checks.md index a17d84700989..704546f24f61 100644 --- a/docs/source/en/pr_checks.md +++ b/docs/source/en/pr_checks.md @@ -123,6 +123,7 @@ Tests are split across parallel CI jobs, and each job picks up files by path pat - `pipelines_torch`: pipeline tests - `tests_training_ci`: training loop tests - `tests_tensor_parallel_ci`: tensor parallel tests +- `tests_fsdp_ci`: FSDP tests ### Slow tests diff --git a/docs/source/ro/pr_checks.md b/docs/source/ro/pr_checks.md index bceb75d73e60..e452ed4b08ac 100644 --- a/docs/source/ro/pr_checks.md +++ b/docs/source/ro/pr_checks.md @@ -123,6 +123,7 @@ Testele sunt împărțite pe job-uri CI paralele, iar fiecare job preia fișiere - `pipelines_torch`: teste de pipeline - `tests_training_ci`: teste pentru loop-ul de antrenare - `tests_tensor_parallel_ci`: teste de tensor parallelism +- `tests_fsdp_ci`: teste FSDP ### Testele slow diff --git a/pyproject.toml b/pyproject.toml index f9b1a5f17a33..1ff95cdd7904 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -83,6 +83,7 @@ markers = [ "generate: marks tests that use the GenerationTesterMixin", "is_training_test: marks tests that use the TrainingTesterMixin (deselect with '-m \"not is_training_test\"')", "is_tensor_parallel_test: marks tests that use the TensorParallelTesterMixin (deselect with '-m \"not is_tensor_parallel_test\"')", + "is_fsdp_test: marks tests that use the FSDPTesterMixin (deselect with '-m \"not is_fsdp_test\"')", ] log_cli = 1 log_cli_level = "WARNING" diff --git a/src/transformers/configuration_utils.py b/src/transformers/configuration_utils.py index 777e4c9cf98d..3bbff67b24df 100755 --- a/src/transformers/configuration_utils.py +++ b/src/transformers/configuration_utils.py @@ -1044,9 +1044,6 @@ 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] @@ -1193,6 +1190,7 @@ def _remove_keys_not_serialized(self, d: dict[str, Any]) -> None: "ignore_keys_at_rope_validation", "base_model_tp_plan", "base_model_pp_plan", + "distributed_config", ]: d.pop(key_to_remove, None) diff --git a/src/transformers/distributed/__init__.py b/src/transformers/distributed/__init__.py index e179a2306156..eb770a5f8399 100644 --- a/src/transformers/distributed/__init__.py +++ b/src/transformers/distributed/__init__.py @@ -20,6 +20,15 @@ _import_structure = { "configuration_utils": ["DistributedConfig"], "fsdp": ["is_fsdp_enabled", "is_fsdp_managed_module", "verify_fsdp_plan"], + "mixin": ["DistributedMixin"], + "utils": [ + "distribute_model", + "gather_full_state_dict", + "initialize_fully_sharded_data_parallelism", + "load_optimizer_distributed", + "save_model_checkpoint_distributed", + "save_optimizer_distributed", + ], } @@ -28,6 +37,15 @@ DistributedConfig, ) from .fsdp import is_fsdp_enabled, is_fsdp_managed_module, verify_fsdp_plan + from .mixin import DistributedMixin + from .utils import ( + distribute_model, + gather_full_state_dict, + initialize_fully_sharded_data_parallelism, + load_optimizer_distributed, + save_model_checkpoint_distributed, + save_optimizer_distributed, + ) else: import sys diff --git a/src/transformers/distributed/configuration_utils.py b/src/transformers/distributed/configuration_utils.py index 5eee35ef0fa5..946b61b6547a 100644 --- a/src/transformers/distributed/configuration_utils.py +++ b/src/transformers/distributed/configuration_utils.py @@ -63,12 +63,35 @@ def __post_init__(self): 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})" - ) + if self.tp_size > 1 and self.fsdp_size > 1: + raise ValueError( + "FSDP+TP is not supported yet. " + "Use DistributedConfig(fsdp_size=N) or DistributedConfig(tp_size=N), not both. " + "2D support will come soon." + ) + + def validate(self) -> None: + """Validate against the live process group. Call before distributed load/train.""" + if self.tp_size is None and self.fsdp_size is None: + return + + if self.tp_size <= 1 and self.fsdp_size <= 1: + return + + if not is_torch_available(): + raise RuntimeError("PyTorch is required to use DistributedConfig.") + + if not torch.distributed.is_available() or not torch.distributed.is_initialized(): + raise RuntimeError( + "torch.distributed must be initialized before using DistributedConfig with tp_size > 1 or " + "fsdp_size > 1. Call dist.init_process_group(...) first, or launch with torchrun." + ) + + 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})" + ) @classmethod def from_dict(cls, config_dict: dict, **kwargs) -> "DistributedConfig": diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py index 1eacfa209317..c9ffd66fff7a 100644 --- a/src/transformers/distributed/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -30,7 +30,7 @@ if is_torch_available(): import torch -if is_torch_available() and is_torch_greater_or_equal("2.6"): +if is_torch_available() and is_torch_greater_or_equal("2.7"): from torch.distributed._composable.fsdp import fully_shard from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy @@ -191,7 +191,7 @@ def verify_fsdp_plan(module_names: list[str], fsdp_plan: dict[str, str] | None) logger.warning(f"The following FSDP rules were not applied to any module: {unused_rules}") -def apply_fully_sharded_data_parallel( +def apply_fully_sharded_data_parallelism( model: nn.Module, fsdp_mesh: torch.distributed.device_mesh.DeviceMesh ) -> nn.Module: """ @@ -200,8 +200,8 @@ def apply_fully_sharded_data_parallel( 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") + if not is_torch_greater_or_equal("2.7"): + raise OSError("FSDP2 requires torch>=2.7") fsdp_plan = dict(getattr(model, "_fsdp_plan", None) or {}) if not fsdp_plan: diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py new file mode 100644 index 000000000000..c41f70466d28 --- /dev/null +++ b/src/transformers/distributed/mixin.py @@ -0,0 +1,274 @@ +# Copyright 2026 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 json +import os +import re +import warnings +from typing import TYPE_CHECKING + +from ..integrations.tensor_parallel import ( + ALL_PARALLEL_STYLES, + gather_state_dict_for_save, + initialize_tensor_parallelism, +) +from ..utils import SAFE_WEIGHTS_INDEX_NAME, is_torch_available, is_torch_greater_or_equal, logging +from ..utils.hub import create_and_tag_model_card +from .configuration_utils import DistributedConfig +from .fsdp import is_fsdp_managed_module +from .utils import ( + _get_torch_distributed_rank, + _is_torch_distributed_initialized, + distribute_model, + gather_full_state_dict, + initialize_fully_sharded_data_parallelism, + save_model_checkpoint_distributed, +) + + +logger = logging.get_logger(__name__) + + +if TYPE_CHECKING: + import torch.nn as nn + +if is_torch_available(): + import torch + + _torch_distributed_available = torch.distributed.is_available() +else: + _torch_distributed_available = False + + +class DistributedMixin: + """Distributed save/load hooks for [`PreTrainedModel`]. + + Stateless heavy lifting stays in `transformers.distributed.*` and + `integrations.tensor_parallel`. This mixin owns orchestration and instance state. + """ + + _device_mesh = None + _tp_plan: dict[str, str] | None = None + _tp_size = None + _pp_plan: dict[str, tuple[str, str]] = None + _fsdp_plan: dict[str, str] | None = None + + @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: + if not self._ep_plan: + raise ValueError( + f"Expert parallelism was requested (`enable_expert_parallel=True`), but " + f"`{self.__class__.__name__}` does not define an expert-parallel plan. Add a " + f"`base_model_ep_plan` to its config, or disable expert parallelism." + ) + 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 + + @tp_plan.setter + def tp_plan(self, plan: dict[str, str] | None): + if plan is None: + self._tp_plan = {} + return + if not isinstance(plan, dict): + raise ValueError("Can only set a dictionary as `tp_plan`") + + 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())}" + ) + + model_param_names = [name for name, _ in self.named_parameters()] + for layer_pattern in plan.keys(): + regex_pattern = layer_pattern.replace("*", r"\d+") + pattern_matched = False + for param_name in model_param_names: + if re.match(regex_pattern, param_name): + pattern_matched = True + break + if not pattern_matched: + warnings.warn( + f"Layer pattern '{layer_pattern}' does not match any parameters in the model. This rule may not " + "be applied during tensor parallelization, or may lead to dimension mismatches" + ) + + self._tp_plan = plan + + @pp_plan.setter + def pp_plan(self, plan: dict[str, tuple[str, str]] | None): + if plan is None: + self._pp_plan = {} + return + if not isinstance(plan, dict): + raise ValueError("Can only set a dictionary as `pp_plan`") + + self._pp_plan = plan + + @classmethod + def prepare_distribute_model( + cls, + distributed_config: DistributedConfig | dict | None, + *, + device_mesh=None, + device_map=None, + ) -> tuple[DistributedConfig | None, object, object]: + """Parse ``distributed_config``, init TP/FSDP mesh, and validate.""" + if distributed_config is None: + return None, device_map, device_mesh + + if isinstance(distributed_config, dict): + distributed_config = DistributedConfig.from_dict(distributed_config) + + if distributed_config.tp_size > 1: + if distributed_config.tp_plan is None: + distributed_config.tp_plan = "auto" + device_map, device_mesh = initialize_tensor_parallelism( + distributed_config.tp_plan, + tp_size=distributed_config.tp_size, + device_mesh=device_mesh, + device_map=device_map, + ) + elif distributed_config.fsdp_size > 1: + device_map, device_mesh = initialize_fully_sharded_data_parallelism(distributed_config) + + distributed_config.validate() + return distributed_config, device_map, device_mesh + + @classmethod + def maybe_distribute_model( + cls, + model: nn.Module, + distributed_config: DistributedConfig | None, + device_mesh, + ): + """Apply TP or FSDP2 after model init, before weight loading.""" + if _torch_distributed_available and device_mesh is not None: + return distribute_model(model, distributed_config, device_mesh) + return model + + def should_save_on_this_rank(self, is_main_process: bool) -> bool: + """Return whether this rank should write checkpoint files.""" + save_on_this_rank = is_main_process + if _is_torch_distributed_initialized(): + save_on_this_rank = save_on_this_rank and _get_torch_distributed_rank() == 0 + return save_on_this_rank + + def save_distributed_checkpoint( + self, + model_to_save, + save_directory: str | os.PathLike, + *, + push_to_hub: bool = False, + save_on_this_rank: bool = True, + repo_id: str | None = None, + files_timestamps: dict | None = None, + commit_message: str | None = None, + token: str | bool | None = None, + create_pr: bool = False, + ) -> None: + """Save an FSDP-wrapped model via DCP and optionally push to the Hub.""" + if not is_torch_greater_or_equal("2.7"): + raise OSError("save_pretrained(..., distributed_checkpoint=True) requires torch>=2.7.") + if not is_fsdp_managed_module(model_to_save): + raise ValueError( + "save_pretrained(..., distributed_checkpoint=True) is only supported for FSDP-wrapped models." + ) + if getattr(model_to_save, "_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(model_to_save, save_directory) + + if push_to_hub and save_on_this_rank: + model_card = create_and_tag_model_card(repo_id, self.model_tags, token=token) + model_card.save(os.path.join(save_directory, "README.md")) + self._upload_modified_files( + save_directory, + repo_id, + files_timestamps, + commit_message=commit_message, + token=token, + create_pr=create_pr, + ) + + def save_gathered_checkpoint_index( + self, + save_directory: str | os.PathLike, + *, + state_dict_split, + weight_map: dict | None, + weights_name: str, + variant: str | None, + max_shard_size: int | str, + ) -> None: + """Write sharded safetensors index after a gathered state-dict save.""" + # Save index if sharded + index = None + if state_dict_split.is_sharded: + index = { + "metadata": {"total_parameters": self.num_parameters(), **state_dict_split.metadata}, + "weight_map": weight_map, + } + + 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: + from ..modeling_utils import _add_variant + + 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}." + ) + + def _gather_tp_state_dict_for_save( + self, + local_state_dict: dict, + *, + is_checkpoint_writer: bool = True, + ) -> tuple[dict, bool]: + """All-gather TP-sharded weights for checkpoint writing.""" + full_state_dict = gather_state_dict_for_save(local_state_dict, self._tp_plan, self._device_mesh, self._tp_size) + if not is_checkpoint_writer: + full_state_dict = {} + return full_state_dict, True + + def _gather_fsdp_state_dict_for_save(self, model_to_save) -> tuple[dict, bool]: + """Gather FSDP-sharded weights to full CPU tensors on rank 0.""" + if not _is_torch_distributed_initialized(): + raise ValueError( + "Saving an FSDP-wrapped model requires torch.distributed to be initialized. " + "Call save_pretrained from every rank after init_process_group." + ) + return gather_full_state_dict(model_to_save), True diff --git a/src/transformers/distributed/utils.py b/src/transformers/distributed/utils.py new file mode 100644 index 000000000000..3ed5eab1a16d --- /dev/null +++ b/src/transformers/distributed/utils.py @@ -0,0 +1,228 @@ +# Copyright 2026 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 ..integrations.tensor_parallel import apply_tensor_parallelism +from ..utils import is_torch_available, is_torch_greater_or_equal +from .fsdp import apply_fully_sharded_data_parallelism + + +if TYPE_CHECKING: + import torch.nn as nn + + from .configuration_utils import DistributedConfig + +if is_torch_available(): + import torch + + _torch_distributed_available = torch.distributed.is_available() +else: + _torch_distributed_available = False + +if is_torch_available() and is_torch_greater_or_equal("2.7"): + import torch.distributed.checkpoint as dcp + from torch.distributed.checkpoint.hf_storage import HuggingFaceStorageWriter + from torch.distributed.checkpoint.state_dict import ( + StateDictOptions, + get_model_state_dict, + get_optimizer_state_dict, + set_optimizer_state_dict, + ) + + +def _is_torch_distributed_initialized() -> bool: + if not _torch_distributed_available: + return False + return hasattr(torch.distributed, "is_initialized") and torch.distributed.is_initialized() + + +def _get_torch_distributed_rank() -> int: + if not _is_torch_distributed_initialized(): + return 0 + return torch.distributed.get_rank() + + +def _get_torch_distributed_world_size() -> int: + if not _is_torch_distributed_initialized() or not hasattr(torch.distributed, "get_world_size"): + return 1 + return torch.distributed.get_world_size() + + +def is_local_dist_rank_0() -> bool: + return _is_torch_distributed_initialized() and int(os.environ.get("LOCAL_RANK", "-1")) == 0 + + +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 _is_torch_distributed_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 initialize_fully_sharded_data_parallelism(distributed_config: DistributedConfig): + 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.7"): + raise OSError("FSDP2 requires `torch>=2.7`.") + + device_type = torch._C._get_accelerator().type + _ensure_torch_distributed(device_type) + + world_size = torch.distributed.get_world_size() + if device_type != "cpu": + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + getattr(torch, device_type).set_device(local_rank) + device_map = torch.device(device_type, local_rank) + else: + device_map = torch.device(device_type) + + fsdp_size = distributed_config.fsdp_size + + assert world_size == fsdp_size, f"world_size ({world_size}) must be equal to fsdp_size ({fsdp_size})" + + dims, names = [], [] + if fsdp_size > 1: + dims.append(fsdp_size) + names.append("fsdp") + + # 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 device_map, mesh + + +def distribute_model( + model, + distributed_config: DistributedConfig, + device_mesh, +) -> nn.Module: + """Apply TP or FSDP2 to `model` based on ``distributed_config`` (mutually exclusive for now).""" + model.config.distributed_config = distributed_config + model._device_mesh = device_mesh + + if distributed_config.tp_size > 1: + model = apply_tensor_parallelism( + model, + distributed_config.tp_plan, + distributed_config, + device_mesh, + ) + elif distributed_config.fsdp_size > 1: + fsdp_mesh = device_mesh["fsdp"] if device_mesh.ndim > 1 else device_mesh + model = apply_fully_sharded_data_parallelism(model, fsdp_mesh) + + return model + + +def gather_full_state_dict(model) -> dict[str, torch.Tensor]: + """Gather FSDP-sharded params to full plain CPU tensors. + + Only rank 0 accumulates the result; other ranks return ``{}``. + """ + options = StateDictOptions(full_state_dict=True, cpu_offload=True) + full_state_dict = get_model_state_dict(model, options=options) + if _get_torch_distributed_rank() == 0: + return full_state_dict + return {} + + +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. + """ + 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) + 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 save_optimizer_distributed(model, optimizer, checkpoint_dir: str) -> None: + """Save optimizer state via DCP.""" + optimizer_state_dict = get_optimizer_state_dict(model, optimizer) + 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.""" + optimizer_state_dict = get_optimizer_state_dict(model, optimizer) + dcp.load({"optimizer": optimizer_state_dict}, checkpoint_id=checkpoint_dir) + set_optimizer_state_dict(model, optimizer, optimizer_state_dict) diff --git a/src/transformers/generation/continuous_batching/distributed.py b/src/transformers/generation/continuous_batching/distributed.py index fd4bd78a2338..86871e8040bd 100644 --- a/src/transformers/generation/continuous_batching/distributed.py +++ b/src/transformers/generation/continuous_batching/distributed.py @@ -161,7 +161,8 @@ def maybe_warn_nccl_graph_mixing(self) -> None: if tp_on and graph_mixing_not_disabled: logger.warning( "NCCL_GRAPH_MIXING_SUPPORT was not set to '0' before init_process_group: performance will be harmed. " - "Construct your `ContinuousBatchingConfig(...)` BEFORE calling `from_pretrained(tp_plan='auto')`, or " + "Construct your `ContinuousBatchingConfig(...)` BEFORE calling " + "`from_pretrained(distributed_config=DistributedConfig(tp_size=...))`, or " "set NCCL_GRAPH_MIXING_SUPPORT=0 in the launch environment." ) diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index 02e5d910b589..d19e0dcac3c5 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -108,7 +108,6 @@ def initialize_tensor_parallelism( tp_device = torch.device(device_type) device_map = device_type or {} - tp_size = tp_size if tp_size is not None else torch.distributed.get_world_size() device_mesh = torch.distributed.init_device_mesh(tp_device.type, (tp_size,)) else: if device_mesh.ndim > 1: @@ -118,10 +117,9 @@ def initialize_tensor_parallelism( "Please provide a valid `device_mesh`." ) device_mesh = device_mesh["tp"] - tp_size = device_mesh.size() device_map = torch.device(f"{device_mesh.device_type}:{int(os.environ['LOCAL_RANK'])}") - return device_map, device_mesh, tp_size + return device_map, device_mesh def replace_layer_number_by_wildcard(name: str) -> str: @@ -1609,9 +1607,9 @@ def verify_tp_plan(expected_keys: list[str], tp_plan: dict[str, str] | None): logger.warning(f"The following layers were not sharded: {', '.join(unsharded_layers)}") -def distribute_model(model, tp_plan, distributed_config, device_mesh, tp_size): - """Distribute a model according to the TP plan.""" - model._tp_size = tp_size +def apply_tensor_parallelism(model, tp_plan, distributed_config, device_mesh): + """Apply tensor parallelism to a model according to the TP plan.""" + model._tp_size = distributed_config.tp_size model._device_mesh = device_mesh if distributed_config is not None: if isinstance(distributed_config, dict): diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 3f14d74979fe..699ac3499680 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -16,7 +16,6 @@ import copy import functools import inspect -import json import os import re import sys @@ -52,6 +51,14 @@ revert_weight_conversion, ) from .distributed import DistributedConfig +from .distributed.fsdp import is_fsdp_managed_module +from .distributed.mixin import DistributedMixin +from .distributed.utils import ( + _distributed_barrier, + _get_torch_distributed_world_size, + _is_torch_distributed_initialized, + is_local_dist_rank_0, +) 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 @@ -76,11 +83,7 @@ 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, ) @@ -195,24 +198,6 @@ def is_quantized(self) -> bool: return self.hf_quantizer is not None -def _is_torch_distributed_initialized() -> bool: - return ( - _torch_distributed_available - and hasattr(torch.distributed, "is_initialized") - and torch.distributed.is_initialized() - ) - - -def _get_torch_distributed_world_size() -> int: - if not _is_torch_distributed_initialized() or not hasattr(torch.distributed, "get_world_size"): - return 1 - return torch.distributed.get_world_size() - - -def is_local_dist_rank_0(): - return _is_torch_distributed_initialized() and int(os.environ.get("LOCAL_RANK", "-1")) == 0 - - @contextmanager def set_quantized_state(): global _is_quantized @@ -1196,7 +1181,9 @@ def set_output_embeddings(self, new_embeddings): self.lm_head = new_embeddings -class PreTrainedModel(nn.Module, EmbeddingAccessMixin, ModuleUtilsMixin, PushToHubMixin, PeftAdapterMixin): +class PreTrainedModel( + nn.Module, EmbeddingAccessMixin, ModuleUtilsMixin, PushToHubMixin, PeftAdapterMixin, DistributedMixin +): r""" Base class for all models. @@ -1257,22 +1244,6 @@ class PreTrainedModel(nn.Module, EmbeddingAccessMixin, ModuleUtilsMixin, PushToH # Model's compatible flash kernels (e.g., "kernels-community/flash-mla") defaulting to the first in the list _compatible_flash_implementations: list[str] | None = None - # Tensor-parallelism-related properties - # A tensor parallel plan of the form `{"model.layer.mlp.param": "colwise"}` to be applied to the model when TP is enabled. - # 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 - # 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 - # models, this attribute is currently defined in respective model code. For base models, it comes from - # `config.base_model_pp_plan` during `post_init`. - _pp_plan: dict[str, tuple[str, str]] = None - # FSDP2 sharding plan of the form `{"layers.*": "free_full_weight"}`. For top-level models, this attribute is - # defined on the head class (e.g. `*ForCausalLM`). For base models, it comes from `config.base_model_fsdp_plan` - # during `post_init`. - _fsdp_plan: dict[str, str] = None - # Advanced functionalities support supports_gradient_checkpointing: bool = False _can_compile_fullgraph: bool = False @@ -1476,75 +1447,6 @@ def post_init(self): self.init_weights() self._backward_compatibility_gradient_checkpointing() - @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: - if not self._ep_plan: - raise ValueError( - f"Expert parallelism was requested (`enable_expert_parallel=True`), but " - f"`{self.__class__.__name__}` does not define an expert-parallel plan. Add a " - f"`base_model_ep_plan` to its config, or disable expert parallelism." - ) - 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 - - @tp_plan.setter - def tp_plan(self, plan: dict[str, str] | None): - if plan is None: - self._tp_plan = {} - return - 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()] - for layer_pattern in plan.keys(): - # Convert pattern to regex (replace * with .*) - regex_pattern = layer_pattern.replace("*", r"\d+") - pattern_matched = False - for param_name in model_param_names: - if re.match(regex_pattern, param_name): - pattern_matched = True - break - if not pattern_matched: - warnings.warn( - f"Layer pattern '{layer_pattern}' does not match any parameters in the model. This rule may not " - "be applied during tensor parallelization, or may lead to dimension mismatches" - ) - - # Set the plan - self._tp_plan = plan - - @pp_plan.setter - def pp_plan(self, plan: dict[str, tuple[str, str]] | None): - if plan is None: - self._pp_plan = {} - return - if not isinstance(plan, dict): - raise ValueError("Can only set a dictionary as `pp_plan`") - - self._pp_plan = plan - def dequantize(self, dtype=None): """ Potentially dequantize the model in case it has been quantized by a quantization method that support @@ -3414,6 +3316,7 @@ def save_pretrained( token: str | bool | None = None, save_peft_format: bool = True, save_original_format: bool = True, + distributed_checkpoint: bool = False, **kwargs, ): """ @@ -3437,7 +3340,7 @@ def save_pretrained( namespace). max_shard_size (`int` or `str`, *optional*, defaults to `"50GB"`): The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size - lower than this size. If expressed as a string, needs to be digits followed by a unit (like `"5MB"`). + lower than this size. If expressed as a string, needs be digits followed by a unit (like `"5MB"`). @@ -3459,6 +3362,11 @@ def save_pretrained( For backward compatibility with the previous versions of `transformers` you can save the checkpoint with its reverse mapping. The reverse mapping needs to exists even if the model was loaded from a None legacy checkpoint. + distributed_checkpoint (`bool`, *optional*, defaults to `False`): + When saving an FSDP-wrapped model, use the distributed checkpoint (DCP) path instead of gathering weights + to CPU first. Every rank must call this method; rank 0 writes the consolidated Hugging Face safetensors. + When `False`, FSDP weights are gathered to CPU on rank 0 via `gather_full_state_dict` before writing. + Native FSDP requires `torch>=2.7`. kwargs (`dict[str, Any]`, *optional*): Additional key word arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method. """ @@ -3505,6 +3413,8 @@ def save_pretrained( # Only save the model itself if we are using distributed training model_to_save = unwrap_model(self) + save_on_this_rank = self.should_save_on_this_rank(is_main_process) + # save the string version of dtype to the config, e.g. convert torch.float32 => "float32" # we currently don't use this setting automatically, but may start to use with v5 dtype = model_to_save.dtype @@ -3516,11 +3426,11 @@ 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.is_remote_code(): + if save_on_this_rank and self.is_remote_code(): custom_object_save(self, save_directory, config=self.config) # Save the config - if is_main_process: + if save_on_this_rank: if not _hf_peft_config_loaded: model_to_save.config.save_pretrained(save_directory) if self.can_generate(): @@ -3553,10 +3463,36 @@ def save_pretrained( current_peft_config = self.peft_config[active_adapter] current_peft_config.save_pretrained(save_directory) - # Get the model state_dict + if distributed_checkpoint: + hub_kwargs = {} + if push_to_hub: + hub_kwargs = { + "repo_id": repo_id, + "files_timestamps": files_timestamps, + "commit_message": commit_message, + "create_pr": create_pr, + } + self.save_distributed_checkpoint( + model_to_save, + save_directory, + push_to_hub=push_to_hub, + save_on_this_rank=save_on_this_rank, + token=token, + **hub_kwargs, + ) + return + + needs_dist_barrier = False if state_dict is None: state_dict = model_to_save.state_dict() + if self._tp_size is not None: + state_dict, needs_dist_barrier = self._gather_tp_state_dict_for_save( + state_dict, is_checkpoint_writer=save_on_this_rank + ) + elif is_fsdp_managed_module(model_to_save): + state_dict, needs_dist_barrier = self._gather_fsdp_state_dict_for_save(model_to_save) + # if any model parameters are offloaded, we need to know it for later is_offloaded = False if ( @@ -3581,10 +3517,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) @@ -3622,7 +3554,7 @@ 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) @@ -3639,73 +3571,59 @@ def save_pretrained( ) # 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() - - # As explained above, for offloaded scenarios, weight format could not be reverted before due to meta weights, - # so do it now after they were loaded onto cpu. For one-weight-to-many operations, it may be an issue, but usually the shards - # contain all the necessary params, except if we are quite unlucky on the sharding. The failure surface is (very few models - # with one-weight-to-many + offloading to disk + unlucky sharding), so it will almost never happen - if is_offloaded and save_original_format and not _hf_peft_config_loaded: - try: - shard_state_dict = revert_weight_conversion(model_to_save, shard_state_dict) - # Save the weight_map, since some names etc may have changed due to conversion compared to initial `state_dict_split` - if state_dict_split.is_sharded: - weight_map.update({k: os.path.basename(shard_file)} for k in shard_state_dict.keys()) # ty: ignore[unresolved-attribute] - except Exception: - raise RuntimeError( - "We could not revert some weight conversions because of offlading, and several weights needed for a single " - "conversion operation living in different shard files. Try reducing `max_shard_size` a bit, or worst case " - "set `save_original_format=False`." - ) - - # 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 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() + + # As explained above, for offloaded scenarios, weight format could not be reverted before due to meta weights, + # so do it now after they were loaded onto cpu. For one-weight-to-many operations, it may be an issue, but usually the shards + # contain all the necessary params, except if we are quite unlucky on the sharding. The failure surface is (very few models + # with one-weight-to-many + offloading to disk + unlucky sharding), so it will almost never happen + if is_offloaded and save_original_format and not _hf_peft_config_loaded: + try: + shard_state_dict = revert_weight_conversion(model_to_save, shard_state_dict) + # Save the weight_map, since some names etc may have changed due to conversion compared to initial `state_dict_split` + if state_dict_split.is_sharded: + weight_map.update({k: os.path.basename(shard_file)} for k in shard_state_dict.keys()) # ty: ignore[unresolved-attribute] + except Exception: + raise RuntimeError( + "We could not revert some weight conversions because of offlading, and several weights needed for a single " + "conversion operation living in different shard files. Try reducing `max_shard_size` a bit, or worst case " + "set `save_original_format=False`." + ) - # Save index if sharded - index = None - if state_dict_split.is_sharded: - index = { - "metadata": {"total_parameters": self.num_parameters(), **state_dict_split.metadata}, - "weight_map": weight_map, - } + # 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}." + self.save_gathered_checkpoint_index( + save_directory, + state_dict_split=state_dict_split, + weight_map=weight_map, + weights_name=weights_name, + variant=variant, + max_shard_size=max_shard_size, ) - 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) @@ -3721,6 +3639,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 needs_dist_barrier: + _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 [] @@ -4117,6 +4042,11 @@ 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. + distributed_config ([`~transformers.distributed.configuration_utils.DistributedConfig`], *optional*): + Configuration for native distributed loading with tensor parallelism or FSDP2. Pass + `DistributedConfig(tp_size=N)` for tensor parallelism, or + `DistributedConfig(fsdp_size=N)` for FSDP2. Requires `torchrun` and an initialized + process group when `tp_size > 1` or `fsdp_size > 1`. Mutually exclusive with `device_map`. 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. @@ -4217,8 +4147,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) @@ -4227,9 +4155,6 @@ 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" - # 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"]: _ = kwargs.pop(name, None) @@ -4266,9 +4191,9 @@ def from_pretrained( ": 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 + if distributed_config is not None: + distributed_config, device_map, device_mesh = cls.prepare_distribute_model( + distributed_config, device_mesh=device_mesh, device_map=device_map ) if gguf_file is not None and not is_accelerate_available(): @@ -4313,6 +4238,9 @@ def from_pretrained( model_kwargs = kwargs commit_hash = getattr(config, "_commit_hash", commit_hash) + if distributed_config is not None: + config.distributed_config = distributed_config + download_kwargs_with_commit["commit_hash"] = commit_hash # Because some composite configs call super().__init__ before instantiating the sub-configs, we need this call @@ -4422,8 +4350,7 @@ def from_pretrained( # Obtain the weight conversion mapping for this model if any are registered and apply to all submodels recursively 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) + model = cls.maybe_distribute_model(model, distributed_config, device_mesh) # Prepare the full device map if device_map is not None: diff --git a/src/transformers/models/cohere2_moe/configuration_cohere2_moe.py b/src/transformers/models/cohere2_moe/configuration_cohere2_moe.py index 32ef6fdf66ad..83af4aa21894 100644 --- a/src/transformers/models/cohere2_moe/configuration_cohere2_moe.py +++ b/src/transformers/models/cohere2_moe/configuration_cohere2_moe.py @@ -90,6 +90,11 @@ class Cohere2MoeConfig(PreTrainedConfig): "layers.*.mlp.experts.down_proj": "grouped_gemm", "layers.*.mlp.experts": "moe_tp_experts", } + 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 diff --git a/src/transformers/models/cohere2_moe/modeling_cohere2_moe.py b/src/transformers/models/cohere2_moe/modeling_cohere2_moe.py index e0290b835948..bdd0ed5b596f 100644 --- a/src/transformers/models/cohere2_moe/modeling_cohere2_moe.py +++ b/src/transformers/models/cohere2_moe/modeling_cohere2_moe.py @@ -583,6 +583,7 @@ class Cohere2MoeForCausalLM(Cohere2MoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_gather_output"} _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_moe/modular_cohere2_moe.py b/src/transformers/models/cohere2_moe/modular_cohere2_moe.py index b5d4c3b49f57..57c76cb0805a 100644 --- a/src/transformers/models/cohere2_moe/modular_cohere2_moe.py +++ b/src/transformers/models/cohere2_moe/modular_cohere2_moe.py @@ -279,6 +279,8 @@ def forward( class Cohere2MoeForCausalLM(Cohere2ForCausalLM): + _fsdp_plan = {"lm_head": "keep_full_weight"} + def forward( self, input_ids: torch.LongTensor | None = None, diff --git a/src/transformers/testing_utils.py b/src/transformers/testing_utils.py index 287386b1d5e7..577ab154d3ee 100644 --- a/src/transformers/testing_utils.py +++ b/src/transformers/testing_utils.py @@ -322,6 +322,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): @@ -404,6 +405,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. diff --git a/tests/causal_lm_tester.py b/tests/causal_lm_tester.py index 3bb99dddd7ed..2efe6d1bcb70 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/generation/test_continuous_batching.py b/tests/generation/test_continuous_batching.py index 84784a95fbb6..e1bb046923c7 100644 --- a/tests/generation/test_continuous_batching.py +++ b/tests/generation/test_continuous_batching.py @@ -1740,13 +1740,14 @@ def _tp_continuous_batching_worker( use_cuda_graph: bool, use_async_batching: bool, ) -> None: - """Loads `model_id` with `tp_plan="auto"`, checks three TP-specific paths in the same process: (a) direct + """Loads `model_id` with `DistributedConfig(tp_size=...)`, checks three TP-specific paths in the same process: (a) direct broadcasts via `DistributedHelper`, (b) per-rank parity of CB-generated tokens via `dist.all_gather_object`, and (c) reproducibility across two CB runs sharing the same seed. Rank 0 owns all the assertions; the other ranks only need to participate in the collectives.""" import torch import torch.distributed as dist + from transformers.distributed import DistributedConfig from transformers.generation.continuous_batching.distributed import DistributedHelper tokenizer = AutoTokenizer.from_pretrained(model_id, padding_side="left") @@ -1754,7 +1755,10 @@ def _tp_continuous_batching_worker( tokenizer.pad_token = tokenizer.eos_token model = AutoModelForCausalLM.from_pretrained( - model_id, attn_implementation=attn_implementation, tp_plan="auto", dtype=torch.float32 + model_id, + attn_implementation=attn_implementation, + distributed_config=DistributedConfig(tp_size=int(os.environ["WORLD_SIZE"])), + dtype=torch.float32, ).eval() # Direct broadcast tests: only rank 0's value should propagate to every TP rank @@ -1832,6 +1836,8 @@ def _tp_cancellation_worker( import torch + from transformers.distributed import DistributedConfig + cb_config = ContinuousBatchingConfig(use_cuda_graph=use_cuda_graph, use_async_batching=use_async_batching) tokenizer = AutoTokenizer.from_pretrained(model_id, padding_side="left") @@ -1839,7 +1845,10 @@ def _tp_cancellation_worker( tokenizer.pad_token = tokenizer.eos_token model = AutoModelForCausalLM.from_pretrained( - model_id, attn_implementation=attn_implementation, tp_plan="auto", dtype=torch.float32 + model_id, + attn_implementation=attn_implementation, + distributed_config=DistributedConfig(tp_size=int(os.environ["WORLD_SIZE"])), + dtype=torch.float32, ).eval() chat = [{"role": "user", "content": "Tell me a long story about a robot exploring the galaxy."}] diff --git a/tests/models/deepseek_v4/test_modeling_deepseek_v4.py b/tests/models/deepseek_v4/test_modeling_deepseek_v4.py index 3cb4205d7567..23aca5b1c4ed 100644 --- a/tests/models/deepseek_v4/test_modeling_deepseek_v4.py +++ b/tests/models/deepseek_v4/test_modeling_deepseek_v4.py @@ -468,7 +468,9 @@ def main() -> int: dtype="auto", attn_implementation="eager", experts_implementation=LOADTIME_DISPATCH, - distributed_config=DistributedConfig(enable_expert_parallel=True), + distributed_config=DistributedConfig( + tp_size=int(os.environ["WORLD_SIZE"]), enable_expert_parallel=True + ), ) model.eval() tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) @@ -545,7 +547,9 @@ def main() -> int: dtype="auto", attn_implementation="eager", experts_implementation=LOADTIME_DISPATCH, - distributed_config=DistributedConfig(enable_expert_parallel=True), + distributed_config=DistributedConfig( + tp_size=int(os.environ["WORLD_SIZE"]), enable_expert_parallel=True + ), ) model.eval() tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) diff --git a/tests/models/glm_moe_dsa/test_modeling_glm_moe_dsa.py b/tests/models/glm_moe_dsa/test_modeling_glm_moe_dsa.py index aa87525eede2..45f760fc2f2d 100644 --- a/tests/models/glm_moe_dsa/test_modeling_glm_moe_dsa.py +++ b/tests/models/glm_moe_dsa/test_modeling_glm_moe_dsa.py @@ -13,6 +13,7 @@ # limitations under the License. """Testing suite for the PyTorch GlmMoeDsa model.""" +import os import unittest import torch @@ -27,6 +28,7 @@ is_torch_available, set_seed, ) +from transformers.distributed import DistributedConfig from transformers.testing_utils import ( require_torch, require_torch_accelerator, @@ -204,7 +206,7 @@ def test_glm_moe_dsa_fp8_inference(self): model = AutoModelForCausalLM.from_pretrained( model_id, quantization_config=quantization_config, - tp_plan="auto", + distributed_config=DistributedConfig(tp_size=int(os.environ["WORLD_SIZE"])), attn_implementation="eager", ) diff --git a/tests/models/gpt_oss/test_modeling_gpt_oss.py b/tests/models/gpt_oss/test_modeling_gpt_oss.py index 5b3c13f8c25a..1dcd911bdbe4 100644 --- a/tests/models/gpt_oss/test_modeling_gpt_oss.py +++ b/tests/models/gpt_oss/test_modeling_gpt_oss.py @@ -159,6 +159,7 @@ def distributed_worker(quantized, model_size, kernels, attn_impl, mode): import os from transformers import AutoModelForCausalLM, AutoTokenizer + from transformers.distributed import DistributedConfig from transformers.testing_utils import torch_device def generate_config_key(quantized, model, kernels, attn_impl, mode): @@ -179,7 +180,7 @@ def generate_config_key(quantized, model, kernels, attn_impl, mode): model = AutoModelForCausalLM.from_pretrained( model_id, dtype="auto", - tp_plan="auto", # distributed inference + distributed_config=DistributedConfig(tp_size=int(os.environ["WORLD_SIZE"])), use_kernels=kernels, ).to(torch_device) model.set_attn_implementation(attn_impl) diff --git a/tests/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py new file mode 100644 index 000000000000..2a966328fc22 --- /dev/null +++ b/tests/test_fsdp_mixin.py @@ -0,0 +1,627 @@ +# Copyright 2026 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 contextlib import contextmanager + +from parameterized import parameterized + +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_torch_greater_or_equal, +) +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 +DDP_FSDP_RTOL = 1e-5 +DDP_FSDP_ATOL = 1e-5 + +# Set to None to run distributed FSDP tests for every model with a plan. +FSDP_DISTRIBUTED_TEST_MODEL_TYPES = {"cohere2_moe"} + + +# ============================================================================= +# 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_available_fsdp_workers(): + if _get_distributed_device_type() == "cpu": + return os.cpu_count() or 1 + return backend_device_count(_get_distributed_device_type()) + + +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) + + +@contextmanager +def _distributed_tmpdir(rank): + if rank == 0: + tmpdir_obj = tempfile.TemporaryDirectory() + tmpdir_list = [tmpdir_obj.name] + else: + tmpdir_obj = None + tmpdir_list = [None] + dist.broadcast_object_list(tmpdir_list, src=0) + try: + yield tmpdir_list[0] + finally: + if rank == 0 and tmpdir_obj is not None: + tmpdir_obj.cleanup() + + +@contextmanager +def _deterministic_init_model_dir(rank, config, dtype): + with _distributed_tmpdir(rank) as model_dir: + if rank == 0: + set_seed(SEED) + model = AutoModelForCausalLM.from_config(config).to(dtype) + model.save_pretrained(model_dir) + del model + dist.barrier() + yield model_dir + + +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() + + +# ============================================================================= +# Training 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 _run_training_steps(model, optimizer, batches, *, track_grad_norms=True): + """Forward/backward/step over batches. Returns (losses, grad_norms).""" + losses, grad_norms = [], [] + for input_ids, labels in batches: + optimizer.zero_grad() + loss = model(input_ids=input_ids, labels=labels, use_cache=False).loss + loss.backward() + if track_grad_norms: + grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=float("inf")) + grad_norms.append(grad_norm) + optimizer.step() + losses.append(loss.detach().item()) + return losses, grad_norms + + +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()} + accelerator_module = backend_torch_accelerator_module(_get_distributed_device_type()) + if accelerator_module is not None and hasattr(accelerator_module, "get_rng_state"): + accel_rng = accelerator_module.get_rng_state() + if accel_rng is not None: + rng["accel"] = accel_rng + 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: + accelerator_module = backend_torch_accelerator_module(_get_distributed_device_type()) + if accelerator_module is not None and hasattr(accelerator_module, "set_rng_state"): + accelerator_module.set_rng_state(rng["accel"]) + + +def _checkpoint_and_resume(pre_model, pre_optimizer, dtype, distributed_config, lr): + """Save model+optimizer, scramble RNG, reload and restore training state.""" + rank = dist.get_rank() + with _distributed_tmpdir(rank) as tmpdir: + model_dir = os.path.join(tmpdir, "model") + training_state_dir = os.path.join(tmpdir, "training_state") + + pre_model.save_pretrained(model_dir, is_main_process=(rank == 0)) + _save_training_state(pre_model, pre_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() + return resumed_model, resumed_optimizer + + +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. + 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 = _run_training_steps(ddp_model, optimizer, batches) + if dist.get_rank() != 0: + state_dict = {} + else: + # Only rank 0 returns data to match gather_full_state_dict semantics. + state_dict = {k: v.clone().detach().cpu() for k, v in ddp_model.module.state_dict().items()} + + 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, +): + distributed_config = DistributedConfig(fsdp_size=dist.get_world_size()) + + # Phase 1: Pre-checkpoint run + _set_determinism(SEED) + 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 = _run_training_steps( + pre_ckpt_model, pre_ckpt_optimizer, batches[:checkpoint_step] + ) + + # Phase 2: Save checkpoint, then load into a fresh model + resumed_model, resumed_optimizer = _checkpoint_and_resume( + pre_ckpt_model, pre_ckpt_optimizer, dtype, distributed_config, lr + ) + + # Phase 3: Post-checkpoint run + post_ckpt_losses, post_ckpt_grad_norms = _run_training_steps( + resumed_model, resumed_optimizer, batches[checkpoint_step:] + ) + + return ( + pre_ckpt_losses + post_ckpt_losses, + pre_ckpt_grad_norms + post_ckpt_grad_norms, + gather_full_state_dict(resumed_model), + ) + + +# ============================================================================= +# Distributed test implementations (top-level for pickling by mp.spawn) +# ============================================================================= + + +def _test_fsdp2_save_load_impl(rank, config_class, config_dict): + """Save FSDP2 model via save_pretrained, load via from_pretrained, compare state dicts.""" + init_test_logger() + + config = config_class.from_dict(config_dict) + distributed_config = DistributedConfig(fsdp_size=dist.get_world_size()) + + with _deterministic_init_model_dir(rank, config, torch.float32) as init_dir: + _set_determinism(SEED) + model = AutoModelForCausalLM.from_pretrained(init_dir, distributed_config=distributed_config) + dist.barrier() + + state_dict_before = gather_full_state_dict(model) + + with _distributed_tmpdir(rank) as tmpdir: + model.save_pretrained(tmpdir, is_main_process=(rank == 0)) + dist.barrier() + new_model = AutoModelForCausalLM.from_pretrained(tmpdir, distributed_config=distributed_config) + dist.barrier() + + state_dict_after = gather_full_state_dict(new_model) + for key in state_dict_before: + assert key in state_dict_after, f"After save/load: Key {key} missing after load" + torch.testing.assert_close( + state_dict_before[key], + state_dict_after[key], + rtol=0, + atol=0, + msg=f"After save/load: Weight mismatch for {key}", + ) + + if rank == 0: + logger.debug(f"FSDP2 save/load test passed: all {len(state_dict_before)} parameters match exactly.") + + +def _test_fsdp2_save_load_dcp_impl(rank, config_class, config_dict): + """Save FSDP2 model via save_pretrained(distributed_checkpoint=True), reload, compare state dicts.""" + init_test_logger() + + config = config_class.from_dict(config_dict) + distributed_config = DistributedConfig(fsdp_size=dist.get_world_size()) + + with _deterministic_init_model_dir(rank, config, torch.float32) as init_dir: + _set_determinism(SEED) + model = AutoModelForCausalLM.from_pretrained(init_dir, distributed_config=distributed_config) + dist.barrier() + + state_dict_before = gather_full_state_dict(model) + + with _distributed_tmpdir(rank) as tmpdir: + model.save_pretrained(tmpdir, is_main_process=(rank == 0), distributed_checkpoint=True) + dist.barrier() + new_model = AutoModelForCausalLM.from_pretrained(tmpdir, distributed_config=distributed_config) + dist.barrier() + + state_dict_after = gather_full_state_dict(new_model) + for key in state_dict_before: + assert key in state_dict_after, f"After DCP save/load: Key {key} missing after load" + torch.testing.assert_close( + state_dict_before[key], + state_dict_after[key], + rtol=0, + atol=0, + msg=f"After DCP save/load: Weight mismatch for {key}", + ) + + if rank == 0: + logger.debug(f"FSDP2 DCP 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_sharded_data_parallel wraps exactly the right modules.""" + init_test_logger() + + config = config_class.from_dict(config_dict) + config.tie_word_embeddings = tie_word_embeddings + distributed_config = DistributedConfig(fsdp_size=dist.get_world_size()) + + with _deterministic_init_model_dir(rank, config, torch.float32) as init_dir: + _set_determinism(SEED) + model = AutoModelForCausalLM.from_pretrained(init_dir, distributed_config=distributed_config) + dist.barrier() + + adapted_fsdp_plan = _resolve_tied_embed_lm_head_plan(model._fsdp_plan, model) + reshard_targets, no_reshard_targets = expand_fsdp_plan(model, adapted_fsdp_plan) + expected_targets = {""} | {name for name, _ in reshard_targets + no_reshard_targets} + actual_targets = {name for name, module in model.named_modules() if type(module).__name__.startswith("FSDP")} + + if rank == 0: + logger.debug(f" Weights tied: {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, 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 + + device = _get_rank_device(rank) + config = config_class.from_dict(config_dict) + config.tie_word_embeddings = tie_word_embeddings + + checkpoint_step = NUM_STEPS // 2 + batches = _build_repeated_training_batches(config, device, NUM_STEPS) + + with _deterministic_init_model_dir(rank, config, dtype) as init_model_dir: + 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, + ) + + for step in range(len(ddp_losses)): + torch.testing.assert_close( + torch.tensor(ddp_losses[step]), + torch.tensor(fsdp_losses[step]), + rtol=DDP_FSDP_RTOL, + atol=DDP_FSDP_ATOL, + msg=f"Loss mismatch at step {step}: DDP={ddp_losses[step]}, FSDP2={fsdp_losses[step]}", + ) + torch.testing.assert_close( + torch.tensor(ddp_grad_norms[step]), + torch.tensor(fsdp_grad_norms[step]), + rtol=DDP_FSDP_RTOL, + atol=DDP_FSDP_ATOL, + msg=f"Grad norm mismatch at step {step}: DDP={ddp_grad_norms[step]}, FSDP2={fsdp_grad_norms[step]}", + ) + + for key in ddp_state_dict: + assert key in fsdp_state_dict, f"Key {key} missing from FSDP2 state dict" + torch.testing.assert_close( + ddp_state_dict[key], + fsdp_state_dict[key], + rtol=DDP_FSDP_RTOL, + atol=DDP_FSDP_ATOL, + msg=f"Weight mismatch for {key}: DDP vs FSDP2", + ) + + if rank == 0: + logger.debug("DDP and FSDP2 comparison checks passed.") + + +# ============================================================================= +# Mixin class +# ============================================================================= + + +class FSDPTesterMixin(ABC): + fsdp_nproc_per_node: int = 2 + # TODO(3outeille): do we put the CONSTANTS in the mixin class ? + + @property + @abstractmethod + def model_tester(self): + """The model tester instance (e.g., CausalLMModelTester).""" + ... + + def _skip_if_insufficient_devices(self): + 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 _has_fsdp_plan(self) -> bool: + config = self.model_tester.get_config() + return hasattr(config, "base_model_fsdp_plan") and config.base_model_fsdp_plan is not None + + def _skip_if_fsdp_distributed_not_enabled(self): + if not self._has_fsdp_plan(): + self.skipTest("Model does not have an FSDP plan (base_model_fsdp_plan)") + + config = self.model_tester.get_config() + # Only top-N models are tested, set FSDP_DISTRIBUTED_TEST_MODEL_TYPES = None to run all tests. + if ( + FSDP_DISTRIBUTED_TEST_MODEL_TYPES is not None + and config.model_type not in FSDP_DISTRIBUTED_TEST_MODEL_TYPES + ): + self.skipTest( + f"FSDP distributed tests are not enabled for model_type={config.model_type!r} " + f"(enabled: {sorted(FSDP_DISTRIBUTED_TEST_MODEL_TYPES)}). Set FSDP_DISTRIBUTED_TEST_MODEL_TYPES = None to run all tests." + ) + + def _create_model_on_meta(self, config): + """Instantiate a model on the meta device (no memory allocated).""" + auto_classes = [AutoModelForCausalLM, AutoModelForSeq2SeqLM] # TODO(3outeille): why 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"): + 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, "moe_intermediate_size"): + config.moe_intermediate_size = 32 + if hasattr(config, "vocab_size_per_layer_input"): + config.vocab_size_per_layer_input = config.vocab_size + return type(config), config.to_diff_dict() + + def _run_fsdp2_distributed_test(self, test_name, test_impl, *test_args, **test_kwargs): + self._skip_if_insufficient_devices() + self._skip_if_fsdp_distributed_not_enabled() + + config_class, config_dict = self._get_tiny_config() + func_args = (config_class, config_dict, *test_args) + + results_file = tempfile.mktemp(suffix=".json") + # port binding + 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']}") + + @is_fsdp_test + def test_fsdp_plan_declared(self): + """The model exposes a non-empty `_fsdp_plan` derived from config + class-level overrides.""" + if not self._has_fsdp_plan(): + self.skipTest("Model does not have an FSDP plan (base_model_fsdp_plan)") + + model = self._create_model_on_meta(self.model_tester.get_config()) + self.assertTrue(model._fsdp_plan, f"No _fsdp_plan declared for {type(model).__name__}") + + @parameterized.expand(["untied", "tied"]) + @require_torch_greater_or_equal("2.7") + @is_fsdp_test + def test_fsdp2_sharding_structure(self, label): + self._run_fsdp2_distributed_test( + f"test_fsdp2_sharding_structure_{label}", + _test_fsdp2_sharding_structure_impl, + label == "tied", + ) + + @require_torch_greater_or_equal("2.7") + @is_fsdp_test + def test_fsdp2_save_load(self): + self._run_fsdp2_distributed_test("test_fsdp2_save_load", _test_fsdp2_save_load_impl) + + @require_torch_greater_or_equal("2.7") + @is_fsdp_test + def test_fsdp2_save_load_dcp(self): + self._run_fsdp2_distributed_test("test_fsdp2_save_load_dcp", _test_fsdp2_save_load_dcp_impl) + + @parameterized.expand(["untied", "tied"]) + @require_torch_greater_or_equal("2.7") + @is_fsdp_test + def test_fsdp2_plan_vs_ddp(self, label): + self._run_fsdp2_distributed_test( + f"test_fsdp2_plan_vs_ddp_{label}", + _test_fsdp2_plan_vs_ddp_impl, + label == "tied", + ) diff --git a/tests/test_tensor_parallel_mixin.py b/tests/test_tensor_parallel_mixin.py index a40649ad883c..42b23d9ce2c9 100644 --- a/tests/test_tensor_parallel_mixin.py +++ b/tests/test_tensor_parallel_mixin.py @@ -118,7 +118,10 @@ def _load_tp_and_reference_models(model_path, model_class): Returns: tuple: (model_tp, model_ref, device) """ - model_tp = model_class.from_pretrained(model_path, tp_plan="auto") + model_tp = model_class.from_pretrained( + model_path, + distributed_config=DistributedConfig(tp_size=dist.get_world_size()), + ) dist.barrier() device = model_tp.device @@ -146,7 +149,7 @@ def _verify_tp_sharding(rank, model_tp, model_ref): # 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) + 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, ( @@ -227,7 +230,7 @@ def _test_tp_backward_impl(rank, model_path, model_class, atol, rtol): 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) + 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) @@ -297,7 +300,11 @@ def _test_tp_generation_quantized_impl(_rank, model_path, model_class, max_new_t quantization_config = TorchAoConfig(Float8WeightOnlyConfig()) - model_tp = model_class.from_pretrained(model_path, tp_plan="auto", quantization_config=quantization_config) + model_tp = model_class.from_pretrained( + model_path, + distributed_config=DistributedConfig(tp_size=dist.get_world_size()), + quantization_config=quantization_config, + ) dist.barrier() device = model_tp.device @@ -349,7 +356,7 @@ 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), + distributed_config=DistributedConfig(tp_size=dist.get_world_size(), enable_expert_parallel=True), ) dist.barrier() diff --git a/tests/trainer/distributed/test_trainer_distributed_fsdp.py b/tests/trainer/distributed/test_trainer_distributed_fsdp.py index 250ef854dcfc..aaaafe18d820 100644 --- a/tests/trainer/distributed/test_trainer_distributed_fsdp.py +++ b/tests/trainer/distributed/test_trainer_distributed_fsdp.py @@ -144,7 +144,7 @@ def test_move_missing_keys_fsdp_non_rank0_moves_meta_to_cpu(self): with ( patch("transformers.modeling_utils.is_fsdp_enabled", return_value=True), - patch("transformers.modeling_utils.is_local_dist_rank_0", return_value=False), + patch("transformers.distributed.utils.is_local_dist_rank_0", return_value=False), ): model._move_missing_keys_from_meta_to_device( missing_keys=set(), device_map=None, device_mesh=None, hf_quantizer=None @@ -161,7 +161,7 @@ def test_fsdp_non_rank0_end_to_end_no_reinit(self): with ( patch("transformers.modeling_utils.is_fsdp_enabled", return_value=True), - patch("transformers.modeling_utils.is_local_dist_rank_0", return_value=False), + patch("transformers.distributed.utils.is_local_dist_rank_0", return_value=False), ): model._move_missing_keys_from_meta_to_device( missing_keys=set(), device_map=None, device_mesh=None, hf_quantizer=None diff --git a/utils/tests_fetcher.py b/utils/tests_fetcher.py index 14fdc786ade7..dc4318128926 100644 --- a/utils/tests_fetcher.py +++ b/utils/tests_fetcher.py @@ -1100,6 +1100,7 @@ def parse_commit_message(commit_message: str) -> dict[str, bool]: "tests_non_model": r"tests/(?!peft_integration/)[^/]*?/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)", "tests_peft_integration": r"tests/peft_integration/test_.*\.py", }