From 7c843391b762cfe7c36882980992fcd557cac672 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 25 Mar 2026 09:09:34 +0000 Subject: [PATCH 001/116] init --- tmp.py | 1 + 1 file changed, 1 insertion(+) create mode 100644 tmp.py diff --git a/tmp.py b/tmp.py new file mode 100644 index 000000000000..05db1b840b44 --- /dev/null +++ b/tmp.py @@ -0,0 +1 @@ +"&&" From a5c25548e9cdb29b079d35692835e8920bc2586a Mon Sep 17 00:00:00 2001 From: 3outeille Date: Mon, 13 Apr 2026 14:10:58 +0000 Subject: [PATCH 002/116] FSDP2 (fully_shard) integration - Add apply_fully_shard_data_parallel() with auto/manual mode block detection - FSDP vs DDP loss/grad parity tests - Distributed test helpers (testing_utils.py) - is_fsdp_enabled(), is_fsdp_managed_module() utilities - Minimal FSDP hooks in from_pretrained - FSDP-aware flash attention check --- src/transformers/distributed/utils.py | 50 ++ src/transformers/integrations/__init__.py | 4 +- src/transformers/integrations/accelerate.py | 2 +- src/transformers/integrations/fsdp.py | 515 ++++++++++- src/transformers/integrations/moe.py | 2 +- .../modeling_flash_attention_utils.py | 19 +- src/transformers/modeling_utils.py | 35 +- src/transformers/testing_utils.py | 118 ++- tests/causal_lm_tester.py | 8 +- tests/test_fsdp_mixin.py | 849 ++++++++++++++++++ tests/test_modeling_common.py | 10 +- tests/test_tensor_parallel_mixin.py | 4 +- 12 files changed, 1561 insertions(+), 55 deletions(-) create mode 100644 src/transformers/distributed/utils.py create mode 100644 tests/test_fsdp_mixin.py diff --git a/src/transformers/distributed/utils.py b/src/transformers/distributed/utils.py new file mode 100644 index 000000000000..40278c8fcf2a --- /dev/null +++ b/src/transformers/distributed/utils.py @@ -0,0 +1,50 @@ +# 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, strtobool + + +if TYPE_CHECKING: + import torch.nn as nn + +if is_torch_available(): + import torch + + +def is_fsdp_enabled() -> bool: + 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: + if not is_torch_available(): + return False + if not torch.distributed.is_available(): + return False + try: + from torch.distributed.fsdp import FullyShardedDataParallel + except ImportError: + return False + return isinstance(module, FullyShardedDataParallel) or getattr(module, "_is_fsdp_managed_module", False) diff --git a/src/transformers/integrations/__init__.py b/src/transformers/integrations/__init__.py index 336db3773f76..e3515eab24b1 100755 --- a/src/transformers/integrations/__init__.py +++ b/src/transformers/integrations/__init__.py @@ -50,7 +50,7 @@ "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"], + "fsdp": ["is_fsdp_enabled"], "ggml": [ "GGUF_CONFIG_DEFAULTS_MAPPING", "GGUF_CONFIG_MAPPING", @@ -209,7 +209,7 @@ 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 .fsdp import is_fsdp_enabled from .ggml import ( GGUF_CONFIG_DEFAULTS_MAPPING, GGUF_CONFIG_MAPPING, diff --git a/src/transformers/integrations/accelerate.py b/src/transformers/integrations/accelerate.py index c2b7fa603570..a5f9835ad1ca 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.utils 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 index 7cda7ad55acc..10936490ee03 100644 --- a/src/transformers/integrations/fsdp.py +++ b/src/transformers/integrations/fsdp.py @@ -15,46 +15,519 @@ import inspect import os -from typing import TYPE_CHECKING +from typing import Any, Literal -from ..utils import is_torch_available, strtobool +from ..utils import is_torch_available, is_torch_greater_or_equal, logging from ..utils.quantization_config import QuantizationMethod -if TYPE_CHECKING: - from torch import nn +if is_torch_available() and is_torch_greater_or_equal("2.5"): + import torch + import torch.distributed as dist + import torch.distributed.checkpoint as dcp + from torch.distributed._composable.fsdp import fully_shard + from torch.distributed.checkpoint.hf_storage import HuggingFaceStorageWriter + from torch.distributed.checkpoint.state_dict import get_model_state_dict + from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy, OffloadPolicy + from torch.distributed.tensor import DTensor + +logger = logging.get_logger(__name__) -def is_fsdp_managed_module(module: nn.Module) -> bool: +def initialize_fsdp( + fsdp_plan: dict[str, Any] | None, + device_mesh=None, + device_map=None, +): + """ + Sets up the device mesh for FSDP2 (Fully Sharded Data Parallel). + This function is called when the model is loaded and fsdp_plan is set. + + Args: + fsdp_plan: Optional FSDP config dict with an explicit "mode". + device_mesh: Optional pre-created DeviceMesh for FSDP. + device_map: Optional device map. + + Returns: + Tuple of (device_map, device_mesh, fsdp_size) + """ if not is_torch_available(): - return False + raise ImportError("PyTorch is required for FSDP support") - import torch + if fsdp_plan is None: + return device_map, device_mesh, None + + if not is_torch_greater_or_equal("2.5"): + raise OSError("FSDP2 is only supported for `torch>=2.5`.") + + if device_mesh is None: + # Detect the accelerator on the machine + device_type = torch._C._get_accelerator().type + current_device = getattr(torch, device_type) + + if not dist.is_initialized(): + try: + rank = int(os.environ["RANK"]) + local_rank = int(os.environ["LOCAL_RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + + backend_map = {"cuda": "nccl", "cpu": "gloo", "xpu": "xccl", "hpu": "hccl"} + backend = backend_map.get(device_type) + if device_type == "cpu" and int(os.environ.get("CCL_WORKER_COUNT", "0")): + backend = "ccl" + if device_type == "xpu" and not is_torch_greater_or_equal("2.8", accept_dev=True): + backend = "ccl" + + dist.init_process_group(backend=backend, rank=rank, world_size=world_size) + if device_type != "cpu": + current_device.set_device(local_rank) + + except Exception as e: + raise OSError( + "We tried to initialize torch.distributed for you, but it failed. Make " + "sure you init torch distributed in your script to use `fsdp_plan`." + ) from e + + if device_type != "cpu": + current_device.set_device(int(os.environ["LOCAL_RANK"])) + index = current_device.current_device() + fsdp_device = torch.device(device_type, index) + device_map = fsdp_device + else: + fsdp_device = torch.device(device_type) + device_map = device_type or {} + + fsdp_size = dist.get_world_size() + device_mesh = torch.distributed.init_device_mesh(fsdp_device.type, (fsdp_size,), mesh_dim_names=("dp_shard",)) + else: + # Use provided device mesh + if device_mesh.ndim > 1: + if "dp_shard" not in device_mesh.mesh_dim_names: + raise ValueError( + "When using `fsdp_plan` with n-d `device_mesh`, it must contain a 'dp_shard' dimension. " + "Please provide a valid `device_mesh`." + ) + device_mesh = device_mesh["dp_shard"] + fsdp_size = device_mesh.size() + device_map = torch.device(f"{device_mesh.device_type}:{int(os.environ['LOCAL_RANK'])}") + + return device_map, device_mesh, fsdp_size + + +def get_transformer_block_classes(model): + """ + Identifies transformer block classes in a model for FSDP wrapping. + These are typically the repeated layers that benefit from FSDP sharding. + + Returns a set of module classes that should be wrapped with fully_shard(). + """ + block_classes = set() + + # Common transformer block class names + block_names = { + "DecoderLayer", + "EncoderLayer", + "TransformerBlock", + "Block", + "Layer", + } + + for module in model.modules(): + class_name = module.__class__.__name__ + # Use endswith to avoid false positives (e.g. "Layer" matching "LayerNorm") + for block_name in block_names: + if class_name.endswith(block_name): + block_classes.add(type(module)) + break + + # Filter out nested block classes (e.g. SparseMoeBlock inside DecoderLayer). + # We only want to FSDP-wrap the outermost block classes. If a class like + # MoeBlock only ever appears inside a DecoderLayer, we skip it. + if len(block_classes) > 1: + # Collect the dotted module paths for each candidate class. + # i.e: {DecoderLayer: ["layers.0", "layers.1"], + # MoeBlock: ["layers.0.moe", "layers.1.moe"]} + paths_by_class = {} + for name, module in model.named_modules(): + cls = type(module) + if cls in block_classes: + paths_by_class.setdefault(cls, []).append(name) - if not torch.distributed.is_available(): - return False + def _is_nested_inside_other_class(cls): + # A class is "inner" if every one of its instances lives under + # an instance of a different candidate class in the module tree. + paths = paths_by_class.get(cls, []) + if not paths: + return False + for path in paths: + has_parent = any( + path.startswith(parent_path + ".") + for other_cls, parent_paths in paths_by_class.items() + if other_cls is not cls + for parent_path in parent_paths + ) + if not has_parent: + return False + return True - import torch.distributed.fsdp + # Keep only the outer (non-nested) classes. + block_classes = {cls for cls in block_classes if not _is_nested_inside_other_class(cls)} + + return block_classes + + +def _get_auto_policy_kwargs(fsdp_plan: dict[str, Any]) -> dict[str, Any]: + """Parse auto-mode fsdp_plan into fully_shard policy kwargs.""" + policy_kwargs = {} + if fsdp_plan.get("cpu_offload"): + policy_kwargs["offload_policy"] = CPUOffloadPolicy() + if fsdp_plan.get("mixed_precision"): + policy_kwargs["mp_policy"] = MixedPrecisionPolicy( + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + output_dtype=None, + ) + return policy_kwargs + + +def _auto_shard_input_embedding(input_embed, is_weights_tied: bool, device_mesh, auto_policy_kwargs): + # Shard input embeddings (only when not tied). + # When tied, the shared weight is grouped with the final norm in step 3. + if input_embed is None or is_weights_tied: + return + fully_shard(input_embed, mesh=device_mesh, reshard_after_forward=True, **auto_policy_kwargs) + logger.debug(f"Applied fully_shard to input embeddings ({type(input_embed).__name__})") + + +def _auto_shard_transformer_blocks(model, block_classes, device_mesh, auto_policy_kwargs): + for name, module in model.named_modules(): + if type(module) in block_classes: + fully_shard(module, mesh=device_mesh, reshard_after_forward=True, **auto_policy_kwargs) + logger.debug(f"Applied fully_shard to {name} ({type(module).__name__})") + + +def _find_final_norm(model, decoder_layer_names): + """Find the final normalization layer before the output head. + + Searches only within the base model scope (e.g. ``model.*``) so that + norms inside the output head / prediction head (e.g. ``lm_head.norm``) + are excluded. + """ + base_prefix = model.base_model_prefix # e.g. "model" + final_norm = None + for name, module in model.named_modules(): + if "Norm" not in type(module).__name__: + continue + # Only consider norms inside the base model (skip root-level heads like lm_head) + if base_prefix and not name.startswith(base_prefix + ".") and name != base_prefix: + continue + if any(name.startswith(layer_name + ".") for layer_name in decoder_layer_names): + continue + final_norm = module + return final_norm + + +def _auto_get_tail_modules(model, decoder_layer_names, input_embed, output_embed, is_weights_tied: bool) -> list: + # Group final norm + output head. + # NOTE(3outeille): Small optimization by forcing reshard_after_forward=False for the final norm and output head. + # Otherwise, that would mean reshard/freeing full params after the last forward and immediately re-all-gathering + # them in the backward pass, which is wasteful. Better to keep them gathered for reuse. + # Untied: [final_norm, lm_head] + # Tied: [final_norm, embed_tokens] - embed_tokens.weight IS lm_head.weight. + tail_modules = [] + + final_norm = _find_final_norm(model, decoder_layer_names) + + if final_norm is not None: + tail_modules.append(final_norm) + + if is_weights_tied: + if input_embed is not None: + tail_modules.append(input_embed) + elif output_embed is not None: + tail_modules.append(output_embed) + + return tail_modules + + +def _auto_shard_tail_modules(tail_modules, device_mesh, auto_policy_kwargs): + if len(tail_modules) > 1: + fully_shard(tail_modules, mesh=device_mesh, reshard_after_forward=False, **auto_policy_kwargs) + logger.debug(f"Applied fully_shard to {[type(m).__name__ for m in tail_modules]} grouped (reshard=False)") + elif len(tail_modules) == 1: + fully_shard(tail_modules[0], mesh=device_mesh, reshard_after_forward=False, **auto_policy_kwargs) + logger.debug(f"Applied fully_shard to {type(tail_modules[0]).__name__} (reshard=False)") + + +def _parse_manual_plan_entry( + entry: list[str], +) -> tuple[bool, MixedPrecisionPolicy | None, CPUOffloadPolicy | None]: + """ + Returns: + tuple[bool, MixedPrecisionPolicy | None, CPUOffloadPolicy | None]: + - bool: whether to reshard after forward + - MixedPrecisionPolicy | None: mixed precision policy + - CPUOffloadPolicy | None: cpu offload policy + """ - return isinstance(module, torch.distributed.fsdp.FullyShardedDataParallel) or getattr( - module, "_is_fsdp_managed_module", False + if not isinstance(entry, list): + raise ValueError( + f"Manual fsdp_plan values must be a list of strings combining strategy/policies, got {type(entry)}" + ) + items = entry + + strategy: Literal["free_full_weight", "keep_full_weight"] | None = None + offload_policy: CPUOffloadPolicy | None = None + mp_policy: MixedPrecisionPolicy | None = None + + for item in items: + if not isinstance(item, str): + raise ValueError( + f"fsdp_plan option must be a string, got {type(item)}. " + "Supported: 'free_full_weight', 'keep_full_weight', 'cpu_offload', 'mixed_precision'." + ) + token = item.lower() + if token in {"free_full_weight", "keep_full_weight"}: + strategy = token + elif token == "cpu_offload": + offload_policy = CPUOffloadPolicy() + elif token == "mixed_precision": + # TODO(3outeille): add support for different dtypes + mp_policy = MixedPrecisionPolicy( + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + output_dtype=torch.bfloat16, + ) + else: + raise ValueError( + "Unknown fsdp_plan option " + f"{item!r}. Supported: 'free_full_weight', 'keep_full_weight', 'cpu_offload', 'mixed_precision'." + ) + + if strategy is None: + strategy = "free_full_weight" + + return strategy != "keep_full_weight", mp_policy, offload_policy + + +def _iter_manual_plan_targets(model, pattern, name_to_module, already_sharded_names): + if pattern in name_to_module: + target = name_to_module[pattern] + if isinstance(target, (torch.nn.ModuleList, torch.nn.ModuleDict, torch.nn.Sequential)): + # (ModuleList, ModuleDict, Sequential) don't have a forward() that gets called - + # the model loops over their children directly. So when a pattern matches a + # container, we shard each child instead. + for child_name, child in target.named_children(): + yield f"{pattern}.{child_name}", child + else: + yield pattern, target + return + + # Prefix match: "model.layers" matches "model.layers.0", etc. + for name, module in model.named_modules(): + if name in already_sharded_names or isinstance( + module, (torch.nn.ModuleList, torch.nn.ModuleDict, torch.nn.Sequential) + ): + continue + if name != pattern and not name.startswith(pattern + "."): + continue + if any( + name.startswith(already_sharded_names_name + ".") for already_sharded_names_name in already_sharded_names + ): + continue + yield name, module + + +def _parse_fsdp_plan_mode(fsdp_plan: dict[str, Any]) -> Literal["auto", "manual"]: + if isinstance(fsdp_plan, str): + fsdp_plan = {"mode": fsdp_plan} + + if not isinstance(fsdp_plan, dict): + raise ValueError(f"fsdp_plan must be a dict with a 'mode' key, got {type(fsdp_plan)}") + + mode = fsdp_plan.get("mode") + if mode not in {"auto", "manual"}: + raise ValueError("fsdp_plan['mode'] must be either 'auto' or 'manual'.") + + return mode + + +def _get_manual_plan_modules(fsdp_plan: dict[str, Any]) -> dict[str, list[str]]: + modules = fsdp_plan.get("modules") + if not isinstance(modules, dict): + raise ValueError("Manual fsdp_plan must define a 'modules' dict.") + return modules + + +def apply_fsdp2( + model, + device_mesh, + fsdp_plan: dict[str, Any] | str | None, +): + """ + Apply FSDP2 (fully_shard) to a model following TorchTitan's approach. + fsdp_plan: + Explicit FSDP config dict with a required "mode" key, or a string + shorthand such as "auto" or "manual". + + Auto mode: + fsdp_plan = "auto" + + Auto mode (equivalent): + fsdp_plan = {"mode": "auto"} + + Auto mode with optional policies: + fsdp_plan = {"mode": "auto", "cpu_offload": False, "mixed_precision": True} + + Manual mode: + fsdp_plan = { + "mode": "manual", + "modules": { + "model.embed_tokens": ["free_full_weight"], + "model.layers.0.self_attn": ["free_full_weight", "cpu_offload", "mixed_precision"], + "model.layers.0.mlp": ["free_full_weight"], + "model.norm": ["keep_full_weight"], + "lm_head": ["keep_full_weight"], + }, + } + """ + if not is_torch_available(): + raise ImportError("PyTorch is required for FSDP support") + + if not is_torch_greater_or_equal("2.5"): + raise OSError("FSDP2 requires torch>=2.5") + + if device_mesh is None: + raise ValueError("device_mesh is required for FSDP2") + + if isinstance(fsdp_plan, str): + fsdp_plan = {"mode": fsdp_plan} + + input_embed = getattr(model, "get_input_embeddings", lambda: None)() + output_embed = getattr(model, "get_output_embeddings", lambda: None)() + is_weights_tied = ( + input_embed is not None + and output_embed is not None + and hasattr(input_embed, "weight") + and hasattr(output_embed, "weight") + and input_embed.weight is output_embed.weight ) + mode = _parse_fsdp_plan_mode(fsdp_plan) + + if mode == "auto": + auto_policy_kwargs = _get_auto_policy_kwargs(fsdp_plan) + + block_classes = get_transformer_block_classes(model) + # Need to collect decoder layer names for norm detection. + decoder_layer_names = {name for name, module in model.named_modules() if type(module) in block_classes} -def is_fsdp_enabled(): - if is_torch_available(): - import torch + if not block_classes: + logger.warning( + "Could not auto-detect transformer block classes for FSDP. Applying FSDP only to root module." + ) + else: + _auto_shard_input_embedding(input_embed, is_weights_tied, device_mesh, auto_policy_kwargs) - 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 + _auto_shard_transformer_blocks(model, block_classes, device_mesh, auto_policy_kwargs) + + tail_modules = _auto_get_tail_modules( + model, decoder_layer_names, input_embed, output_embed, is_weights_tied + ) + _auto_shard_tail_modules(tail_modules, device_mesh, auto_policy_kwargs) + + # Shard root model + fully_shard(model, mesh=device_mesh, **auto_policy_kwargs) + + logger.info( + f"FSDP2 applied to model: {len(block_classes)} block type(s), {len(decoder_layer_names)} decoder layers" ) - return False + else: + # fsdp_plan = { + # "mode": "manual", + # "modules": { + # "model.layers.0.self_attn": ["free_full_weight"], # reshard_after_forward=True + # "model.norm": ["keep_full_weight"], # reshard_after_forward=False + # "model.layers.0.mlp": ["free_full_weight", "cpu_offload", "mixed_precision"], + # }, + # } + + name_to_module = dict(model.named_modules()) + already_sharded_names: set[str] = set() + root_mp_policy = MixedPrecisionPolicy() + root_offload_policy = OffloadPolicy() + + for pattern, entry in _get_manual_plan_modules(fsdp_plan).items(): + reshard, mp_policy, offload_policy = _parse_manual_plan_entry(entry) + if mp_policy is not None: + root_mp_policy = mp_policy + if offload_policy is not None: + root_offload_policy = offload_policy + + for name, module in _iter_manual_plan_targets(model, pattern, name_to_module, already_sharded_names): + if name in already_sharded_names: + continue + shard_kwargs = {"mesh": device_mesh, "reshard_after_forward": reshard} + if mp_policy is not None: + shard_kwargs["mp_policy"] = mp_policy + if offload_policy is not None: + shard_kwargs["offload_policy"] = offload_policy + fully_shard(module, **shard_kwargs) + already_sharded_names.add(name) + logger.debug(f"Applied fully_shard to {name}") + + # Shard root model with the same policies as sub-modules. + # MixedPrecisionPolicy.output_dtype casting happens in post_forward + # for every fully_shard-wrapped module, even with no direct parameters. + fully_shard(model, mesh=device_mesh, mp_policy=root_mp_policy, offload_policy=root_offload_policy) + + # Used by generation code to detect FSDP and enable synced_gpus. + model._is_fsdp_managed_module = True + + if is_weights_tied and hasattr(model, "tie_weights"): + # Re-tie weights. + # fully_shard replaces nn.Parameter objects (swapping data for DTensor shards), + # which breaks weight tying (e.g. lm_head.weight is no longer embed_tokens.weight). + # Re-tying makes lm_head._parameters["weight"] point to the new DTensor parameter + # so gradients accumulate correctly into a single buffer. + model.tie_weights() + + return model + + +# TODO(3outeille): probably remove this function. Will be handled when someone tackle PEFT + FSDP. +def save_fsdp_model(model, save_directory): + """Save FSDP2 model weights as HF safetensors via DCP distributed save + consolidation. + + Each rank saves its DTensor shard in parallel, then rank 0 consolidates + into standard HF-compatible safetensors files. + """ + model_sd = get_model_state_dict(model) + + # Clone tensors sharing storage (tied weights) — safetensors refuses aliased tensors + seen_data_ptrs = {} + for key in list(model_sd.keys()): + tensor = model_sd[key] + t = tensor._local_tensor if isinstance(tensor, DTensor) else tensor + ptr = t.data_ptr() + if ptr in seen_data_ptrs: + model_sd[key] = tensor.clone() + else: + seen_data_ptrs[ptr] = key + + dcp.save( + model_sd, + storage_writer=HuggingFaceStorageWriter( + path=save_directory, + save_distributed=True, + enable_consolidation=True, + ), + ) +# ========================= PEFT compatibility ========================= +# TODO(3outeille): make sure new FSDP works with PEFT def get_fsdp_ckpt_kwargs(): """ Returns checkpoint kwargs for FSDP model saving. diff --git a/src/transformers/integrations/moe.py b/src/transformers/integrations/moe.py index 70178dd1fa7e..df04841c3d78 100644 --- a/src/transformers/integrations/moe.py +++ b/src/transformers/integrations/moe.py @@ -357,7 +357,7 @@ def _grouped_linear( out = _grouped_mm(input, weight, offs=offs) else: # (S, input_dim) @ grouped (num_experts, output_dim, input_dim).T -> (S, output_dim) - out = _grouped_mm(input, weight.transpose(-2, -1), offs=offs) + out = _grouped_mm(input, weight.transpose(-2, -1).contiguous(), offs=offs) if bias is not None: # We should be able to pass bias to the grouped_mm call, but it's not yet supported. diff --git a/src/transformers/modeling_flash_attention_utils.py b/src/transformers/modeling_flash_attention_utils.py index 9211ccb19a9e..0454b6dfefae 100644 --- a/src/transformers/modeling_flash_attention_utils.py +++ b/src/transformers/modeling_flash_attention_utils.py @@ -74,8 +74,10 @@ def is_flash_attn_available(): 2: { "flash_attn_version": 2, "general_availability_check": is_flash_attn_2_available, - "pkg_availability_check": lambda *args, **kwargs: importlib.util.find_spec("flash_attn") is not None - and "flash-attn" in [pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn"]], + "pkg_availability_check": lambda *args, **kwargs: ( + importlib.util.find_spec("flash_attn") is not None + and "flash-attn" in [pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn"]] + ), "supported_devices": ( (is_torch_cuda_available, "cuda"), (is_torch_mlu_available, "mlu"), @@ -93,16 +95,21 @@ def is_flash_attn_available(): 3: { "flash_attn_version": 3, "general_availability_check": is_flash_attn_3_available, - "pkg_availability_check": lambda *args, **kwargs: importlib.util.find_spec("flash_attn_interface") is not None - and "flash-attn-3" in [pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn_interface"]], + "pkg_availability_check": lambda *args, **kwargs: ( + importlib.util.find_spec("flash_attn_interface") is not None + and "flash-attn-3" + in [pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn_interface"]] + ), "supported_devices": ((is_torch_cuda_available, "cuda"),), "cuda_min_major_version": 8, # Ampere }, 4: { "flash_attn_version": 4, "general_availability_check": is_flash_attn_4_available, - "pkg_availability_check": lambda *args, **kwargs: importlib.util.find_spec("flash_attn") is not None - and "flash-attn-4" in [pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn"]], + "pkg_availability_check": lambda *args, **kwargs: ( + importlib.util.find_spec("flash_attn") is not None + and "flash-attn-4" in [pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn"]] + ), "supported_devices": ((is_torch_cuda_available, "cuda"),), "cuda_min_major_version": 9, # Hopper }, diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 81ff067b1470..2a7304650640 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -51,9 +51,10 @@ revert_weight_conversion, ) from .distributed import DistributedConfig +from .distributed.utils import is_fsdp_enabled 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, @@ -68,6 +69,7 @@ from .integrations.flash_attention import flash_attention_forward from .integrations.flash_paged import paged_attention_forward from .integrations.flex_attention import flex_attention_forward +from .integrations.fsdp import apply_fsdp2 from .integrations.hub_kernels import allow_all_hub_kernels, is_kernel from .integrations.peft import maybe_load_adapters from .integrations.sdpa_attention import sdpa_attention_forward @@ -3316,6 +3318,27 @@ def save_pretrained( current_peft_config = self.peft_config[active_adapter] current_peft_config.save_pretrained(save_directory) + # FSDP2 models: use DCP distributed save + consolidation for safetensors. + # All ranks must call this collectively. Config/generation_config are + # already saved above (guarded by is_main_process). + if getattr(self, "_is_fsdp_managed_module", False): + from .integrations.fsdp import save_fsdp_model + + save_fsdp_model(model_to_save, save_directory) + + if push_to_hub: + 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, + ) + return + # Get the model state_dict if state_dict is None: state_dict = model_to_save.state_dict() @@ -3932,6 +3955,7 @@ def from_pretrained( gguf_file = kwargs.pop("gguf_file", None) tp_plan = kwargs.pop("tp_plan", None) tp_size = kwargs.pop("tp_size", None) + fsdp_plan = kwargs.pop("fsdp_plan", 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) @@ -4137,6 +4161,15 @@ def from_pretrained( model.eval() # Set model in evaluation mode to deactivate Dropout modules by default model.set_use_kernels(use_kernels, kernel_config) + # Apply FSDP2 if configured (must be after weight loading) + if fsdp_plan is not None: + if device_mesh is None: + raise ValueError( + "`fsdp_plan` was provided but no device mesh is available. " + "Pass `device_mesh` to `from_pretrained`." + ) + model = apply_fsdp2(model, device_mesh, fsdp_plan) + # If it is a model with generation capabilities, attempt to load generation files (generation config, # custom generate function) if model.can_generate() and hasattr(model, "adjust_generation_fn") and not gguf_file: diff --git a/src/transformers/testing_utils.py b/src/transformers/testing_utils.py index 863242a695c6..f4e19e558727 100644 --- a/src/transformers/testing_utils.py +++ b/src/transformers/testing_utils.py @@ -231,7 +231,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 @@ -285,6 +288,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): @@ -351,6 +355,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 @@ -367,6 +387,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. @@ -4129,6 +4165,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.""" @@ -4168,8 +4243,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 @@ -4179,7 +4255,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) @@ -4197,7 +4273,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() @@ -4208,26 +4284,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/tests/causal_lm_tester.py b/tests/causal_lm_tester.py index b3398f13c393..18c2bf082ae7 100644 --- a/tests/causal_lm_tester.py +++ b/tests/causal_lm_tester.py @@ -29,6 +29,7 @@ ) from .test_configuration_common import ConfigTester +from .test_fsdp_mixin import FSDPTesterMixin from .test_modeling_common import ( GenerationTesterMixin, ModelTesterMixin, @@ -307,7 +308,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/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py new file mode 100644 index 000000000000..d7a0a4ca3340 --- /dev/null +++ b/tests/test_fsdp_mixin.py @@ -0,0 +1,849 @@ +# 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.distributed.checkpoint as dcp + import torch.multiprocessing as mp + from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner + from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict + from torch.distributed.tensor import DTensor + from torch.nn.parallel import DistributedDataParallel as DDP + + from transformers.integrations.fsdp import ( + _find_final_norm, + apply_fsdp2, + get_transformer_block_classes, + initialize_fsdp, + ) + + +# ============================================================================= +# Constants +# ============================================================================= + +BATCH_SIZE = 2 +SEQ_LEN = 64 +NUM_STEPS = 20 +LR = 3e-4 +SEED = 42 +FSDP_TOP_MODEL_NAMES = { + # Dense + "gpt2", + "qwen3", + "phi", + "llama", + "modernbert_decoder", + "olmo3", + "phi3", + "mistral", + "lfm2", + "gemma2", + # MoE + "gpt_oss", + "glm_moe_dsa", + "qwen3_moe", + "glm4_moe_lite", + "qwen3_5_moe", + "deepseek_v2", + "qwen3_next", + "mixtral", + "qwen2_moe", + "phimoe", +} + + +# ============================================================================= +# 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_fsdp2_state_dict(model): + """Gather FSDP2 sharded parameters into full tensors via DTensor.full_tensor().""" + state_dict = {} + for name, tensor in model.state_dict().items(): + if isinstance(tensor, DTensor): + state_dict[name] = tensor.full_tensor().clone().detach().cpu() + else: + state_dict[name] = tensor.clone().detach().cpu() + return state_dict + + +def _gather_ddp_state_dict(model): + return {k: v.clone().detach().cpu() for k, v in model.module.state_dict().items()} + + +def _build_manual_fsdp_plan(config, device, policy_options=None): + """Build a default manual FSDP2 plan from model structure.""" + policy_options = policy_options or [] + set_seed(SEED) + model = AutoModelForCausalLM.from_config(config).to(device) + named_modules = dict(model.named_modules()) + id_to_name = {id(module): name for name, module in named_modules.items()} + block_classes = get_transformer_block_classes(model) + + decoder_layer_names = {name for name, module in named_modules.items() if type(module) in block_classes} + layer_prefixes = {".".join(name.split(".")[:-1]) for name in decoder_layer_names} + assert layer_prefixes, "Expected at least one decoder layer prefix for manual FSDP plan." + + input_embed = model.get_input_embeddings() + output_embed = model.get_output_embeddings() + weights_tied = ( + input_embed is not None + and output_embed is not None + and hasattr(input_embed, "weight") + and hasattr(output_embed, "weight") + and input_embed.weight is output_embed.weight + ) + embed_name = id_to_name.get(id(input_embed)) if input_embed is not None else None + output_name = id_to_name.get(id(output_embed)) if output_embed is not None else None + final_norm = _find_final_norm(model, decoder_layer_names) + norm_name = id_to_name.get(id(final_norm)) if final_norm is not None else None + + module_plan = {name: ["free_full_weight", *policy_options] for name in layer_prefixes} + + if norm_name: + module_plan[norm_name] = ["keep_full_weight", *policy_options] + + if weights_tied: + if embed_name: + module_plan[embed_name] = ["keep_full_weight", *policy_options] + else: + if embed_name: + module_plan[embed_name] = ["free_full_weight", *policy_options] + if output_name: + module_plan[output_name] = ["keep_full_weight", *policy_options] + + del model + return {"mode": "manual", "modules": module_plan} + + +def _save_init_pretrained(rank, config, dtype): + """Save a deterministic initial model to a shared tmpdir for from_pretrained loading.""" + tmpdir, tmpdir_obj = _create_shared_tmpdir(rank) + if rank == 0: + set_seed(SEED) + model = AutoModelForCausalLM.from_config(config).to(dtype) + model.save_pretrained(tmpdir) + del model + dist.barrier() + return tmpdir, tmpdir_obj + + +def _save_training_state(model, optimizer, training_state_dir): + """Save optimizer + RNG states as distcp (for training resume only).""" + _, optim_sd = get_state_dict(model, optimizer) + training_state = { + "optim": optim_sd, + "cpu_rng_state": torch.get_rng_state(), + } + accelerator_rng_state = _get_accelerator_rng_state() + if accelerator_rng_state is not None: + training_state["accelerator_rng_state"] = accelerator_rng_state + dcp.save(training_state, checkpoint_id=training_state_dir) + + +def _load_training_state(model, optimizer, training_state_dir): + """Load optimizer + RNG states from distcp (model weights loaded separately via from_pretrained).""" + model_sd, optim_sd = get_state_dict(model, optimizer) + loaded_training_state = { + "optim": optim_sd, + "cpu_rng_state": torch.empty_like(torch.get_rng_state()), + } + accelerator_rng_state = _get_accelerator_rng_state() + if accelerator_rng_state is not None: + loaded_training_state["accelerator_rng_state"] = torch.empty_like(accelerator_rng_state) + # MoE models can have sparse optimizer state (experts not selected yet), so + # allow partial optimizer key restoration instead of failing hard on missing keys. + dcp.load( + loaded_training_state, + checkpoint_id=training_state_dir, + planner=DefaultLoadPlanner(allow_partial_load=True), + ) + set_state_dict( + model, + optimizer, + model_state_dict=model_sd, + optim_state_dict=loaded_training_state["optim"], + ) + torch.set_rng_state(loaded_training_state["cpu_rng_state"]) + _set_accelerator_rng_state(loaded_training_state.get("accelerator_rng_state")) + + +def train_ddp(rank, batches, lr, device, dtype, init_model_dir): + _set_determinism(SEED) + model = AutoModelForCausalLM.from_pretrained(init_model_dir, torch_dtype=dtype, attn_implementation="eager").to( + device + ) + # MoE/conditional-routing variants) may not use all params on + # every step, and DDP would otherwise fail. Specifying find_unused_parameters=True allows running backward on a subgraph of the model. + ddp_kwargs = {"find_unused_parameters": True} + if device.type != "cpu": + ddp_kwargs["device_ids"] = [rank] + ddp_model = DDP(model, **ddp_kwargs) + ddp_model.train() + optimizer = torch.optim.Adam(ddp_model.parameters(), lr=lr) + + losses, grad_norms = [], [] + for input_ids, labels in batches: + optimizer.zero_grad() + output = ddp_model(input_ids=input_ids, labels=labels, use_cache=False) + loss = output.loss + loss.backward() + grad_norm = torch.nn.utils.clip_grad_norm_(ddp_model.parameters(), max_norm=float("inf")) + optimizer.step() + + losses.append(loss.detach().item()) + grad_norms.append(grad_norm) + + state_dict = _gather_ddp_state_dict(ddp_model) + + del optimizer, ddp_model, model + backend_empty_cache(_get_distributed_device_type()) + dist.barrier() + + return losses, grad_norms, state_dict + + +def train_fsdp2( + rank, + batches, + lr, + dtype, + init_model_dir, + checkpoint_step, + fsdp_plan, +): + # -- Phase 1: Pre-checkpoint run -- train only the first `checkpoint_step` steps, then save + _set_determinism(SEED) + _, device_mesh, _ = initialize_fsdp(fsdp_plan=fsdp_plan) + pre_ckpt_model = AutoModelForCausalLM.from_pretrained( + init_model_dir, + torch_dtype=dtype, + fsdp_plan=fsdp_plan, + device_mesh=device_mesh, + attn_implementation="eager", + ) + pre_ckpt_model.train() + pre_ckpt_optimizer = torch.optim.Adam(pre_ckpt_model.parameters(), lr=lr) + + pre_ckpt_losses, pre_ckpt_grad_norms = [], [] + for step in range(0, checkpoint_step): + input_ids, labels = batches[step] + pre_ckpt_optimizer.zero_grad() + output = pre_ckpt_model(input_ids=input_ids, labels=labels, use_cache=False) + loss = output.loss + loss.backward() + grad_norm = torch.nn.utils.clip_grad_norm_(pre_ckpt_model.parameters(), max_norm=float("inf")) + pre_ckpt_optimizer.step() + + pre_ckpt_losses.append(loss.detach().item()) + pre_ckpt_grad_norms.append(grad_norm) + + # -- Phase 2: Save checkpoint, then load into a fresh model + # tmpdir/ + # model/ <- HF safetensors via save_pretrained (DCP + consolidation) + # training_state/ <- distcp (optimizer + RNG) + tmpdir, tmpdir_obj = _create_shared_tmpdir(rank) + try: + model_dir = os.path.join(tmpdir, "model") + training_state_dir = os.path.join(tmpdir, "training_state") + + pre_ckpt_model.save_pretrained(model_dir, is_main_process=(rank == 0)) + _save_training_state(pre_ckpt_model, pre_ckpt_optimizer, training_state_dir) + dist.barrier() + + # Intentionally scramble RNG to prove checkpoint restore works + _set_determinism(SEED + 1234) + resumed_model = AutoModelForCausalLM.from_pretrained( + model_dir, + torch_dtype=dtype, + fsdp_plan=fsdp_plan, + device_mesh=device_mesh, + attn_implementation="eager", + ) + resumed_model.train() + resumed_optimizer = torch.optim.Adam(resumed_model.parameters(), lr=lr) + + _load_training_state(resumed_model, resumed_optimizer, training_state_dir) + dist.barrier() + finally: + if rank == 0: + tmpdir_obj.cleanup() + + # -- Phase 3: Post-checkpoint run -- continue training the remaining steps from the resumed model + post_ckpt_losses, post_ckpt_grad_norms = [], [] + for step in range(checkpoint_step, len(batches)): + input_ids, labels = batches[step] + resumed_optimizer.zero_grad() + output = resumed_model(input_ids=input_ids, labels=labels, use_cache=False) + loss = output.loss + loss.backward() + grad_norm = torch.nn.utils.clip_grad_norm_(resumed_model.parameters(), max_norm=float("inf")) + resumed_optimizer.step() + + post_ckpt_losses.append(loss.detach().item()) + post_ckpt_grad_norms.append(grad_norm) + + combined_losses = pre_ckpt_losses + post_ckpt_losses + combined_grad_norms = pre_ckpt_grad_norms + post_ckpt_grad_norms + combined_state_dict = _gather_fsdp2_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) + + auto_plan = {"mode": "auto"} + + init_tmpdir, init_tmpdir_obj = _save_init_pretrained(rank, config, torch.float32) + try: + _, device_mesh, _ = initialize_fsdp(fsdp_plan=auto_plan) + _set_determinism(SEED) + model = AutoModelForCausalLM.from_pretrained( + init_tmpdir, + fsdp_plan=auto_plan, + device_mesh=device_mesh, + attn_implementation="eager", + ) + dist.barrier() + finally: + if rank == 0 and init_tmpdir_obj is not None: + init_tmpdir_obj.cleanup() + model.train() + optimizer = torch.optim.Adam(model.parameters(), lr=LR) + + for input_ids, labels in batches: + optimizer.zero_grad() + output = model(input_ids=input_ids, labels=labels, use_cache=False) + output.loss.backward() + optimizer.step() + + state_dict_before = _gather_fsdp2_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, + fsdp_plan=auto_plan, + device_mesh=device_mesh, + attn_implementation="eager", + ) + dist.barrier() + finally: + if rank == 0: + tmpdir_obj.cleanup() + + state_dict_after = _gather_fsdp2_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_fsdp2(fsdp_plan={"mode": "auto"}) wraps exactly the right modules. + + Expected FSDP targets: + UNTIED TIED + ────── ──── + 1. embed_tokens (reshard=True) 1. (skip — embed goes to step 3) + 2. layers[i] (reshard=True) 2. layers[i] (reshard=True) + 3. [norm, lm_head] (reshard=False) 3. [norm, embed_tokens] (reshard=False) + 4. root 4. root + """ + init_test_logger() + + config = config_class.from_dict(config_dict) + config.tie_word_embeddings = tie_word_embeddings + + auto_plan = {"mode": "auto"} + device_map, device_mesh, _ = initialize_fsdp(fsdp_plan=auto_plan) + + set_seed(SEED) + model = AutoModelForCausalLM.from_config(config).to(device_map) + + block_classes = get_transformer_block_classes(model) + assert block_classes, "get_transformer_block_classes found no block classes" + + decoder_layer_names = {name for name, module in model.named_modules() if type(module) in block_classes} + assert len(decoder_layer_names) > 0, "Expected at least one transformer block instance" + + id_to_name = {id(module): name for name, module in model.named_modules()} + + input_embed = model.get_input_embeddings() + output_embed = model.get_output_embeddings() + final_norm = _find_final_norm(model, decoder_layer_names) + weights_tied = ( + input_embed is not None + and output_embed is not None + and hasattr(input_embed, "weight") + and hasattr(output_embed, "weight") + and input_embed.weight is output_embed.weight + ) + + embed_name = id_to_name.get(id(input_embed)) + output_name = id_to_name.get(id(output_embed)) + norm_name = id_to_name.get(id(final_norm)) + + expected_targets = {""} | decoder_layer_names | {embed_name} | {norm_name} + if not weights_tied: + expected_targets |= {output_name} + + model = apply_fsdp2(model, device_mesh, fsdp_plan=auto_plan) + + actual_targets = {name for name, module in model.named_modules() if type(module).__name__.startswith("FSDP")} + + if rank == 0: + logger.debug(f" Weights tied: {weights_tied}") + logger.debug(f" Expected FSDP targets: {sorted(expected_targets)}") + logger.debug(f" Actual FSDP targets: {sorted(actual_targets)}") + + missing = expected_targets - actual_targets + extra = actual_targets - expected_targets + assert not missing and not extra, ( + f"FSDP target mismatch.\n" + f" Missing (expected but not wrapped): {sorted(missing)}\n" + f" Extra (wrapped but not expected): {sorted(extra)}" + ) + + if rank == 0: + logger.debug(f" FSDP sharding structure OK ({len(actual_targets)} targets)") + + +def _test_fsdp2_plan_vs_ddp_impl( + rank, config_class, config_dict, tie_word_embeddings, plan_mode, policy_options=None, dtype=None +): + """Validate DDP-vs-FSDP2 trace matching for either auto or manual plan mode.""" + init_test_logger() + + if dtype is None: + dtype = torch.float32 + + policy_options = policy_options or [] + assert "mixed_precision" not in policy_options, ( + "Use the mixed-precision specific tests when enabling mixed_precision policy." + ) + + device = _get_rank_device(rank) + config = config_class.from_dict(config_dict) + config.tie_word_embeddings = tie_word_embeddings + + if plan_mode == "auto": + fsdp_plan = { + "mode": "auto", + "cpu_offload": "cpu_offload" in policy_options, + "mixed_precision": "mixed_precision" in policy_options, + } + test_label = f"FSDP2(auto{'+policies' if policy_options else ''})" + elif plan_mode == "manual": + fsdp_plan = _build_manual_fsdp_plan(config, device, policy_options=policy_options) + test_label = f"FSDP2(manual{'+policies' if policy_options else ''})" + else: + raise ValueError(f"Unsupported plan_mode '{plan_mode}'. Expected 'auto' or 'manual'.") + + checkpoint_step = NUM_STEPS // 2 + init_model_dir, init_tmpdir_obj = _save_init_pretrained(rank, config, dtype) + batches = _build_repeated_training_batches(config, device, NUM_STEPS) + try: + ddp_losses, ddp_grad_norms, ddp_state_dict = train_ddp(rank, batches, LR, device, dtype, init_model_dir) + + fsdp_losses, fsdp_grad_norms, fsdp_state_dict = train_fsdp2( + rank, + batches, + LR, + dtype, + init_model_dir=init_model_dir, + checkpoint_step=checkpoint_step, + fsdp_plan=fsdp_plan, + ) + finally: + if rank == 0 and init_tmpdir_obj is not None: + init_tmpdir_obj.cleanup() + + for step in range(len(ddp_losses)): + torch.testing.assert_close( + torch.tensor(ddp_losses[step]), + torch.tensor(fsdp_losses[step]), + rtol=1e-5, + atol=1e-5, + msg=f"Loss mismatch at step {step}: DDP={ddp_losses[step]}, {test_label}={fsdp_losses[step]}", + ) + torch.testing.assert_close( + torch.tensor(ddp_grad_norms[step]), + torch.tensor(fsdp_grad_norms[step]), + rtol=1e-5, + atol=1e-5, + msg=f"Grad norm mismatch at step {step}: DDP={ddp_grad_norms[step]}, {test_label}={fsdp_grad_norms[step]}", + ) + + for key in ddp_state_dict: + assert key in fsdp_state_dict, f"Key {key} missing from {test_label} state dict" + torch.testing.assert_close( + ddp_state_dict[key], + fsdp_state_dict[key], + rtol=1e-5, + atol=1e-5, + msg=f"Weight mismatch for {key}: DDP vs {test_label}", + ) + + if rank == 0: + logger.debug(f"DDP and {test_label} comparison checks passed.") + + +# ============================================================================= +# Mixin class +# ============================================================================= + + +class FSDPTesterMixin(ABC): + fsdp_nproc_per_node: int = 2 + skip_fsdp_tests: bool = False + + @property + @abstractmethod + def model_tester(self): + """The model tester instance (e.g., CausalLMModelTester).""" + ... + + def _skip_if_fsdp_disabled(self): + if self.skip_fsdp_tests: + self.skipTest("FSDP tests disabled for this model (skip_fsdp_tests=True)") + + def _skip_if_insufficient_devices(self): + self._skip_if_fsdp_disabled() + available_workers = _get_available_fsdp_workers() + if available_workers < self.fsdp_nproc_per_node: + self.skipTest(f"Need at least {self.fsdp_nproc_per_node} FSDP workers, have {available_workers}") + + def _get_fsdp_model_name(self): + module_parts = self.__class__.__module__.split(".") + if len(module_parts) >= 3 and module_parts[0] == "tests" and module_parts[1] == "models": + return module_parts[2] + return None + + def _skip_if_fsdp_model_not_selected(self): + model_name = self._get_fsdp_model_name() + if model_name not in FSDP_TOP_MODEL_NAMES: + model_label = model_name or self.__class__.__module__ + self.skipTest( + "FSDP mixin coverage is currently limited to the top-10 dense and top-10 MoE model suites " + f"(skipping {model_label})." + ) + + def _create_model_on_meta(self, config): + """Instantiate a model on the meta device (no memory allocated).""" + auto_classes = [AutoModelForCausalLM, AutoModelForSeq2SeqLM] + for auto_cls in auto_classes: + try: + with torch.device("meta"): + return auto_cls.from_config(config) + except Exception: + continue + self.skipTest(f"Cannot instantiate model with any Auto class for config {type(config).__name__}") + + def _get_tiny_config(self): + """Get config class and serialized dict for passing to spawned processes.""" + config = self.model_tester.get_config() + config.vocab_size = 256 + config.hidden_size = 64 + config.intermediate_size = 128 + if hasattr(config, "ffn_config"): + # Keep nested FFN projections consistent with resized hidden size. + if hasattr(config.ffn_config, "ffn_hidden_size"): + config.ffn_config.ffn_hidden_size = config.hidden_size + if hasattr(config.ffn_config, "hidden_size"): + config.ffn_config.hidden_size = config.intermediate_size + if hasattr(config, "num_attention_heads"): + config.num_attention_heads = 4 + if hasattr(config, "num_key_value_heads"): + config.num_key_value_heads = 4 + if hasattr(config, "vocab_size_per_layer_input"): + config.vocab_size_per_layer_input = config.vocab_size + # `to_diff_dict()` avoids nested config pollution (e.g. DBRX ffn_config receiving + # generic PretrainedConfig keys that its constructor rejects). + config_dict = config.to_diff_dict() + return type(config), config_dict + + def _run_fsdp2_distributed_test(self, test_name, test_impl, *test_args, **test_kwargs): + self._skip_if_fsdp_model_not_selected() + self._skip_if_insufficient_devices() + + config_class, config_dict = self._get_tiny_config() + func_args = (config_class, config_dict, *test_args) + + results_file = tempfile.mktemp(suffix=".json") + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("", 0)) + port = s.getsockname()[1] + + try: + mp.spawn( + _fsdp_global_wrapper, + args=(test_name, test_impl, func_args, test_kwargs, self.fsdp_nproc_per_node, port, results_file), + nprocs=self.fsdp_nproc_per_node, + ) + + with open(results_file) as f: + result = json.load(f) + finally: + if os.path.exists(results_file): + os.unlink(results_file) + + if result["error"] is not None: + self.fail(f"FSDP test '{test_name}' failed:\n{result['error']}") + + # ========================================================================= + # Test: get_transformer_block_classes (CPU, meta device) + # ========================================================================= + + @is_fsdp_test + def test_get_transformer_block_classes(self): + """get_transformer_block_classes() finds >= 1 block class for the model.""" + self._skip_if_fsdp_disabled() + self._skip_if_fsdp_model_not_selected() + start_time = time.perf_counter() + logger.info("[FSDP] Starting test: test_get_transformer_block_classes") + status = "FAIL" + try: + config = self.model_tester.get_config() + model = self._create_model_on_meta(config) + + block_classes = get_transformer_block_classes(model) + self.assertTrue(len(block_classes) > 0, f"No block classes found for {type(config).__name__}") + + for cls in block_classes: + count = sum(1 for m in model.modules() if type(m) is cls) + self.assertGreater(count, 0, f"Block class {cls.__name__} has no instances in model") + status = "PASS" + finally: + logger.info( + "[FSDP] %s test: test_get_transformer_block_classes (%.1fs)", status, time.perf_counter() - start_time + ) + + @is_fsdp_test + @require_fsdp + def test_fsdp2_sharding_structure_untied(self): + self._run_fsdp2_distributed_test( + "test_fsdp2_sharding_structure_untied", _test_fsdp2_sharding_structure_impl, False + ) + + @is_fsdp_test + @require_fsdp + def test_fsdp2_sharding_structure_tied(self): + self._run_fsdp2_distributed_test( + "test_fsdp2_sharding_structure_tied", _test_fsdp2_sharding_structure_impl, True + ) + + @is_fsdp_test + @require_fsdp + def test_fsdp2_save_load(self): + self._run_fsdp2_distributed_test("test_fsdp2_save_load", _test_fsdp2_save_load_impl) + + @is_fsdp_test + @require_fsdp + def test_fsdp2_auto_plan_vs_ddp_untied(self): + self._run_fsdp2_distributed_test( + "test_fsdp2_auto_plan_vs_ddp_untied", _test_fsdp2_plan_vs_ddp_impl, False, "auto" + ) + + @is_fsdp_test + @require_fsdp + def test_fsdp2_auto_plan_vs_ddp_tied(self): + self._run_fsdp2_distributed_test( + "test_fsdp2_auto_plan_vs_ddp_tied", _test_fsdp2_plan_vs_ddp_impl, True, "auto" + ) + + @is_fsdp_test + @require_fsdp + def test_fsdp2_manual_plan_vs_ddp_untied(self): + self._run_fsdp2_distributed_test( + "test_fsdp2_manual_plan_vs_ddp_untied", _test_fsdp2_plan_vs_ddp_impl, False, "manual" + ) + + @is_fsdp_test + @require_fsdp + def test_fsdp2_manual_plan_vs_ddp_tied(self): + self._run_fsdp2_distributed_test( + "test_fsdp2_manual_plan_vs_ddp_tied", _test_fsdp2_plan_vs_ddp_impl, True, "manual" + ) diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py index 9dbf44c03c12..f9b3e48a55fd 100755 --- a/tests/test_modeling_common.py +++ b/tests/test_modeling_common.py @@ -2533,8 +2533,10 @@ def test_can_use_safetensors(self): torch.testing.assert_close( v, reloaded_state[k], - msg=lambda x: f"{model_class.__name__}: Tensor {k}: {x}.\n{v}\nvs\n{reloaded_state[k]}\n" - "This probably means that it was not set with the correct value when tying.", + msg=lambda x: ( + f"{model_class.__name__}: Tensor {k}: {x}.\n{v}\nvs\n{reloaded_state[k]}\n" + "This probably means that it was not set with the correct value when tying." + ), ) # Checking the tensor sharing are correct on the new model (weights are properly tied in both cases) @@ -2580,7 +2582,9 @@ def test_load_save_without_tied_weights(self): torch.testing.assert_close( v, reloaded_state[k], - msg=lambda x: f"{model_class.__name__}: Tensor {k}: {x}. Key {k} was serialized: {k in serialized_keys}. If `False`, this means it was probably aliased and safetensors removed it. If `True` it means `_init_weights` overwrote that key", + msg=lambda x: ( + f"{model_class.__name__}: Tensor {k}: {x}. Key {k} was serialized: {k in serialized_keys}. If `False`, this means it was probably aliased and safetensors removed it. If `True` it means `_init_weights` overwrote that key" + ), ) # Checking there was no complain of missing weights diff --git a/tests/test_tensor_parallel_mixin.py b/tests/test_tensor_parallel_mixin.py index 4e4e60159eec..6f142b513b41 100644 --- a/tests/test_tensor_parallel_mixin.py +++ b/tests/test_tensor_parallel_mixin.py @@ -146,7 +146,7 @@ def _verify_tp_sharding(rank, model_tp, model_ref): 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",): + if param_plan == "packed_colwise": expected_size = param_full.size(dim) // world_size assert param.size(dim) == expected_size, ( f"Packed weight {name} sharding incorrect: expected {expected_size}, got {param.size(dim)}" @@ -227,7 +227,7 @@ def _test_tp_backward_impl(rank, model_path, model_class, atol, rtol): 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",): + if param_plan == "packed_colwise": # interleaved slicing grad = get_packed_grad_shard(grad, world_size, rank, dim) else: From 739332cdcc3352a77d5e7daa61aa4c5ab15441f4 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Mon, 13 Apr 2026 14:14:20 +0000 Subject: [PATCH 003/116] DistributedConfig + shard-on-read loading - DtensorShardOperation for range-math shard-on-read - spawn_materialize() enhancements - from_pretrained wiring for distributed config - Shard operation helpers in tensor_parallel - Shard-on-read and LoadStateDictConfig tests --- src/transformers/core_model_loading.py | 342 +++++++++++++++--- .../integrations/tensor_parallel.py | 10 +- src/transformers/modeling_utils.py | 36 +- tests/utils/test_core_model_loading.py | 84 ++++- tests/utils/test_modeling_utils.py | 37 ++ 5 files changed, 427 insertions(+), 82 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index e0310c4abfeb..2b528bd0a829 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -32,7 +32,7 @@ import torch from .integrations.accelerate import get_device, offload_weight -from .integrations.tensor_parallel import ALL_PARALLEL_STYLES +from .integrations.tensor_parallel import ALL_PARALLEL_STYLES, get_tensor_shard from .utils import is_env_variable_true from .utils.loading_report import LoadStateDictInfo from .utils.logging import get_logger, tqdm @@ -41,9 +41,11 @@ _torch_distributed_available = torch.distributed.is_available() if TYPE_CHECKING: - from .integrations.tensor_parallel import TensorParallelLayer from .modeling_utils import LoadStateDictConfig, PreTrainedModel from .quantizers import HfQuantizer +elif _torch_distributed_available: + from torch.distributed.tensor import DTensor + from torch.distributed.tensor._utils import compute_local_shape_and_global_offset logger = get_logger(__name__) @@ -81,6 +83,21 @@ def build_glob_alternation( return alternation, src_group_to_glob, tgt_group_to_glob +def resolve_target_wildcards(source_pattern: str, target_pattern: str, source_key: str) -> str: + if "*" not in target_pattern or "*" not in source_pattern: + return target_pattern + + wildcard_regex = re.escape(source_pattern).replace(r"\*", r"(.*?)") + match = re.fullmatch(wildcard_regex, source_key) + if match is None: + return target_pattern + + resolved_target = target_pattern + for wildcard_value in match.groups(): + resolved_target = resolved_target.replace("*", wildcard_value, 1) + return resolved_target + + class ConversionOps: """Base class for weight conversion operations.""" @@ -316,7 +333,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) @@ -332,11 +349,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 @@ -519,7 +535,7 @@ class WeightTransform: target_patterns: str | list[str] = field(init=True) compiled_sources: re.Pattern = field(init=False) - distributed_operation: TensorParallelLayer | None = None + distributed_operation: Any | None = None quantization_operation: ConversionOps | None = None collected_tensors: dict[str, list[Future]] = field(default_factory=lambda: defaultdict(list), init=False) @@ -612,6 +628,7 @@ def rename_source_key(self, source_key: str) -> tuple[str, str | None]: source_pattern_that_matched = self.source_patterns[int(matching_group_name[1:])] # If we matched, we always replace with the first target pattern, in case we have several (one to many transform) replacement = self.target_patterns[0] + replacement = resolve_target_wildcards(source_pattern_that_matched, replacement, source_key) # Allow capturing groups in patterns, i.e. to add a prefix to all keys (e.g. timm_wrapper, sam3) if r"\1" in replacement: # The index of the internal group we need to replace is the index of the matched named group as it comes @@ -659,7 +676,7 @@ def materialize_tensors(self) -> dict[str, list[torch.Tensor]]: tensors = [future.result() for future in tensors if future.result() is not None] # Sync loading elif callable(tensors[0]): - tensors = [func() for func in tensors] + tensors = [tensor for func in tensors if (tensor := func()) is not None] # Add them to the new dictionary collected_tensors[key] = tensors @@ -766,18 +783,41 @@ def convert( pass if hf_quantizer is not None and self.quantization_operation is not None: - with log_conversion_errors( - layer_name, loading_info, (len(collected_tensors), layer_name), self.quantization_operation - ): - collected_tensors = self.quantization_operation.convert( - collected_tensors, - source_patterns=self.source_patterns, - target_patterns=self.target_patterns, - full_layer_name=layer_name, - config=config, - model=model, - missing_keys=loading_info.missing_keys if loading_info else None, - ) + if len(collected_tensors) > 1 and model is not None: + quantized_tensors = {} + for target_key, tensor in collected_tensors.items(): + if not hf_quantizer.param_needs_quantization(model, target_key): + quantized_tensors[target_key] = tensor + continue + quantize_input = tensor if isinstance(tensor, list) else [tensor] + with log_conversion_errors( + target_key, loading_info, (len(quantize_input), target_key), self.quantization_operation + ): + quantized_tensors.update( + self.quantization_operation.convert( + {target_key: quantize_input}, + source_patterns=self.source_patterns, + target_patterns=[target_key], + full_layer_name=target_key, + config=config, + model=model, + missing_keys=loading_info.missing_keys if loading_info else None, + ) + ) + collected_tensors = quantized_tensors + else: + with log_conversion_errors( + layer_name, loading_info, (len(collected_tensors), layer_name), self.quantization_operation + ): + collected_tensors = self.quantization_operation.convert( + collected_tensors, + source_patterns=self.source_patterns, + target_patterns=self.target_patterns, + full_layer_name=layer_name, + config=config, + model=model, + missing_keys=loading_info.missing_keys if loading_info else None, + ) return collected_tensors @@ -815,10 +855,15 @@ def _job(): return _job -def spawn_tp_materialize( - thread_pool: ThreadPoolExecutor | None, tensor: torch.Tensor, sharding_method, tensor_idx, device=None, dtype=None +def spawn_parallel_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 + """Materialize and shard a tensor according to the active parallelism strategy if `thread_pool` is provided, or return a Callable that will load the tensor synchronously when called.""" def _job(): @@ -832,6 +877,133 @@ def _job(): return _job +@dataclass(slots=True) +class ParallelMaterializationContext: + distributed_operation: Any + tensor_idx: int | None + device: Any + + +def is_dtensor_like(value: Any) -> bool: + return all(hasattr(value, attr) for attr in ("device_mesh", "placements", "to_local")) + + +@dataclass(slots=True) +class FSDPShardOperation: + device_mesh: Any + rank: int + empty_param: Any + placements: tuple[Any, ...] + shard_placement: Any | None = field(init=False, default=None) + local_shape: tuple[int, ...] = field(init=False) + + def __post_init__(self): + shard_placements = [placement for placement in self.placements if placement.is_shard()] + if len(shard_placements) > 1: + raise NotImplementedError( + f"FSDP shard-on-read does not support multiple shard placements yet: {self.placements}" + ) + self.shard_placement = shard_placements[0] if shard_placements else None + if self.shard_placement is not None and len(self.placements) != 1: + raise NotImplementedError( + f"FSDP shard-on-read only supports a single placement today. Got placements={self.placements}." + ) + self.local_shape = self.get_expected_sharded_shape(self.empty_param.shape) + + @classmethod + def from_param(cls, param: Any) -> FSDPShardOperation: + return cls( + device_mesh=param.device_mesh, + rank=param.device_mesh.get_local_rank(), + empty_param=param, + placements=tuple(param.placements), + ) + + def shard_tensor( + self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None + ) -> torch.Tensor | None: + if self.shard_placement is None: + local_tensor = param[...] + else: + param_shape = list(param.shape) if isinstance(param, torch.Tensor) else param.get_shape() + # Mixtral-style converted expert weights first stack individual expert tensors along dim 0 before + # concatenating. Only materialize the experts owned by this rank. + if ( + tensor_idx is not None + and len(self.empty_param.shape) == len(param_shape) + 1 + and self.shard_placement.dim == 0 + ): + local_expert_count = self.local_shape[0] + expert_offset = compute_local_shape_and_global_offset( + self.empty_param.shape, self.device_mesh, self.placements + )[1][0] + if tensor_idx < expert_offset or tensor_idx >= expert_offset + local_expert_count: + return None + local_tensor = param[...] + else: + local_tensor = get_tensor_shard( + param, + self.empty_param, + self.device_mesh, + self.rank, + self.shard_placement.dim, + tensor_idx=tensor_idx, + ) + if local_tensor is None: + return None + return local_tensor.to(device=device, dtype=dtype) + + def get_expected_sharded_shape(self, full_shape: tuple[int, ...] | torch.Size) -> tuple[int, ...]: + local_shape, _ = compute_local_shape_and_global_offset(full_shape, self.device_mesh, self.placements) + return tuple(local_shape) + + def update_module_attributes(self, module: torch.nn.Module): + return None + + +def get_parallel_materialization_context( + mapping: WeightTransform, + renamed_key: str, + source_pattern: str, + empty_param: Any, + device_mesh: Any, + parallel_plan: dict[str, Any], + parallel_pattern_matcher: re.Pattern | None, + parallel_pattern_by_group_name: dict[str, str] | None, + device_map: dict[str, Any], +) -> ParallelMaterializationContext | None: + tensor_idx = ( + len(mapping.collected_tensors.get(source_pattern, [])) + if isinstance(mapping, WeightConverter) and isinstance(mapping.operations[0], MergeModulelist) + else None + ) + + if ( + device_mesh + and parallel_plan + and parallel_pattern_matcher is not None + and parallel_pattern_by_group_name is not None + ): + if matched_parallel_pattern := parallel_pattern_matcher.search(renamed_key): + matched_parallel_pattern = parallel_pattern_by_group_name[matched_parallel_pattern.lastgroup] + if getattr(mapping, "distributed_operation", None) is None: + parallel_layer = ALL_PARALLEL_STYLES[parallel_plan[matched_parallel_pattern]].__class__ + mapping.distributed_operation = parallel_layer( + device_mesh=device_mesh, rank=device_mesh.get_local_rank(), empty_param=empty_param.clone() + ) + return ParallelMaterializationContext(mapping.distributed_operation, tensor_idx, device_map[""]) + + if is_dtensor_like(empty_param): + if getattr(mapping, "distributed_operation", None) is None: + mapping.distributed_operation = FSDPShardOperation.from_param(empty_param) + return ParallelMaterializationContext( + mapping.distributed_operation, + tensor_idx, + get_device(device_map, renamed_key, valid_torch_device=True), + ) + + return None + def dot_natural_key(s: str): """Sort key for state-dict names: split on ``"."`` and sort digits numerically and strings alphabetically. We emit a tuple at each point to sort ints @@ -900,7 +1072,7 @@ def set_param_for_module( target_name: str, param_value: torch.Tensor, loading_info: LoadStateDictInfo, - distributed_operation: TensorParallelLayer | None, + distributed_operation: Any | None, hf_quantizer: HfQuantizer, ): module_path, _, param_name = target_name.rpartition(".") @@ -915,15 +1087,18 @@ 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 is_dtensor_like(ref): 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: + # Determine expected shape: for TP/FSDP shard-on-read, use the local shard shape; otherwise, use full shape + if is_dtensor_like(ref): + local_shape, _ = compute_local_shape_and_global_offset(ref.shape, ref.device_mesh, ref.placements) + expected_shape = torch.Size(local_shape) + elif distributed_operation is not None: expected_shape = torch.Size(distributed_operation.get_expected_sharded_shape(ref.shape)) else: expected_shape = ref.shape @@ -931,11 +1106,29 @@ def set_param_for_module( if ref is not None and param_value.shape != expected_shape and hf_quantizer is None: loading_info.mismatched_keys.add((target_name, param_value.shape, expected_shape)) else: - # 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) + if is_dtensor_like(ref): + local_param = param_value.detach() if isinstance(param_value, torch.nn.Parameter) else param_value + fsdp_param = DTensor.from_local( + local_param.contiguous(), + ref.device_mesh, + ref.placements, + run_check=False, + shape=ref.shape, + stride=tuple(ref.stride()), + ) + with torch.no_grad(): + if ref.is_meta: + fsdp_param = torch.nn.Parameter(fsdp_param, requires_grad=ref.requires_grad) + torch.utils.swap_tensors(ref, fsdp_param) + else: + ref.copy_(fsdp_param) + ref._is_hf_initialized = True + else: + # 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( @@ -1002,6 +1195,30 @@ def rename_source_key( return renamed_key, source_pattern +def concretize_target_patterns( + converter: WeightConverter, + source_key: str, + source_pattern: str, + prefix: str | None, + meta_state_dict: dict | None, +) -> WeightConverter: + concrete_targets = [] + for target_pattern in converter.target_patterns: + concrete_target = resolve_target_wildcards(source_pattern, target_pattern, source_key) + if prefix is not None and meta_state_dict is not None: + if ( + concrete_target.startswith(prefix) + and meta_state_dict.get(re.sub(f"^{prefix}.", "", concrete_target, count=1)) is not None + ): + concrete_target = re.sub(f"^{prefix}.", "", concrete_target, count=1) + elif meta_state_dict.get(f"{prefix}.{concrete_target}") is not None: + concrete_target = f"{prefix}.{concrete_target}" + concrete_targets.append(concrete_target) + + object.__setattr__(converter, "target_patterns", concrete_targets) + return converter + + def convert_and_load_state_dict_in_model( model: PreTrainedModel, state_dict: dict[str, Any], @@ -1156,10 +1373,25 @@ def convert_and_load_state_dict_in_model( # 2. finally, collect the tensor into the proper converter if renamed_key in meta_model_state_dict: - empty_param = meta_model_state_dict.get(renamed_key) + empty_param = meta_model_state_dict[renamed_key] + try: + empty_param = model.get_parameter_or_buffer(renamed_key) + except (AttributeError, KeyError): + if getattr(model, "_is_fsdp_managed_module", False): + raise RuntimeError( + f"FSDP shard-on-read requires the live parameter for {renamed_key!r}, " + f"but get_parameter_or_buffer() failed." + ) # If we enter here, we have a WeightConverter operation to perform if source_pattern is not None: new_converter = deepcopy(pattern_to_converter[source_pattern]) + new_converter = concretize_target_patterns( + new_converter, + original_key, + source_pattern, + prefix, + meta_model_state_dict, + ) # each target key gets its own converter instance mapping = param_name_to_load.setdefault(renamed_key, new_converter) # Otherwise, only potential renaming @@ -1200,29 +1432,27 @@ 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 + # 4. Handle parallel shard-on-read or device_map placement future_or_tensor = None - if device_mesh and tp_plan: - if matched_tp_pattern := tp_plan_alt.search(renamed_key): - matched_tp_pattern = tp_plan_by_group_name[matched_tp_pattern.lastgroup] - if getattr(mapping, "distributed_operation", None) is None: - tp_layer = ALL_PARALLEL_STYLES[model.tp_plan[matched_tp_pattern]].__class__ - mapping.distributed_operation = tp_layer( - device_mesh=device_mesh, rank=device_mesh.get_local_rank(), empty_param=empty_param.clone() - ) - shard_index = ( - len(mapping.collected_tensors.get(source_pattern, [])) - if isinstance(mapping, WeightConverter) and isinstance(mapping.operations[0], MergeModulelist) - else None - ) - future_or_tensor = spawn_tp_materialize( - thread_pool, - tensor, - mapping.distributed_operation, - shard_index, - device_map[""], - _dtype, - ) + if parallel_context := get_parallel_materialization_context( + mapping=mapping, + renamed_key=renamed_key, + source_pattern=source_pattern, + empty_param=empty_param, + device_mesh=device_mesh, + parallel_plan=tp_plan, + parallel_pattern_matcher=tp_plan_alt if tp_plan else None, + parallel_pattern_by_group_name=tp_plan_by_group_name if tp_plan else None, + device_map=device_map, + ): + future_or_tensor = spawn_parallel_materialize( + thread_pool, + tensor, + parallel_context.distributed_operation, + parallel_context.tensor_idx, + parallel_context.device, + _dtype, + ) if future_or_tensor is None: param_device = get_device(device_map, renamed_key, valid_torch_device=True) diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index 39a2e696941b..c378ebbac227 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -1506,8 +1506,8 @@ 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.""" +def distribute_model(model, tp_plan, distributed_config, device_mesh, tp_size, fsdp_plan=None): + """Attach distributed runtime hooks before checkpoint loading.""" model._tp_size = tp_size model._device_mesh = device_mesh if distributed_config is not None: @@ -1517,7 +1517,7 @@ def distribute_model(model, tp_plan, distributed_config, device_mesh, tp_size): # Set the new requested tp_plan on the model if isinstance(tp_plan, dict): model.tp_plan = tp_plan - model_plan = model.tp_plan + model_plan = model.tp_plan if tp_plan is not None or tp_size is not None else None if model_plan is not None and _torch_distributed_available: for v in model_plan.values(): if v not in ALL_PARALLEL_STYLES: @@ -1533,4 +1533,8 @@ def distribute_model(model, tp_plan, distributed_config, device_mesh, tp_size): device_mesh, ) module._is_hooked = True + if fsdp_plan is not None: + from .fsdp import apply_fsdp2 + + model = apply_fsdp2(model, device_mesh, fsdp_plan) return model diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 2a7304650640..b6feec6f7ba6 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -69,7 +69,7 @@ from .integrations.flash_attention import flash_attention_forward from .integrations.flash_paged import paged_attention_forward from .integrations.flex_attention import flex_attention_forward -from .integrations.fsdp import apply_fsdp2 +from .integrations.fsdp import initialize_fsdp from .integrations.hub_kernels import allow_all_hub_kernels, is_kernel from .integrations.peft import maybe_load_adapters from .integrations.sdpa_attention import sdpa_attention_forward @@ -178,6 +178,7 @@ class LoadStateDictConfig: dtype_plan: dict = field(default_factory=dict) hf_quantizer: HfQuantizer | None = None device_mesh: Optional["torch.distributed.device_mesh.DeviceMesh"] = None + tp_plan: dict[str, str] | None = None weights_only: bool = True weight_mapping: list[WeightConverter | WeightRenaming] | None = None @@ -4003,11 +4004,21 @@ def from_pretrained( ": PartialState().process_index} where PartialState comes from accelerate library" ) + if fsdp_plan is not None and (tp_plan is not None or tp_size is not None): + raise ValueError("Combining `fsdp_plan` with tensor parallel loading is not supported yet.") + 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 fsdp_plan is not None: + device_map, device_mesh, _ = initialize_fsdp( + fsdp_plan=fsdp_plan, + device_mesh=device_mesh, + device_map=device_map, + ) + if gguf_file is not None and not is_accelerate_available(): raise ValueError("accelerate is required when loading a GGUF file `pip install accelerate`.") @@ -4132,14 +4143,17 @@ 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) + if _torch_distributed_available and device_mesh is not None and (tp_plan is not None or fsdp_plan is not None): + model = distribute_model(model, tp_plan, distributed_config, device_mesh, tp_size, fsdp_plan=fsdp_plan) # Prepare the full device map - if device_map is not None: + if isinstance(device_map, dict): device_map = _get_device_map(model, device_map, max_memory, hf_quantizer) + elif device_map is not None: + device_map = {"": device_map} # Finalize model weight initialization + active_tp_plan = getattr(model, "_tp_plan", None) if tp_size is not None else None load_config = LoadStateDictConfig( pretrained_model_name_or_path=pretrained_model_name_or_path, ignore_mismatched_sizes=ignore_mismatched_sizes, @@ -4151,6 +4165,7 @@ def from_pretrained( dtype_plan=dtype_plan, hf_quantizer=hf_quantizer, device_mesh=device_mesh, + tp_plan=active_tp_plan, weights_only=weights_only, weight_mapping=weight_conversions, use_safetensors=use_safetensors, @@ -4161,15 +4176,6 @@ def from_pretrained( model.eval() # Set model in evaluation mode to deactivate Dropout modules by default model.set_use_kernels(use_kernels, kernel_config) - # Apply FSDP2 if configured (must be after weight loading) - if fsdp_plan is not None: - if device_mesh is None: - raise ValueError( - "`fsdp_plan` was provided but no device mesh is available. " - "Pass `device_mesh` to `from_pretrained`." - ) - model = apply_fsdp2(model, device_mesh, fsdp_plan) - # If it is a model with generation capabilities, attempt to load generation files (generation config, # custom generate function) if model.can_generate() and hasattr(model, "adjust_generation_fn") and not gguf_file: @@ -4226,7 +4232,7 @@ def _load_pretrained_model( expected_keys = list(model.state_dict().keys()) if expected_keys is None else expected_keys if logger.level >= logging.WARNING: - verify_tp_plan(expected_keys, getattr(model, "_tp_plan", None)) + verify_tp_plan(expected_keys, load_config.tp_plan) # This offload index if for params explicitly on the "disk" in the device_map disk_offload_index = None @@ -4289,7 +4295,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, ) diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py index 3e8c18b1d351..942dcdc99b11 100644 --- a/tests/utils/test_core_model_loading.py +++ b/tests/utils/test_core_model_loading.py @@ -23,6 +23,7 @@ Chunk, Concatenate, ErnieFuseAndSplitTextVisionExperts, + FSDPShardOperation, MergeModulelist, PermuteForRope, WeightConverter, @@ -31,6 +32,7 @@ convert_and_load_state_dict_in_model, rename_source_key, revert_weight_conversion, + spawn_parallel_materialize, ) from transformers.modeling_utils import LoadStateDictConfig from transformers.utils.import_utils import is_triton_available @@ -214,7 +216,66 @@ def __init__(self, add_extra_moe=False): self.mlp = DummyMLP() +class FakeMesh: + def __init__(self, world_size: int, rank: int): + self.shape = (world_size,) + self._rank = rank + + def get_local_rank(self): + return self._rank + + def get_coordinate(self): + return (self._rank,) + + class TestConvertAndLoadStateDict(unittest.TestCase): + def test_fsdp_shard_aware_mixtral_conversion_uses_only_local_experts(self): + shard_op = FSDPShardOperation( + device_mesh=FakeMesh(world_size=2, rank=0), + rank=0, + empty_param=torch.empty((2, 4, 2)), + placements=(torch.distributed.tensor.placement_types.Shard(0),), + ) + 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_parallel_materialize(None, tensor, shard_op, idx, device="cpu", dtype=None), + ) + + 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_parallel_materialize(None, tensor, shard_op, idx, device="cpu", dtype=None), + ) + + 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() model.config = PretrainedConfig() @@ -467,6 +528,10 @@ def __init__(self): self, "quantization_config", SimpleNamespace(weight_block_size=bs) ), "param_needs_quantization": lambda self, _model, param_name: param_name.endswith("q_proj.weight"), + "get_quantize_ops": lambda self: __import__( + "transformers.integrations.finegrained_fp8", + fromlist=["Fp8Quantize"], + ).Fp8Quantize(self), "pre_quantized": False, }, ) @@ -499,11 +564,11 @@ def __init__(self): 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) @@ -514,11 +579,14 @@ def __init__(self): 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_ernie4_5_vl_moe_conversion(self): diff --git a/tests/utils/test_modeling_utils.py b/tests/utils/test_modeling_utils.py index 7366845c4d78..59177fec5061 100644 --- a/tests/utils/test_modeling_utils.py +++ b/tests/utils/test_modeling_utils.py @@ -431,6 +431,43 @@ 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 = [] + + def fake_distribute_model(model, tp_plan, distributed_config, device_mesh, tp_size, fsdp_plan=None): + call_order.append("distribute") + self.assertEqual(fsdp_plan, {"mode": "auto"}) + model._tp_plan = {"model.layers.*.mlp.experts.gate_up_proj": "packed_colwise"} + model._is_fsdp_managed_module = True + return model + + def fake_load_pretrained_model(model, state_dict, checkpoint_files, load_config, expected_keys=None): + call_order.append("load") + self.assertEqual(load_config.device_mesh, "fake-mesh") + self.assertEqual(load_config.device_map, {"": torch.device("cpu")}) + self.assertIsNone(load_config.tp_plan) + return mock.Mock(), None + + with ( + patch( + "transformers.modeling_utils.initialize_fsdp", return_value=(torch.device("cpu"), "fake-mesh", 2) + ), + patch("transformers.modeling_utils.distribute_model", side_effect=fake_distribute_model), + 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, fsdp_plan={"mode": "auto"}) + + self.assertEqual(call_order, ["distribute", "load"]) + def test_hub_retry(self): @hub_retry(max_attempts=2) def test_func(): From 11b55a200941fb053819fd62e1f8b54ec337ee71 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Mon, 13 Apr 2026 14:19:34 +0000 Subject: [PATCH 004/116] TPStyle API + dense model tensor parallelism - Replace hook-based TP with DTensor-based TPStyle API - TPStyle dataclass with dense kinds: colwise, rowwise, vocab - apply_tensor_parallel() using PyTorch parallelize_module - verify_tp_plan() for plan validation - Update dense model configs (llama, mistral, qwen2, phi, glm) to TPStyle - DTensor apply_rotary_pos_emb guard for llama, mistral, qwen3 - Extended DistributedConfig with tp/fsdp size and plan fields - DistributedConfig serialization in configuration_utils - MXFP4 NotImplementedError for DTensor TP - Dense TP tests --- src/transformers/configuration_utils.py | 3 + .../distributed/configuration_utils.py | 118 +- src/transformers/integrations/mxfp4.py | 48 +- .../integrations/tensor_parallel.py | 1552 ++--------------- .../models/glm/configuration_glm.py | 13 +- .../models/llama/configuration_llama.py | 15 +- .../models/llama/modeling_llama.py | 5 + .../models/mistral/configuration_mistral.py | 15 +- .../models/mistral/modeling_mistral.py | 5 + .../models/phi/configuration_phi.py | 17 +- .../models/qwen2/configuration_qwen2.py | 15 +- .../models/qwen3/modeling_qwen3.py | 7 +- tests/tensor_parallel/test_tensor_parallel.py | 32 - tests/test_distributed_config.py | 86 + tests/test_tensor_parallel_mixin.py | 131 +- 15 files changed, 394 insertions(+), 1668 deletions(-) create mode 100644 tests/test_distributed_config.py diff --git a/src/transformers/configuration_utils.py b/src/transformers/configuration_utils.py index 4f58a230e352..2f993e87d4a4 100755 --- a/src/transformers/configuration_utils.py +++ b/src/transformers/configuration_utils.py @@ -1008,6 +1008,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] diff --git a/src/transformers/distributed/configuration_utils.py b/src/transformers/distributed/configuration_utils.py index 7726d9f3290d..e40aed267bcd 100644 --- a/src/transformers/distributed/configuration_utils.py +++ b/src/transformers/distributed/configuration_utils.py @@ -12,99 +12,59 @@ # 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 +import torch +from dataclasses import asdict, dataclass @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 (`str` or `dict`, *optional*): + Tensor parallel sharding plan. Use `"auto"` for the model's default plan. + fsdp_size (`int`, *optional*): + Number of devices for FSDP (data parallelism). If `None` and `tp_size` is set, defaults to 1. + fsdp_plan (`str` or `dict`, *optional*): + FSDP wrapping plan. Use `"auto"` to wrap each transformer layer + root. """ - enable_expert_parallel: bool = False - # TODO: add tp_plan, pp_plan, device_mesh etc.. + tp_size: int = 1 + tp_plan: str | dict[str, str] | None = None + enable_sequence_parallel: bool = False + fsdp_size: int = 1 + fsdp_plan: str | dict | None = None - @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 + def __post_init__(self): + # If a size is set without a plan, default the plan to "auto" + if self.tp_size > 1 and self.tp_plan is None: + self.tp_plan = "auto" + if self.fsdp_size > 1 and self.fsdp_plan is None: + self.fsdp_plan = "auto" + + world_size = torch.distributed.get_world_size() + assert self.tp_size * self.fsdp_size == world_size, f"tp_size ({self.tp_size}) * fsdp_size ({self.fsdp_size}) must be equal to world_size ({world_size})" - # Copied from transformers.utils.quantization_config.QuantizationConfigMixin.to_json_file - 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" + @classmethod + 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}) - writer.write(json_string) + def to_dict(self) -> dict: + return asdict(self) - 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__) + def to_json_string(self) -> str: + return json.dumps(self.to_dict(), indent=2) + "\n" - # 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() + def to_json_file(self, json_file_path: str | os.PathLike): + with open(json_file_path, "w", encoding="utf-8") as f: + f.write(self.to_json_string()) - # Copied from transformers.utils.quantization_config.QuantizationConfigMixin.__repr__ 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/integrations/mxfp4.py b/src/transformers/integrations/mxfp4.py index 67d9420659af..d1e5506cbf20 100644 --- a/src/transformers/integrations/mxfp4.py +++ b/src/transformers/integrations/mxfp4.py @@ -20,7 +20,7 @@ from torch import nn from contextlib import contextmanager -from ..core_model_loading import ConversionOps, _IdentityOp +from ..core_model_loading import ConversionOps from ..quantizers.quantizers_utils import get_module_from_name, should_convert_module @@ -145,10 +145,6 @@ def convert( dequantized = dequantize_convertops(param_data[f"{proj}_blocks"], param_data[f"{proj}_scales"]) return {full_layer_name: dequantized} - @property - def reverse_op(self) -> "ConversionOps": - return _IdentityOp() - class Mxfp4Deserialize(ConversionOps): def __init__(self, hf_quantizer): @@ -511,28 +507,15 @@ def mlp_forward(self, hidden_states): def dequantize(module, param_name, param_value, target_device, dq_param_name, **kwargs): - from ..integrations.tensor_parallel import shard_and_distribute_module - - model = kwargs.get("model") - empty_param = kwargs.get("empty_param") - casting_dtype = kwargs.get("casting_dtype") - to_contiguous = kwargs.get("to_contiguous") - rank = kwargs.get("rank") device_mesh = kwargs.get("device_mesh") + if device_mesh is not None: + raise NotImplementedError( + "MXFP4 quantization is not yet compatible with DTensor-based tensor parallelism. " + "Please disable TP or use a non-quantized model." + ) for proj in ["gate_up_proj", "down_proj"]: if proj in param_name: - if device_mesh is not None: - param_value = shard_and_distribute_module( - model, - param_value, - empty_param, - dq_param_name, - casting_dtype, - to_contiguous, - rank, - device_mesh, - ) blocks_attr = f"{proj}_blocks" scales_attr = f"{proj}_scales" setattr(module, param_name.rsplit(".", 1)[1], param_value) @@ -557,24 +540,19 @@ def load_and_swizzle_mxfp4(module, param_name, param_value, target_device, trito triton_kernels_hub.matmul_ogs.FlexCtx, triton_kernels_hub.matmul_ogs.InFlexData, ) - from ..integrations.tensor_parallel import shard_and_distribute_module - model = kwargs.get("model") - empty_param = kwargs.get("empty_param") - casting_dtype = kwargs.get("casting_dtype") - to_contiguous = kwargs.get("to_contiguous") - rank = kwargs.get("rank") device_mesh = kwargs.get("device_mesh") + if device_mesh is not None: + raise NotImplementedError( + "MXFP4 quantization is not yet compatible with DTensor-based tensor parallelism. " + "Please disable TP or use a non-quantized model." + ) + if "blocks" in param_name: proj = param_name.split(".")[-1].split("_blocks")[0] if "scales" in param_name: proj = param_name.split(".")[-1].split("_scales")[0] - if device_mesh is not None: - shard_and_distribute_module( - model, param_value, empty_param, param_name, casting_dtype, to_contiguous, rank, device_mesh - ) - else: - setattr(module, param_name.rsplit(".", 1)[1], torch.nn.Parameter(param_value, requires_grad=False)) + setattr(module, param_name.rsplit(".", 1)[1], torch.nn.Parameter(param_value, requires_grad=False)) blocks_attr = f"{proj}_blocks" scales_attr = f"{proj}_scales" blocks = getattr(module, blocks_attr) # at this point values were loaded from ckpt diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index c378ebbac227..41f4414c8d91 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -13,22 +13,24 @@ # limitations under the License. from __future__ import annotations -import math -import operator -import os import re -from functools import reduce - -from ..distributed import DistributedConfig -from ..utils import is_torch_greater_or_equal, logging -from ..utils.generic import GeneralInterface +from dataclasses import dataclass +from typing import Literal + +from torch.distributed.tensor import DTensor, Replicate, Shard +from torch.distributed.tensor.parallel import ( + ColwiseParallel, + RowwiseParallel, + parallelize_module, +) +from torch.distributed.tensor.parallel.style import ParallelStyle + +from ..utils import logging from ..utils.import_utils import is_torch_available if is_torch_available(): import torch - import torch.distributed as dist - from torch import nn # Cache this result has it's a C FFI call which can be pretty time-consuming _torch_distributed_available = torch.distributed.is_available() @@ -37,71 +39,6 @@ logger = logging.get_logger(__name__) -def initialize_tensor_parallelism( - tp_plan: str | dict[str, str] | None, tp_size: int | None = None, device_mesh=None, device_map=None -): - r""" - Sets up the device mesh and initialized the backend for tensor parallelism. - This function is called when the model is loaded and the TP plan is set to 'auto'. - """ - if tp_size is not None and tp_plan is None: - raise ValueError("tp_plan has to be set when tp_size is passed.") - if tp_plan is not None and device_map is not None: - raise ValueError("`tp_plan` and `device_map` are mutually exclusive. Choose either one for parallelization.") - if device_mesh is None: - if not is_torch_greater_or_equal("2.5"): - raise OSError("Tensor parallel is only supported for `torch>=2.5`.") - - # Detect the accelerator on the machine. If no accelerator is available, it returns CPU. - device_type = torch._C._get_accelerator().type - if device_type == "mps": - raise RuntimeError("Tensor parallelism is not supported on MPS devices.") - current_device = getattr(torch, device_type) - if not torch.distributed.is_initialized(): - try: - rank = int(os.environ["RANK"]) - local_rank = int(os.environ["LOCAL_RANK"]) - world_size = int(os.environ["WORLD_SIZE"]) - - backend_map = {"cuda": "nccl", "cpu": "gloo", "xpu": "xccl", "hpu": "hccl", "neuron": "neuron"} - backend = backend_map.get(device_type) - - torch.distributed.init_process_group(backend=backend, rank=rank, world_size=world_size) - current_device = getattr(torch, device_type) - if device_type != "cpu": - current_device.set_device(local_rank) - - except Exception as e: - raise OSError( - "We tried to initialize torch.distributed for you, but it failed. Make " - "sure you init torch distributed in your script to use `tp_plan`." - ) from e - - if device_type != "cpu": - current_device.set_device(int(os.environ["LOCAL_RANK"])) - index = current_device.current_device() - tp_device = torch.device(device_type, index) - device_map = tp_device - else: - tp_device = torch.device(device_type) - device_map = device_type or {} - - tp_size = tp_size if tp_size is not None else torch.distributed.get_world_size() - device_mesh = torch.distributed.init_device_mesh(tp_device.type, (tp_size,)) - else: - if device_mesh.ndim > 1: - if "tp" not in device_mesh.mesh_dim_names: - raise ValueError( - "When using `tp_plan` and n-d `device_mesh`, it must contain a 'tp' dimension. " - "Please provide a valid `device_mesh`." - ) - device_mesh = device_mesh["tp"] - tp_size = device_mesh.size() - device_map = torch.device(f"{device_mesh.device_type}:{int(os.environ['LOCAL_RANK'])}") - - return device_map, device_mesh, tp_size - - def replace_layer_number_by_wildcard(name: str) -> str: """ Replace the numbers in the `name` by wildcards, only if they are in-between dots (`.`) or if they are between @@ -135,1150 +72,16 @@ def _get_parameter_tp_plan(parameter_name: str, tp_plan: dict[str, str], is_weig # ============================================================================= -if is_torch_available(): - str_to_dtype = { - "BOOL": torch.bool, - "U8": torch.uint8, - "I8": torch.int8, - "I16": torch.int16, - "F16": torch.float16, - "BF16": torch.bfloat16, - "I32": torch.int32, - "F32": torch.float32, - "F64": torch.float64, - "I64": torch.int64, - "F8_E4M3": torch.float8_e4m3fn, - } - - -def _blocks_to_block_sizes(total_size: int, blocks: int | list[int]) -> list[int]: - """ - Convert block count or proportions to block sizes. - - This function accepts - - - The number of blocks (int), in which case the block size is - total_size//blocks; or - - A list of block sizes (list[int]). - - In the second case, if sum(blocks) < total_size, the ratios between - the block sizes will be preserved. For instance, if blocks is - [2, 1, 1] and total_size is 1024, the returned block sizes are - [512, 256, 256]. - """ - if isinstance(blocks, list): - total_blocks = sum(blocks) - assert total_size % total_blocks == 0, f"Cannot split {total_size} in proportional blocks: {blocks}" - part_size = total_size // total_blocks - return [part_size * block for block in blocks] - else: - assert total_size % blocks == 0, f"Prepacked is not divisible by {blocks}" - single_size = total_size // blocks - return [single_size] * blocks - - -def get_packed_weights(param, empty_param, device_mesh, rank, dim): - """ - When weights are packed (gate_up_proj), we need to make sure each shard gets its correct share. - So if you have: gate_proj ( 16, 5120, 8190) - and up_proj ( 16, 5120, 8190) - packed as gate_up_proj ( 16, 5120, 2 * 8190) - And you shard along the last dimension, you need to interleave the gate and up values: - - Now, if we shard along the last dimension across TP_size (Tensor Parallelism size), we must interleave the values from gate and up projections correctly. - - Let's take TP_size = 4 for an example: - - Packed tensor `gate_up_proj` - --------------------------------------------------------------- - [ G0 G1 G2 G3 | G4 G5 G6 G7 | ... | U0 U1 U2 U3 | U4 U5 U6 U7 | ... ] - ↑─────────────↑ ↑─────────────↑ ↑─────────────↑ ↑─────────────↑ - Gate Slice 0 Gate Slice 1 Up Slice 0 Up Slice 1 - - Explanation: - - The first half of the tensor (left of the center) holds the gate_proj values. - - The second half (right of the center) holds the up_proj values. - - For TP=4, we divide each half into 4 slices. In this example, we show two slices for brevity. - - Each shard receives one slice from the gate part and the corresponding slice from the up part. - - For instance: - • Shard 0 gets: [ Gate Slice 0, Up Slice 0 ] = [ G0, G1, G2, G3, U0, U1, U2, U3 ] - • Shard 1 gets: [ Gate Slice 1, Up Slice 1 ] = [ G4, G5, G6, G7, U4, U5, U6, U7 ] - • … and so on. - - This ensures that each shard receives an equal portion of both gate and up projections, maintaining consistency across tensor parallelism. - """ - slice_ = param - total_size = empty_param.shape[dim] - world_size = device_mesh.size() - block_sizes = _blocks_to_block_sizes(total_size=total_size, blocks=2) - - tensors_slices = [] - block_offset = 0 - for block_size in block_sizes: - shard_block_size = block_size // world_size - start = rank * shard_block_size - stop = (rank + 1) * shard_block_size - tensors_slices += range(block_offset + start, block_offset + stop) - block_offset += block_size - - slice_dtype = slice_.get_dtype() - # Handle F8_E4M3 dtype by converting to float16 before slicing - # Without upcasting, the slicing causes : RuntimeError: "index_cpu" not implemented for 'Float8_e4m3fn' - casted = False - if slice_dtype == "F8_E4M3" or slice_dtype == "F8_E5M2": - slice_ = slice_[...].to(torch.float16) - casted = True - - if dim == 0: - tensor = slice_[tensors_slices, ...] - elif dim == 1 or dim == -2: - tensor = slice_[:, tensors_slices, ...] - elif dim == 2 or dim == -1: - tensor = slice_[..., tensors_slices] - else: - raise ValueError(f"Unsupported dim {dim}, only dim 0, 1 or 2 are supported") - - if casted: +def _to_cpu_fresh(tensor: torch.Tensor) -> torch.Tensor: + """Plain tensor → contiguous CPU tensor with fresh storage for safetensors.""" + if tensor.device.type == "meta": return tensor - else: - return tensor.to(str_to_dtype[slice_dtype]) - - -def repack_weights( - packed_parameter: torch.Tensor, - sharded_dim: int, # The dimension index in the global tensor that was sharded - world_size: int, - num_blocks: int = 2, -) -> torch.Tensor: - """ - Reorders a tensor that was reconstructed from sharded packed weights into its canonical packed format. - - For example, if a weight was packed (e.g., gate_proj and up_proj) and then sharded, - DTensor.full_tensor() might produce an interleaved layout like [G0, U0, G1, U1, ...] - along the sharded dimension. This function reorders it to [G0, G1, ..., U0, U1, ...]. - This is an inverse operation to get_packed_weights. - - Args: - reconstructed_tensor: The tensor reconstructed from DTensor (e.g., via .full_tensor().contiguous()). - sharded_dim: The dimension index in the reconstructed_tensor that was originally sharded. - world_size: The tensor parallel world size. - num_packed_projs: The number of projections that were packed together (e.g., 2 for gate_up_proj). - - Returns: - The reordered tensor in canonical packed format. - """ - - if num_blocks != 2: - raise ValueError( - "Num blocks different from 2 is not supported yet. This is most likely a bug in your implementation as we only pack gate and up projections together." - ) - - actual_sharded_dim = sharded_dim if sharded_dim >= 0 else sharded_dim + packed_parameter.ndim - total_size_on_sharded_dim = packed_parameter.shape[actual_sharded_dim] - original_block_size_on_dim = total_size_on_sharded_dim // num_blocks - shard_chunk_size = original_block_size_on_dim // world_size - - prefix_shape = packed_parameter.shape[:actual_sharded_dim] - suffix_shape = packed_parameter.shape[actual_sharded_dim + 1 :] - - tensor_view = packed_parameter.view( - *prefix_shape, - world_size, - num_blocks, - shard_chunk_size, - *suffix_shape, - ) - - # Permute to bring num_packed_projs first, then world_size, then shard_chunk_size - # This groups all chunks of G together, then all chunks of U together. - # Target order of these middle dimensions: (num_packed_projs, world_size, shard_chunk_size) - # Current order of view's middle dimensions: (world_size, num_packed_projs, shard_chunk_size) - # Absolute indices of the dimensions to be permuted (world_size, num_packed_projs) - axis_ws_abs = len(prefix_shape) - axis_npp_abs = len(prefix_shape) + 1 - - permute_order = list(range(tensor_view.ndim)) - permute_order[axis_ws_abs], permute_order[axis_npp_abs] = permute_order[axis_npp_abs], permute_order[axis_ws_abs] - - tensor_permuted = tensor_view.permute(*permute_order) - - # Reshape back to the original tensor's ndim, with the sharded dimension now correctly ordered as [G_all, U_all]. - # The final shape should be the same as reconstructed_tensor. - final_ordered_tensor = tensor_permuted.reshape_as(packed_parameter) - - return final_ordered_tensor - - -def get_tensor_shard(param, empty_param, device_mesh, rank, dim, tensor_idx: int | None = None): - """ - Generalized tensor sharding across a multi-dimensional device mesh. - Extract only the fraction of the parameter owned by the given `rank` when the parameter would have gone sharding at provided `dim`. - Extraction follows the pytorch `Shard` placement so that sharding and materializing back to full tensor follows `Shard` semantics. - `Shard` follows torch.chunk style sharding of the tensor. We demonstrate some cases below on how sharding happens including some edge cases - such as some ranks having an empty tensor as shard. Below implementation is robut to all these cases. - - Case (1) - empty_param (16, 5120, 8190) - dim 0 - device_mesh.size() 4 - rank 0 gets (4, 5120, 8190) (0 ... 4, 5120, 8190) - rank 1 gets (4, 5120, 8190) (4 ... 8, 5120, 8190) - rank 2 gets (4, 5120, 8190) (8 ... 12, 5120, 8190) - rank 3 gets (4, 5120, 8190) (12 ... 16, 5120, 8190) - - Case (2) - empty_param (16, 5120, 8190) - dim 0 - device_mesh.size() 14 - rank 0 gets (2, 5120, 8190) (0 ... 2, 5120, 8190) - rank 1 gets (2, 5120, 8190) (2 ... 4, 5120, 8190) - rank 2 gets (2, 5120, 8190) (4 ... 6, 5120, 8190) - rank 3 gets (2, 5120, 8190) (6 ... 8, 5120, 8190) - rank 4 gets (2, 5120, 8190) (8 ... 10, 5120, 8190) - rank 5 gets (2, 5120, 8190) (10 ... 12, 5120, 8190) - rank 6 gets (2, 5120, 8190) (12 ... 14, 5120, 8190) - rank 7 gets (2, 5120, 8190) (14 ... 16, 5120, 8190) - rank 8 gets (0, 5120, 8190) - rank 9 gets (0, 5120, 8190) - rank 10 gets (0, 5120, 8190) - rank 11 gets (0, 5120, 8190) - rank 12 gets (0, 5120, 8190) - rank 13 gets (0, 5120, 8190) - - Case (3) - empty_param (16, 5120, 8190) - dim 0 - device_mesh.size() 3 - rank 0 gets (6, 5120, 8190) (0 ... 6, 5120, 8190) - rank 1 gets (6, 5120, 8190) (6 ... 12, 5120, 8190) - rank 2 gets (4, 5120, 8190) (12 ... 16, 5120, 8190) - - In case (2), empty shards are returned with appropriate dimension to allow for operations to work smoothly. - Args: - param (torch.Tensor): The tensor to shard. - empty_param (torch.Tensor): A tensor used for shape reference. - device_mesh (torch.Tensor): Shape [d_0, ..., d_n] representing the mesh. - rank (int): Global rank of the current process/device. - dim (int): Dimension along which to shard the tensor. - """ - param_dim = empty_param.ndim - mesh_shape = device_mesh.shape - world_size = reduce(operator.mul, mesh_shape) - # Get param shape: works for both torch.Tensor and safetensors TensorInfo - param_shape = list(param.shape) if isinstance(param, torch.Tensor) else param.get_shape() - if dim < 0: - dim = param_dim + dim - if empty_param.dim() == 3 and dim == 1 and len(param_shape) == 2: - dim = 0 - elif empty_param.dim() == 3 and dim == 2 and len(param_shape) == 2: - dim = 1 - - shard_size = math.ceil(param_shape[dim] / world_size) - start = rank * shard_size - end = min(start + shard_size, param_shape[dim]) - - if dim >= param_dim: - raise ValueError(f"dim {dim} is out of bounds for tensor of dimension {param_dim}") - - if rank >= world_size: - raise ValueError(f"Rank {rank} is out of bounds for mesh size {world_size}") - - # we have the full tensor not 1 part of it. - # in that case, we just assume that the weight was properly saved - # and thus because we TP if the layer is colwise it should not use this. Layer should be packed_colwise - # to inform that it needs to read form a packed tensor. It will also take care of the module list thingy. - # here we take care of potential chunking / layer split / layer chunking. - # The only "hard" case is? if we collect q,k,v -> merge it into qkv. In that case - # actually we still shard dim=0 does not change - # so only case is if the dim of the empty param is 3 and the shard dim is 0 -> we put the - # tensor on a certain device (with the input tensor_index) - if tensor_idx is not None and empty_param.dim() == 3 and dim == 0 and len(param_shape) == 2: - # special case we don't "shard" just send this entire tensor to the correct rank. - if start <= tensor_idx < end: - # this tensor does need to be materialized on this device: - return param[:] - else: - return torch.empty([], dtype=torch.int64, device=rank) - - slice_indices = [slice(None)] * len(param_shape) - - if start < param_shape[dim]: - slice_indices[dim] = slice(start, end) - param = param[tuple(slice_indices)] - if isinstance(param, list): # TODO handle the modulelist case! - param = [p[:] for p in param] - return param - - param_shape[dim] = 0 - return torch.empty(tuple(param_shape), dtype=torch.int64) # empty allocates memory.... - - -def _split_along_last_dim(x, world_size): - """Split tensor along last dimension into world_size chunks.""" - return torch.chunk(x, world_size, dim=-1) - - -# ============================================================================= -# Distributed Communication Primitives -# ============================================================================= -# -# Naming convention: -# - Functions describe their FORWARD behavior -# - Backward behavior is the "conjugate" operation for gradient flow -# -# Available operations: -# ┌────────────────────┬─────────────────────┬─────────────────────┐ -# │ Function │ Forward │ Backward │ -# ├────────────────────┼─────────────────────┼─────────────────────┤ -# │ all_reduce │ all-reduce (sum) │ identity │ -# │ all_reduce_backward│ identity │ all-reduce (sum) │ -# │ all_gather │ all-gather │ split (local chunk) │ -# │ split │ split (local chunk) │ all-gather │ -# │ reduce_scatter │ reduce-scatter │ all-gather │ -# └────────────────────┴─────────────────────┴─────────────────────┘ -# =================== - - -class _AllReduceBackward(torch.autograd.Function): - """Identity forward, all-reduce backward. Used before colwise layers (f in Megatron).""" - - @staticmethod - def forward(ctx, x, device_mesh): - ctx.device_mesh = device_mesh - return x - - @staticmethod - def backward(ctx, grad_output): - device_mesh = ctx.device_mesh - if device_mesh.size() == 1: - return grad_output, None - grad_output = grad_output.contiguous() - dist.all_reduce(grad_output, op=dist.ReduceOp.SUM, group=device_mesh.get_group()) - return grad_output, None - - -class _AllReduceForward(torch.autograd.Function): - """All-reduce forward, identity backward. Used after rowwise layers (g in Megatron).""" - - @staticmethod - def forward(ctx, x, device_mesh): - if device_mesh.size() == 1: - return x - dist.all_reduce(x, op=dist.ReduceOp.SUM, group=device_mesh.get_group()) - return x - - @staticmethod - def backward(ctx, grad_output): - return grad_output, None - - -class _AllGather(torch.autograd.Function): - """All-gather forward, split backward. Gathers sharded outputs.""" - - @staticmethod - def forward(ctx, x, device_mesh): - ctx.device_mesh = device_mesh - world_size = device_mesh.size() - - if world_size == 1: - return x - - last_dim = x.dim() - 1 - rank = device_mesh.get_local_rank() - group = device_mesh.get_group() - - x = x.contiguous() - tensor_list = [torch.empty_like(x) for _ in range(world_size)] - tensor_list[rank] = x - dist.all_gather(tensor_list, x, group=group) - return torch.cat(tensor_list, dim=last_dim).contiguous() - - @staticmethod - def backward(ctx, grad_output): - device_mesh = ctx.device_mesh - world_size = device_mesh.size() - - if world_size == 1: - return grad_output, None - - rank = device_mesh.get_local_rank() - chunks = _split_along_last_dim(grad_output, world_size) - return chunks[rank].contiguous(), None - - -class _Split(torch.autograd.Function): - """Split forward, all-gather backward. Scatters replicated input.""" - - @staticmethod - def forward(ctx, x, device_mesh): - ctx.device_mesh = device_mesh - world_size = device_mesh.size() - - if world_size == 1: - return x - - rank = device_mesh.get_local_rank() - chunks = _split_along_last_dim(x, world_size) - return chunks[rank].contiguous() - - @staticmethod - def backward(ctx, grad_output): - device_mesh = ctx.device_mesh - world_size = device_mesh.size() - - if world_size == 1: - return grad_output, None - - last_dim = grad_output.dim() - 1 - rank = device_mesh.get_local_rank() - group = device_mesh.get_group() - - grad_output = grad_output.contiguous() - tensor_list = [torch.empty_like(grad_output) for _ in range(world_size)] - tensor_list[rank] = grad_output - dist.all_gather(tensor_list, grad_output, group=group) - return torch.cat(tensor_list, dim=last_dim).contiguous(), None - - -class _ReduceScatter(torch.autograd.Function): - """Reduce-scatter forward, all-gather backward. For sequence parallel.""" - - @staticmethod - def forward(ctx, x, device_mesh): - ctx.device_mesh = device_mesh - world_size = device_mesh.size() - - if world_size == 1: - return x - - last_dim = x.dim() - 1 - group = device_mesh.get_group() - - input_chunks = list(x.chunk(world_size, dim=last_dim)) - output_shape = list(x.shape) - output_shape[last_dim] //= world_size - output = torch.empty(output_shape, dtype=x.dtype, device=x.device) - - dist.reduce_scatter(output, input_chunks, op=dist.ReduceOp.SUM, group=group) - return output - - @staticmethod - def backward(ctx, grad_output): - device_mesh = ctx.device_mesh - world_size = device_mesh.size() - - if world_size == 1: - return grad_output, None - - last_dim = grad_output.dim() - 1 - rank = device_mesh.get_local_rank() - group = device_mesh.get_group() - - grad_output = grad_output.contiguous() - tensor_list = [torch.empty_like(grad_output) for _ in range(world_size)] - tensor_list[rank] = grad_output - dist.all_gather(tensor_list, grad_output, group=group) - return torch.cat(tensor_list, dim=last_dim).contiguous(), None - - -# ============================================================================= -# Convenience wrappers -# ============================================================================= - - -def all_reduce_backward(x, device_mesh): - """Identity forward, all-reduce backward. Use before colwise layers.""" - return _AllReduceBackward.apply(x, device_mesh) - - -def all_reduce_forward(x, device_mesh): - """All-reduce forward, identity backward. Use after rowwise layers.""" - return _AllReduceForward.apply(x, device_mesh) - - -def all_gather(x, device_mesh): - """All-gather forward, split backward.""" - return _AllGather.apply(x, device_mesh) - - -def split(x, device_mesh): - """Split forward, all-gather backward.""" - return _Split.apply(x, device_mesh) - - -def reduce_scatter(x, device_mesh): - """Reduce-scatter forward, all-gather backward.""" - return _ReduceScatter.apply(x, device_mesh) - - -def distribute_module( - module: nn.Module, - device_mesh=None, - input_fn=None, - output_fn=None, -) -> nn.Module: - """ - Copy pasted from torch's function but we remove the communications (partitioning) - as well as buffer registering that is similarly not efficient. - """ - if input_fn is not None: - module.register_forward_pre_hook(lambda mod, inputs: input_fn(mod, inputs, device_mesh)) - if output_fn is not None: - module.register_forward_hook(lambda mod, inputs, outputs: output_fn(mod, outputs, device_mesh)) - return module - - -class TensorParallelLayer: - """General tensor parallel layer for transformers""" - - device_mesh = None - rank = None - empty_param = None - - def __init__(self, device_mesh=None, rank=None, empty_param=None): - self.rank = rank - self.device_mesh = device_mesh - self.empty_param = empty_param - - def _prepare_input_fn(self, mod, inputs, device_mesh): - raise NotImplementedError - - def _prepare_output_fn(self, mod, outputs, device_mesh): - raise NotImplementedError - - def shard_tensor( - self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None - ) -> torch.Tensor: - raise NotImplementedError - - def prepare_module_tp(self, module: nn.Module, device_mesh, **kwargs) -> nn.Module: - distribute_module( - module, - device_mesh, - self._prepare_input_fn, - self._prepare_output_fn, - ) - - def get_expected_sharded_shape(self, full_shape: tuple[int, ...] | torch.Size) -> tuple[int, ...]: - """ - Compute the expected shape after TP sharding for a given full shape. - - Args: - full_shape: The full (unsharded) parameter shape - - Returns: - The expected sharded shape for this rank - """ - # Default: no sharding, return full shape - return tuple(full_shape) - - def update_module_attributes(self, module: nn.Module): - """ - Update module attributes (e.g. in_features, out_features) to reflect sharded dimensions. - - Args: - module: The module to update - - Returns: - None, update the module in-place - """ - pass - - -class ColwiseParallel(TensorParallelLayer): - """ - Column-wise parallel: weight is sharded on dim -2 (output features). - Forward: input replicated -> output sharded on last dim. - If gather_output=True, output is all-gathered to produce full tensor. - """ - - def __init__(self, gather_output: bool = False, **kwargs): - super().__init__(**kwargs) - self.gather_output = gather_output - - def _prepare_input_fn(self, mod, inputs, device_mesh): - input_tensor = inputs[0] if inputs else inputs - return all_reduce_backward(input_tensor, device_mesh) - - def _prepare_output_fn(self, mod, outputs, device_mesh): - if self.gather_output: - return all_gather(outputs, device_mesh) - return outputs - - def shard_tensor( - self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None - ) -> torch.Tensor: - # If only 1 dim, shard this one (usually it's a `bias`) - dim = param.dim() if isinstance(param, torch.Tensor) else len(param.get_shape()) - if dim == 1: - parameter = get_tensor_shard(param, self.empty_param, self.device_mesh, self.rank, -1) - else: - parameter = get_tensor_shard(param, self.empty_param, self.device_mesh, self.rank, -2) - return parameter.to(device=device, dtype=dtype) - - def get_expected_sharded_shape(self, full_shape: tuple[int, ...] | torch.Size) -> tuple[int, ...]: - world_size = self.device_mesh.size() - shape = list(full_shape) - # Colwise shards dim -2, but 1D tensors (bias) shard on dim -1 - dim = -1 if len(shape) == 1 else -2 - dim = len(shape) + dim if dim < 0 else dim - shard_size = math.ceil(shape[dim] / world_size) - start = self.rank * shard_size - end = min(start + shard_size, shape[dim]) - shape[dim] = end - start - return tuple(shape) - - def update_module_attributes(self, module: nn.Module): - # If we gather the output, the output dimension of the module is not sharded, so no need to update out_features. - # Otherwise, we need to update out_features to reflect the sharded dimension. - if not self.gather_output and hasattr(module, "out_features"): - module.out_features = self.get_expected_sharded_shape((module.out_features,))[0] - - -class ReplicatedWithGradAllReduce(TensorParallelLayer): - """ - Replicated parameter with gradient all-reduce. - - For parameters like q_norm/k_norm that sit between colwise and rowwise - layers. The parameter is replicated (not sharded), but its gradient - accumulates from local heads only in TP mode. This class registers a - backward hook to all-reduce the parameter gradient. - """ - - def _prepare_input_fn(self, mod, inputs, device_mesh): - return inputs - - def _prepare_output_fn(self, mod, outputs, device_mesh): - return outputs - - def shard_tensor(self, param, tensor_idx=None, device=None, dtype=None): - return param[...].to(device=device, dtype=dtype) - - def prepare_module_tp(self, module, device_mesh, **kwargs): - # Use a module-level backward hook (not param.register_hook) because parameters are replaced during weight loading after this method runs. - # Module hooks survive parameter replacement. - def _backward_hook(mod, grad_input, grad_output, mesh=device_mesh): - for param in mod.parameters(): - if param.grad is not None: - all_reduce_forward(param.grad, mesh) - - module.register_full_backward_hook(_backward_hook) - - -class MlaKvAProjParallel(TensorParallelLayer): - """ - For MLA attention used in DeepSeek-V2 style models (deepseek_v2, longcat_flash, glm_moe_dsa, glm4_moe_lite): - kv_a_proj_with_mqa output is [kv_lora_rank + qk_rope_head_dim] (can have different naming but important thing - to understand is that it is split) - Example below (from modeling_longcat_flash.py): - - kv_a_proj_with_mqa - | - split - / \ - k_pass k_rot <-- "bypasses kv_b_proj" - | | (goes straight to attention, - kv_a_layernorm | never touches kv_b_proj) - | | - kv_b_proj | - (colwise) | - | | - k_pass k_rot - \\ / - cat - | - key_states - - k_pass is passed to kv_b_proj (colwise) which has built-in all_reduce_backward so we don't have a partial gradient for it. - However, k_rot goes straight to attention, never touches kv_b_proj. So we need to average gradient across all ranks otherwise we only get gradient for one rank (partial gradient). - """ - - def _prepare_output_fn(self, mod, output, device_mesh): - if not hasattr(mod.config, "qk_rope_head_dim"): - raise AttributeError( - f"Config for {type(mod).__name__} does not have `qk_rope_head_dim`. " - "MlaKvAProjParallel requires `qk_rope_head_dim` to be defined in the model config. " - "Please add it to the model's config or update the TP plan mapping." - ) - rope_dim = mod.config.qk_rope_head_dim - pass_output, rope_output = output.split([output.shape[-1] - rope_dim, rope_dim], dim=-1) - rope_output = all_reduce_backward(rope_output, device_mesh) - return torch.cat([pass_output, rope_output], dim=-1) - - def shard_tensor(self, param, tensor_idx=None, device=None, dtype=None): - return param[...].to(device=device, dtype=dtype) - - def prepare_module_tp(self, module, device_mesh, config=None, **kwargs): - module.config = config - distribute_module(module, device_mesh, output_fn=self._prepare_output_fn) - - -class RowwiseParallel(TensorParallelLayer): - """ - Row-wise parallel: weight is sharded on dim -1 (input features). - Forward: input (optionally split) -> output partial -> all-reduce to replicate. - - Args: - split_input: If True, splits replicated input before matmul. Use when input - comes from a non-parallelizable operation (chunk/slice). - Default False (expects pre-sharded input from colwise layer). - """ - - def __init__(self, split_input: bool = False, **kwargs): - super().__init__(**kwargs) - self.split_input = split_input - - def _prepare_input_fn(self, mod, inputs, device_mesh): - if hasattr(mod, "bias") and mod.bias is not None: - mod._bias = mod.bias - mod.bias = None - - input_tensor = inputs[0] if inputs else inputs - - if self.split_input: - # Input is replicated, split it to match sharded weight - return split(input_tensor, device_mesh) - return input_tensor - - def _prepare_output_fn(self, mod, outputs, device_mesh): - outputs = all_reduce_forward(outputs, device_mesh) - if hasattr(mod, "_bias") and mod._bias is not None: - outputs = outputs + mod._bias - return outputs - - def shard_tensor( - self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None - ) -> torch.Tensor: - # If only 1 dim, it should not be sharded (usually it's a `bias`) - dim = param.dim() if isinstance(param, torch.Tensor) else len(param.get_shape()) - if dim == 1: - parameter = param[...] - else: - parameter = get_tensor_shard(param, self.empty_param, self.device_mesh, self.rank, -1) - return parameter.to(device=device, dtype=dtype) - - def get_expected_sharded_shape(self, full_shape: tuple[int, ...] | torch.Size) -> tuple[int, ...]: - # 1D tensors (bias) are NOT sharded in rowwise - if len(full_shape) == 1: - return tuple(full_shape) - world_size = self.device_mesh.size() - shape = list(full_shape) - dim = -1 - dim = len(shape) + dim if dim < 0 else dim - shard_size = math.ceil(shape[dim] / world_size) - start = self.rank * shard_size - end = min(start + shard_size, shape[dim]) - shape[dim] = end - start - return tuple(shape) - - def update_module_attributes(self, module: nn.Module): - if hasattr(module, "in_features"): - # To fall in the 2D case in get_expected_sharded_shape, - # otherwise it will be treated as 1D and not sharded - shape = (1, module.in_features) - module.in_features = self.get_expected_sharded_shape(shape)[1] - - -class PackedColwiseParallel(ColwiseParallel): - """Packed column-wise parallel for fused weights like gate_up_proj.""" - - def shard_tensor( - self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None - ) -> torch.Tensor: - # If only 1 dim, shard this one (usually it's a `bias`) - dim = param.dim() if isinstance(param, torch.Tensor) else len(param.get_shape()) - if dim == 1: - parameter = get_tensor_shard(param, self.empty_param, self.device_mesh, self.rank, -1) - else: - expected_shape = self.get_expected_sharded_shape(self.empty_param.shape) - if dim < len(expected_shape): - # Input is unpacked (e.g., gate_proj that will be concatenated to gate_up_proj) - # Use regular tensor shard - concatenation will happen after - parameter = get_tensor_shard(param, self.empty_param, self.device_mesh, self.rank, -2) - else: - # Input is already packed, use packed sharding - parameter = get_packed_weights(param, self.empty_param, self.device_mesh, self.rank, -2) - return parameter.to(device=device, dtype=dtype) - - -class PackedRowwiseParallel(RowwiseParallel): - """Packed row-wise parallel for fused weights like gate_up_proj.""" - - def shard_tensor( - self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None - ) -> torch.Tensor: - # If only 1 dim, it should not be sharded (usually it's a `bias`) - dim = param.dim() if isinstance(param, torch.Tensor) else len(param.get_shape()) - if dim == 1: - parameter = param[...] - else: - # Check if input tensor is unpacked (shape mismatch with expected packed size) - # This happens when using MergeModulelist + Concatenate for fused weights like gate_up_proj - param_shape = param.shape if isinstance(param, torch.Tensor) else param.get_shape() - expected_packed_dim = self.empty_param.shape[-1] if self.empty_param.dim() >= 1 else 0 - actual_dim = param_shape[-1] if len(param_shape) >= 1 else 0 - - if actual_dim < expected_packed_dim: - # Input is unpacked, use regular tensor shard - parameter = get_tensor_shard(param, self.empty_param, self.device_mesh, self.rank, -1) - else: - # Input is already packed, use packed sharding - parameter = get_packed_weights(param, self.empty_param, self.device_mesh, self.rank, -1) - return parameter.to(device=device, dtype=dtype) - - -class EmbeddingParallel(TensorParallelLayer): - """EmbeddingParallel: shards embedding table, handles masked lookups for vocab parallelism.""" - - def __init__(self, *, embedding_dim_sharding: int = 0, **kwargs): - super().__init__(**kwargs) - self.embedding_dim_sharding = embedding_dim_sharding - - def _prepare_input_fn(self, mod, inputs, device_mesh): - input_tensor = inputs[0] if inputs else inputs - - # For vocab-parallel (dim 0), we need to handle masking and offsetting - if self.embedding_dim_sharding == 0: - rank = device_mesh.get_local_rank() - - # Get vocab range for this rank - # Use weight.shape[0] to get the actual local (sharded) size, not num_embeddings - # which may not be updated after sharding - per_partition_size = mod.weight.shape[0] - vocab_start_index = rank * per_partition_size - vocab_end_index = vocab_start_index + per_partition_size - - # Build mask for out-of-vocabulary tokens - input_mask = (input_tensor < vocab_start_index) | (input_tensor >= vocab_end_index) - mod._input_mask = input_mask - - # Offset input to local indices and mask invalid ones - masked_input = input_tensor.clone() - vocab_start_index - masked_input[input_mask] = 0 # Set to valid local index - - return masked_input - - return input_tensor - - def _prepare_output_fn(self, mod, outputs, device_mesh): - # For vocab-parallel (dim 0), zero out embeddings for out-of-range tokens before all-reduce - if self.embedding_dim_sharding == 0 and hasattr(mod, "_input_mask"): - input_mask = mod._input_mask - # Use multiplication instead of in-place assignment to preserve gradients - mask_expanded = input_mask.unsqueeze(-1).expand_as(outputs) - outputs = outputs * (~mask_expanded).to(outputs.dtype) - del mod._input_mask - - return all_reduce_forward(outputs, device_mesh) - - def shard_tensor( - self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None - ) -> torch.Tensor: - # If only 1 dim, shard this one (usually it's a `bias`) - dim = param.dim() if isinstance(param, torch.Tensor) else len(param.get_shape()) - if dim == 1: - parameter = get_tensor_shard(param, self.empty_param, self.device_mesh, self.rank, -1) - else: - parameter = get_tensor_shard( - param, - self.empty_param, - self.device_mesh, - self.rank, - self.embedding_dim_sharding, - ) - return parameter.to(device=device, dtype=dtype) - - def get_expected_sharded_shape(self, full_shape: tuple[int, ...] | torch.Size) -> tuple[int, ...]: - world_size = self.device_mesh.size() - shape = list(full_shape) - # EmbeddingParallel shards on self.embedding_dim_sharding (default 0) - # 1D tensors (bias) shard on dim -1 - dim = -1 if len(shape) == 1 else self.embedding_dim_sharding - dim = len(shape) + dim if dim < 0 else dim - shard_size = math.ceil(shape[dim] / world_size) - start = self.rank * shard_size - end = min(start + shard_size, shape[dim]) - shape[dim] = end - start - return tuple(shape) - - def update_module_attributes(self, module: nn.Module): - if hasattr(module, "num_embeddings") and self.embedding_dim_sharding == 0: - module.num_embeddings = self.get_expected_sharded_shape((module.num_embeddings,))[0] - if hasattr(module, "embedding_dim") and self.embedding_dim_sharding == 1: - module.embedding_dim = self.get_expected_sharded_shape((module.embedding_dim,))[0] - - -class SequenceParallel(TensorParallelLayer): - """ - Sequence Parallel: input/output sharded on sequence dimension. - Weights are replicated. - """ - - def __init__(self, sequence_dim: int = 1, use_local_output: bool = False, use_dtensor=False, **kwargs): - super().__init__(**kwargs) - self.sequence_dim = sequence_dim - - def _prepare_input_fn(self, mod, inputs, device_mesh): - input_tensor = inputs[0] if inputs else inputs - # For sequence parallel, input is sharded on sequence dim - # All-gather for the layer, then reduce-scatter after - return all_gather(input_tensor, device_mesh) - - def _prepare_output_fn(self, mod, outputs, device_mesh): - return reduce_scatter(outputs, device_mesh) - - def shard_tensor( - self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None - ) -> torch.Tensor: - return param[...].to(device=device, dtype=dtype) - - -class GroupedGemmParallel(TensorParallelLayer): - """ - Applies Expert Parallelism to MoE experts by loading the correct experts on each device. - """ - - def __init__(self, **kwargs): - super().__init__(**kwargs) - - def shard_tensor( - self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None - ) -> torch.Tensor: - global_num_experts = self.empty_param.shape[0] - if global_num_experts % self.device_mesh.size() != 0: - raise ValueError( - f"Global number of experts must be divisible by number of devices: {global_num_experts} % {self.device_mesh.size()} != 0" - ) - local_num_experts = global_num_experts // self.device_mesh.size() - shard_size = local_num_experts - if isinstance(device, torch.device): - device = device.index if device.index is not None else 0 - start = device * shard_size - end = (device + 1) * shard_size - # special case we don't "shard" just send this entire tensor to the correct rank. - shape = param.get_shape() if not isinstance(param, torch.Tensor) else param.shape - if tensor_idx is not None and start <= tensor_idx < end: - # this tensor does need to be materialized on this device: - return param[:].to(device=device) - elif tensor_idx is None: # a bias or a weight, but already merged - return param[start:end].to(device=device, dtype=dtype) - elif len(shape) >= 1 and tensor_idx is not None: - return None - else: # bias case - return param[:].to(device=device, dtype=dtype) - - def get_expected_sharded_shape(self, full_shape: tuple[int, ...] | torch.Size) -> tuple[int, ...]: - # GroupedGemm shards on dim 0 (experts dimension) - world_size = self.device_mesh.size() - shape = list(full_shape) - local_num_experts = shape[0] // world_size - shape[0] = local_num_experts - return tuple(shape) - - def update_module_attributes(self, module: nn.Module): - if hasattr(module, "num_experts"): - module.num_experts = self.get_expected_sharded_shape((module.num_experts,))[0] - - -class RouterParallel(TensorParallelLayer): - """ - Allows to reshape the router scores to support running expert parallel. - """ - - def __init__(self, **kwargs): - super().__init__(**kwargs) - - def _prepare_input_fn(self, mod, inputs, device_mesh): - return inputs[0] if inputs else inputs - - def _prepare_output_fn(self, mod, outputs, device_mesh): - """ - Imagine if you had 4 tokens, top_k = 4, and 128experts. - With EP = 8. The num_local_expert should be 128/8 = 16 - Imagine router_indices being: - [ 52, 42, 119, 67], - [102, 89, 61, 40], - [ 82, 103, 4, 34], - [ 93, 23, 109, 11], - - then you can map which rank should be getting which values - - [3, 2, 7, 4], - [6, 5, 3, 2], - [5, 6, 0, 2], - [5, 1, 6, 0], - - Thus for say rank 0, you fill with 16 (num_local_expert) the index tensor - - [ 16, 16, 16, 16], - [ 16, 16, 16, 16], - [ 16, 16, 4, 16], - [ 16, 16, 16, 11], - - This works well. For another rank you need to make sure you round to num_local_expert - because the next operation will one hot encode the router index vector. - - This allows us to know directly which local expert is hit. - Similarly the scores are indexed with something created form - router_indices. - - The kinda naive training loop that we use for device_map "auto" uses a similar logic. - Here we are just making each rank believe that he is alone, and he computes his part of the hiddenstates. - Mask invalid indices with num_local_expert for one-hot encoding, so the computes will skip the masking index. - """ - ep_rank, ep_size = device_mesh.get_local_rank(), device_mesh.size() - if mod.num_experts % ep_size != 0: - raise ValueError( - f"The number of experts must be divisible by number of ep_size: {mod.num_experts} % {ep_size} != 0" - ) - num_local_experts = mod.num_experts // ep_size - router_logits, router_scores, router_indices = outputs - router_scores = torch.zeros_like(router_logits).scatter_(1, router_indices, router_scores) - router_scores = router_scores[:, ep_rank * num_local_experts : (ep_rank + 1) * num_local_experts] - router_indices = router_indices.masked_fill((router_indices // num_local_experts) != ep_rank, -1) - # As -1 % 1 is 0, we can only use mask fill when num_local_experts is 1 - if num_local_experts > 1: - router_indices = torch.fmod(router_indices, num_local_experts) - else: - router_indices = router_indices.masked_fill(router_indices > 0, 0).masked_fill(router_indices < 0, -1) - router_indices = router_indices.masked_fill(router_indices == -1, num_local_experts) - return router_logits, router_scores, router_indices - - def shard_tensor( - self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None - ) -> torch.Tensor: - return param[...].to(device=device, dtype=dtype) - - -class MoeTensorParalellExperts(TensorParallelLayer): - """ - Note: For tensor parallel, the MoEExpertsParallel TP layer handles gradient sync: - - all_reduce_backward on hidden_states (for colwise gate_up_proj gradient) - - all_reduce_backward on top_k_weights (for router gradient) - - all_reduce_forward on output (for partial expert outputs) - """ - - def __init__(self, **kwargs): - super().__init__(**kwargs) - - def _prepare_input_fn(self, mod, inputs, device_mesh): - # inputs = (hidden_states, top_k_index, top_k_weights) - hidden_states = inputs[0] - top_k_index = inputs[1] - top_k_weights = inputs[2] - - # all_reduce_backward on hidden_states for correct colwise (gate_up_proj) gradient - hidden_states = all_reduce_backward(hidden_states, device_mesh) - - # all_reduce_backward on routing weights for correct router gradient - # This is needed because ∂L/∂routing_weights = ∂L/∂output * partial_expert_output - # and partial_expert_output is different on each GPU before all-reduce - top_k_weights = all_reduce_backward(top_k_weights, device_mesh) - - return (hidden_states, top_k_index, top_k_weights) - - def _prepare_output_fn(self, mod, outputs, device_mesh): - # all_reduce_forward to sum partial expert outputs across GPUs - return all_reduce_forward(outputs, device_mesh) - - def shard_tensor( - self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None - ) -> torch.Tensor: - # This class doesn't shard tensors - sharding is handled by packed_colwise/rowwise - # on the individual weight tensors (gate_up_proj/down_proj) - return param[...].to(device=device, dtype=dtype) - - -class MoeIdentityExpertParallel(TensorParallelLayer): - """ - TP class for zero/identity experts in MoE layers. - - Under TP, the parent MoeTensorParalellExperts does all_reduce_forward (sum) - on the expert module output. Identity experts produce the same output on - every rank, so the sum gives world_size * output. This class divides the - input by world_size to compensate. - """ - - def _prepare_input_fn(self, mod, inputs, device_mesh): - input_tensor = inputs[0] if inputs else inputs - # TODO(fmom): when 2D-device mesh, need to select a //-ism axis to divide the input tensor by. - return input_tensor / device_mesh.size() - - def shard_tensor(self, param, tensor_idx=None, device=None, dtype=None): - return param[...].to(device=device, dtype=dtype) - - def prepare_module_tp(self, module, device_mesh, **kwargs): - distribute_module(module, device_mesh, input_fn=self._prepare_input_fn) - - -class ParallelInterface(GeneralInterface): - # Class instance object, so that a call to `register` can be reflected into all other files correctly, even if - # a new instance is created (in order to locally override a given entry) - _global_mapping = ( - { - "embedding_rowwise": EmbeddingParallel(embedding_dim_sharding=0), - "embedding_colwise": EmbeddingParallel(embedding_dim_sharding=1), - "colwise_gather_output": ColwiseParallel(gather_output=True), - "colwise": ColwiseParallel(), - "rowwise": RowwiseParallel(), - "rowwise_split_input": RowwiseParallel(split_input=True), - "packed_colwise": PackedColwiseParallel(), - "packed_rowwise": PackedRowwiseParallel(), - "sequence_parallel": SequenceParallel(), - "grouped_gemm": GroupedGemmParallel(), - "ep_router": RouterParallel(), - "moe_tp_experts": MoeTensorParalellExperts(), - "moe_identity_expert": MoeIdentityExpertParallel(), - "replicated_with_grad_allreduce": ReplicatedWithGradAllReduce(), - "mla_kv_a_proj": MlaKvAProjParallel(), - } - if is_torch_available() and _torch_distributed_available - else {} - ) - - # Map plan names to sharding dimensions for weights - # For weights: colwise shards dim -2, rowwise shards dim -1 - # For embedding: rowwise shards dim 0 (vocab), colwise shards dim -2 (hidden) - plan_to_weight_dim: dict[str, int | None] = { - "colwise": -2, - "colwise_gather_output": -2, - "packed_colwise": -2, - "rowwise": -1, - "rowwise_split_input": -1, - "packed_rowwise": -1, - "embedding_rowwise": 0, - "embedding_colwise": 1, - "sequence_parallel": None, - "replicated_with_grad_allreduce": None, - "mla_kv_a_proj": None, - } - - # Bias sharding: colwise shards bias, rowwise doesn't (bias is replicated and all-reduced) - plan_to_bias_dim: dict[str, int | None] = { - "colwise": -1, - "colwise_gather_output": -1, - "packed_colwise": -1, - "rowwise": None, - "rowwise_split_input": None, - "packed_rowwise": None, - "embedding_rowwise": None, - "embedding_colwise": None, - "sequence_parallel": None, - "replicated_with_grad_allreduce": None, - "mla_kv_a_proj": None, - } - - @classmethod - def register_plan_to_weight_dim(cls, key: str, value: int | None): - cls.plan_to_weight_dim[key] = value - - @classmethod - def register_plan_to_bias_dim(cls, key: str, value: int | None): - cls.plan_to_bias_dim[key] = value - - -ALL_PARALLEL_STYLES: ParallelInterface = ParallelInterface() + t = tensor.detach() + if t.device.type != "cpu": + t = t.to(device="cpu") + out = torch.empty(t.shape, dtype=t.dtype, device="cpu") + out.copy_(t) + return out.contiguous() # ============================================================================= @@ -1286,200 +89,11 @@ def register_plan_to_bias_dim(cls, key: str, value: int | None): # ============================================================================= -def gather_full_tensor( - local_tensor: torch.Tensor, shard_dim: int, device_mesh: dist.device_mesh.DeviceMesh -) -> torch.Tensor: - """ - All-gather a sharded tensor along the specified dimension to reconstruct the full tensor. - - Args: - local_tensor: The local shard of the tensor on this rank - shard_dim: The dimension along which the tensor was sharded - device_mesh: The device mesh for distributed communication - - Returns: - The full reconstructed tensor (same on all ranks) - """ - world_size = device_mesh.size() - # In case of TP+DP configuration, the TP group should be used for gathering, not the full DP group - process_group = device_mesh.get_group("tp") if "tp" in (device_mesh.mesh_dim_names or {}) else None - - # Normalize negative dimension - if shard_dim < 0: - shard_dim = local_tensor.ndim + shard_dim - - # Gather all shards - gathered_tensors = [torch.empty_like(local_tensor) for _ in range(world_size)] - dist.all_gather(gathered_tensors, local_tensor.contiguous(), group=process_group) - - # Concatenate along the shard dimension - return torch.cat(gathered_tensors, dim=shard_dim) - - -def gather_state_dict_for_save( - state_dict: dict[str, torch.Tensor], - tp_plan: dict[str, str], - device_mesh, - tp_size: int, -) -> dict[str, torch.Tensor]: - """ - Gather sharded tensors to reconstruct full tensors for saving. - - This function all-gathers each sharded tensor along its shard dimension - to reconstruct the full unsharded tensor for checkpoint saving. - - Args: - state_dict: The model state dict with local sharded tensors - tp_plan: The tensor parallel plan mapping layer patterns to shard styles - device_mesh: The device mesh for distributed communication - tp_size: The tensor parallel world size - - Returns: - State dict with full (gathered) tensors - """ - # Use the global mappings from ParallelInterface (can be extended by users) - plan_to_weight_dim = ALL_PARALLEL_STYLES.plan_to_weight_dim - plan_to_bias_dim = ALL_PARALLEL_STYLES.plan_to_bias_dim - - result = {} - for key, tensor in state_dict.items(): - # Find the matching TP plan for this parameter - param_name = key.rsplit(".", 1)[0] if "." in key else key - param_type = key.rsplit(".", 1)[1] if "." in key else None - generic_param_name = re.sub(r"\d+", "*", param_name) - # Also check the full key for nn.Parameter (e.g., MoE experts without .weight suffix) - generic_full_key = re.sub(r"\d+", "*", key) - - # Check if this parameter has a TP plan - current_plan = None - if generic_full_key in tp_plan: - # Full key match (e.g., "model.layers.*.mlp.experts.gate_up_proj" for MoE experts) - current_plan = tp_plan[generic_full_key] - elif generic_param_name in tp_plan: - current_plan = tp_plan[generic_param_name] - elif "." in generic_param_name: - parent_param_name = generic_param_name.rsplit(".", 1)[0] - if parent_param_name in tp_plan: - current_plan = tp_plan[parent_param_name] - - if current_plan is None or current_plan not in plan_to_weight_dim: - # Not sharded, keep as-is - result[key] = tensor - continue - - # Determine sharding dimension based on param type - if param_type == "bias": - shard_dim = plan_to_bias_dim.get(current_plan) - else: - shard_dim = plan_to_weight_dim.get(current_plan) - - if shard_dim is None: - # Replicated, keep as-is - result[key] = tensor - continue - - # Gather full tensor and handle packed weights repacking - full_tensor = gather_full_tensor(tensor, shard_dim, device_mesh) - if current_plan in ("packed_colwise", "packed_rowwise"): - full_tensor = repack_weights(full_tensor, shard_dim, tp_size, 2) - result[key] = full_tensor.contiguous() - - return result - - -def add_tensor_parallel_hooks_to_module( - model, - module, - current_module_plan, - layer_name, - device_mesh, -): - r""" - This function is called in `PretrainedModel.post_init()`. It is responsible of adding hooks - to the modules of the `model`, based on the `PretrainedModel._tp_plan`. - - This is the place where we add the `pre_forward` and `post_forwards` hooks. These are defined - for each `TensorParallelLayer` as `_prepare_input_fn` and `_prepare_output_fn`. - - Args: - model (`PretrainedModel`): The model containing the modules. - module (`nn.Module`): The current module to which we want to add the hooks. - current_module_plan (`str` or `None`): The tensor parallel plan for the current module, if any. - layer_name (`str`): The qualified name of the current module. - device_mesh (`dist.device_mesh.DeviceMesh`): The device mesh for distributed communication. - - """ - if current_module_plan is not None: - tp_layer = ALL_PARALLEL_STYLES[current_module_plan] - try: - tp_layer.prepare_module_tp(module, device_mesh, config=model.config) - except NotImplementedError as e: - logger.warning( - f"Trying to prepare {layer_name}, but it's not supported. Corresponding module: {module} Fix it's TP " - f"plan: {e}" - ) - - module._hf_tp_plan = current_module_plan - module._hf_device_mesh = device_mesh - module.__repr__ = lambda: f"{module.__repr__()}\nTP Plan: {current_module_plan}" - - -def shard_and_distribute_module( - model, param, empty_param, parameter_name, param_casting_dtype, is_contiguous, rank, device_mesh -): - r""" - This function is called in `from_pretrained` when loading a model's checkpoints. - It receives the pointer to the parameter (or the parameter itself) and takes care of "sharding". - All process run this function, so they just load the partition of the tensor that they require. - - Main uses cases: - - column / rowise parallelism, you just shard all the weights of the layer (weight and bias) - - packed layers: you slice the weights, then shard like above - - custom operation: - - you want to add an all-gather at the end of a local layer. - - you want to have a layer that is isolated from the rest of the world (because torch.DTensor does not work well with `.view` for instance) - - """ - param_name, param_type = parameter_name.rsplit(".", 1) if "." in parameter_name else parameter_name - tp_plan = model.tp_plan or {} - module_to_tp = model.get_submodule(param_name) - rank = int(rank) - current_shard_plan = _get_parameter_tp_plan(parameter_name, tp_plan) - - if dist.get_rank() == 0: - if current_shard_plan is None: - logger.info(f"Tensor sharding plan for {param_name} not found, using default 'replicate' plan.") - else: - logger.info(f"Tensor sharding plan for {param_name}: {current_shard_plan}") - - if current_shard_plan is not None: - try: - tp_layer = ALL_PARALLEL_STYLES[current_shard_plan] - tp_layer.empty_param = empty_param - tp_layer.device_mesh = device_mesh - tp_layer.rank = rank - param = tp_layer.shard_tensor(param, tensor_idx=None, dtype=param_casting_dtype, device=rank) - if is_contiguous: - param = param.contiguous() - except NotImplementedError as e: - print( - f"Trying to prepare {parameter_name}, but it's not supported. Corresponding module: {module_to_tp} Fix it's TP plan, current layer: {tp_layer} : {e}" - ) - else: - param = param[:].to(param_casting_dtype) - - # SUPER IMPORTANT we have to use setattr - # otherwise loading is crazy slow - if not isinstance(param, torch.nn.Parameter): - param = torch.nn.Parameter(param, requires_grad=empty_param.is_floating_point()) - setattr(module_to_tp, param_type, param) - tp_layer.update_module_attributes(module_to_tp) - return param - - -def verify_tp_plan(expected_keys: list[str], tp_plan: dict[str, str] | None): +def verify_tp_plan(expected_keys: list[str], tp_plan: dict[str, str | TPStyle] | None): """ Verify the TP plan of the model, log a warning if the layers that were not sharded and the rules that were not applied. + + Only weight-sharding rules (colwise, rowwise, vocab) are checked. """ if tp_plan is None: @@ -1506,35 +120,93 @@ 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, fsdp_plan=None): - """Attach distributed runtime hooks before checkpoint loading.""" - model._tp_size = tp_size - model._device_mesh = device_mesh - if distributed_config is not None: - if isinstance(distributed_config, dict): - distributed_config = DistributedConfig.from_dict(distributed_config) - model.config.distributed_config = distributed_config - # Set the new requested tp_plan on the model - if isinstance(tp_plan, dict): - model.tp_plan = tp_plan - model_plan = model.tp_plan if tp_plan is not None or tp_size is not None else None - if model_plan is not None and _torch_distributed_available: - for v in model_plan.values(): - if v not in ALL_PARALLEL_STYLES: - raise ValueError(f"Unsupported tensor parallel style {v}. Supported styles are {ALL_PARALLEL_STYLES}") - for name, module in model.named_modules(): - if not getattr(module, "_is_hooked", False): - plan = _get_parameter_tp_plan(parameter_name=name, tp_plan=model_plan, is_weight=False) - add_tensor_parallel_hooks_to_module( - model, - module, - plan, - name, - device_mesh, - ) - module._is_hooked = True - if fsdp_plan is not None: - from .fsdp import apply_fsdp2 +@dataclass(frozen=True) +class TPStyle: + kind: Literal["colwise", "rowwise", "vocab"] + comm: Literal["none", "allreduce", "reduce_scatter"] + sequence_dim: int = 1 + use_local_output: bool = True + + def to_dtensor_style(self) -> ParallelStyle: + """Convert to the corresponding PyTorch DTensor ParallelStyle.""" + if self.kind == "colwise": + match self.comm: + case "none": + return ColwiseParallel( + input_layouts=Replicate(), output_layouts=Shard(-1), use_local_output=self.use_local_output + ) + elif self.kind == "rowwise": + match self.comm: + case "allreduce": + return RowwiseParallel( + input_layouts=Shard(-1), + output_layouts=Replicate(), + use_local_output=self.use_local_output, + ) + case "reduce_scatter": + return RowwiseParallel( + input_layouts=Shard(-1), + output_layouts=Shard(1), + use_local_output=self.use_local_output, + ) + elif self.kind == "vocab": + match self.comm: + case "allreduce": + return RowwiseParallel( + input_layouts=Replicate(), + output_layouts=Replicate(), + use_local_output=self.use_local_output, + ) + case "reduce_scatter": + return RowwiseParallel( + input_layouts=Replicate(), output_layouts=Shard(1), use_local_output=self.use_local_output + ) + raise ValueError( + f"Invalid TPStyle({self.kind!r}, {self.comm!r}). Valid combinations:\n" + f" colwise: none\n" + f" rowwise: allreduce, reduce_scatter\n" + f" vocab: allreduce, reduce_scatter" + ) + + def __str__(self): + if self.comm == "none": + return self.kind + return f"{self.kind}_{self.comm}" + + +def apply_tensor_parallel(model, tp_mesh, tp_plan): + """Apply tensor parallelism using PyTorch's parallelize_module. + + Converts the wildcard tp_plan from model config into a concrete plan + for ``parallelize_module``. Plan values are ``TPStyle`` instances. + """ + if tp_plan is None: + return model + + if tp_plan == "auto": + base_plan = model.config.base_model_tp_plan or {} + + # Prefix base model keys (e.g. "layers.*.q_proj" → "model.layers.*.q_proj") + # Top-level keys like "lm_head" are kept as-is. + base_model_prefix = model.base_model_prefix + tp_plan = {} + for k, v in base_plan.items(): + is_top_level = hasattr(model, k.split(".")[0]) + tp_plan[k if is_top_level else f"{base_model_prefix}.{k}"] = v + + parallelize_plan = {} + + for name, _ in model.named_modules(): + style_value = _get_parameter_tp_plan(parameter_name=name, tp_plan=tp_plan, is_weight=False) + if style_value is None: + continue + + if isinstance(style_value, TPStyle): + dtensor_style = style_value.to_dtensor_style() + parallelize_plan[name] = dtensor_style + else: + parallelize_plan[name] = style_value + + parallelize_module(model, tp_mesh, parallelize_plan) - model = apply_fsdp2(model, device_mesh, fsdp_plan) return model diff --git a/src/transformers/models/glm/configuration_glm.py b/src/transformers/models/glm/configuration_glm.py index 98525012a23b..d9d07638f6cb 100644 --- a/src/transformers/models/glm/configuration_glm.py +++ b/src/transformers/models/glm/configuration_glm.py @@ -17,6 +17,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -40,12 +41,12 @@ class GlmConfig(PreTrainedConfig): model_type = "glm" 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.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.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_up_proj": TPStyle("packed_colwise", "none"), # fused gate/up shards stay local for chunk + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/llama/configuration_llama.py b/src/transformers/models/llama/configuration_llama.py index 6960a6970592..f0d0d8a97197 100644 --- a/src/transformers/models/llama/configuration_llama.py +++ b/src/transformers/models/llama/configuration_llama.py @@ -21,6 +21,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring from ...utils.type_validators import interval @@ -47,13 +48,13 @@ class LlamaConfig(PreTrainedConfig): 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.o_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/llama/modeling_llama.py b/src/transformers/models/llama/modeling_llama.py index 9d659c7c6f08..3a966dc30798 100644 --- a/src/transformers/models/llama/modeling_llama.py +++ b/src/transformers/models/llama/modeling_llama.py @@ -21,6 +21,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache @@ -163,6 +164,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/mistral/configuration_mistral.py b/src/transformers/models/mistral/configuration_mistral.py index c57193f58d7b..2b64d139f145 100644 --- a/src/transformers/models/mistral/configuration_mistral.py +++ b/src/transformers/models/mistral/configuration_mistral.py @@ -16,6 +16,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring, logging @@ -46,13 +47,13 @@ class MistralConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `MistralModel` 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", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/mistral/modeling_mistral.py b/src/transformers/models/mistral/modeling_mistral.py index b79dea36c9e9..3d992a8e8aa3 100644 --- a/src/transformers/models/mistral/modeling_mistral.py +++ b/src/transformers/models/mistral/modeling_mistral.py @@ -9,6 +9,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache @@ -76,6 +77,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/phi/configuration_phi.py b/src/transformers/models/phi/configuration_phi.py index 2a65f0f16aec..c04e91344b45 100644 --- a/src/transformers/models/phi/configuration_phi.py +++ b/src/transformers/models/phi/configuration_phi.py @@ -17,6 +17,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -46,12 +47,12 @@ class PhiConfig(PreTrainedConfig): model_type = "phi" 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.dense": "rowwise", - "layers.*.mlp.fc1": "colwise", - "layers.*.mlp.fc2": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.dense": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.fc1": TPStyle("colwise", "none"), + "layers.*.mlp.fc2": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -66,8 +67,8 @@ class PhiConfig(PreTrainedConfig): 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/qwen2/configuration_qwen2.py b/src/transformers/models/qwen2/configuration_qwen2.py index ae41af7b211f..c48a7639fc5e 100644 --- a/src/transformers/models/qwen2/configuration_qwen2.py +++ b/src/transformers/models/qwen2/configuration_qwen2.py @@ -16,6 +16,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -44,13 +45,13 @@ class Qwen2Config(PreTrainedConfig): # Default tensor parallel plan for base model `Qwen2` 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", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/qwen3/modeling_qwen3.py b/src/transformers/models/qwen3/modeling_qwen3.py index 91715a33cf9d..30e7f5d2d9e1 100644 --- a/src/transformers/models/qwen3/modeling_qwen3.py +++ b/src/transformers/models/qwen3/modeling_qwen3.py @@ -23,6 +23,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache @@ -176,6 +177,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed @@ -441,7 +446,6 @@ 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"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): @@ -452,7 +456,6 @@ def __init__(self, config): # Initialize weights and apply final processing self.post_init() - @can_return_tuple @auto_docstring def forward( diff --git a/tests/tensor_parallel/test_tensor_parallel.py b/tests/tensor_parallel/test_tensor_parallel.py index 91770b683e45..b2e19d91bd72 100644 --- a/tests/tensor_parallel/test_tensor_parallel.py +++ b/tests/tensor_parallel/test_tensor_parallel.py @@ -25,42 +25,10 @@ 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): diff --git a/tests/test_distributed_config.py b/tests/test_distributed_config.py new file mode 100644 index 000000000000..53e9a063c123 --- /dev/null +++ b/tests/test_distributed_config.py @@ -0,0 +1,86 @@ +import json +import tempfile + +from transformers.distributed import DistributedConfig + + +class TestDistributedConfig: + def test_2d_parallelism(self): + dc = DistributedConfig(tp_size=2, tp_plan="auto", fsdp_size=2, fsdp_plan="auto") + assert dc.tp_size == 2 + assert dc.fsdp_size == 2 + assert dc.tp_plan == "auto" + assert dc.fsdp_plan == "auto" + + 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 == "auto" # size given → plan defaults to "auto" + + def test_fsdp_only_defaults_tp_to_1(self): + dc = DistributedConfig(fsdp_size=4) + assert dc.tp_size == 1 + assert dc.fsdp_size == 4 + assert dc.fsdp_plan == "auto" # size given → plan defaults to "auto" + assert dc.tp_plan == "auto" # tp_size got set to 1 → plan also defaults + + def test_empty_config(self): + dc = DistributedConfig() + assert dc.tp_size is None + assert dc.fsdp_size is None + assert dc.tp_plan is None + assert dc.fsdp_plan is None + + def test_from_dict(self): + dc = DistributedConfig.from_dict({"tp_size": 2, "fsdp_size": 4, "tp_plan": "auto"}) + assert dc.tp_size == 2 + assert dc.fsdp_size == 4 + assert dc.tp_plan == "auto" + + 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": "auto", "fsdp_size": 4, "fsdp_plan": "auto"} + + 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, tp_plan="auto") + 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"] == "auto" + + def test_roundtrip_dict(self): + original = DistributedConfig(tp_size=2, tp_plan="auto", fsdp_size=4, fsdp_plan="auto") + 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_tensor_parallel_mixin.py b/tests/test_tensor_parallel_mixin.py index 6f142b513b41..c138be86ad68 100644 --- a/tests/test_tensor_parallel_mixin.py +++ b/tests/test_tensor_parallel_mixin.py @@ -16,6 +16,7 @@ from abc import ABC, abstractmethod from transformers import TorchAoConfig, set_seed +from transformers.distributed import DistributedConfig from transformers.integrations.tensor_parallel import _get_parameter_tp_plan from transformers.testing_utils import ( is_tensor_parallel_test, @@ -32,8 +33,24 @@ import torch import torch.distributed as dist import torch.multiprocessing as mp + from torch.distributed.tensor import DTensor, Replicate 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 = tensor.redistribute(placements=(Replicate(),)) + return tensor.to_local() + return tensor + def _find_free_port(): """Find a free port by binding a socket and releasing it.""" @@ -68,6 +85,10 @@ def get_packed_grad_shard(grad, world_size, rank, dim): return grad.index_select(dim, torch.tensor(indices, device=grad.device)) +def _is_packed_colwise_plan(plan) -> bool: + return plan == "packed_colwise" or getattr(plan, "kind", None) == "packed_colwise" + + def _global_wrapper(rank, func, tp, port, func_args, func_kwargs): """Wrapper to set up distributed environment and run the test function.""" @@ -111,50 +132,69 @@ def wrapper(*args, **kwargs): return _init_distributed_inner -def _load_tp_and_reference_models(model_path, model_class): +def _load_tp_and_reference_models(model_path, model_class, enable_sequence_parallel=False): """Load TP model and non-TP reference model for comparison. Returns: tuple: (model_tp, model_ref, device) """ - model_tp = model_class.from_pretrained(model_path, tp_plan="auto") + tp_size = dist.get_world_size() + distributed_config = DistributedConfig(tp_size=tp_size, tp_plan="auto", enable_sequence_parallel=enable_sequence_parallel) + model_tp = model_class.from_pretrained( + model_path, distributed_config=distributed_config, attn_implementation="sdpa" + ) dist.barrier() device = model_tp.device - model_ref = model_class.from_pretrained(model_path) + model_ref = model_class.from_pretrained(model_path, attn_implementation="sdpa") model_ref = model_ref.to(device) return model_tp, model_ref, device +def _get_active_tp_plan(model_tp): + distributed_config = getattr(model_tp.config, "distributed_config", None) + tp_plan = getattr(distributed_config, "tp_plan", None) + + if tp_plan == "auto": + return getattr(model_tp, "_tp_plan", None) or {} + + return tp_plan or getattr(model_tp, "_tp_plan", None) or {} + + def _verify_tp_sharding(rank, model_tp, model_ref): """Verify TP sharding by comparing parameter shapes between TP and reference models. + For DTensor params, uses the local tensor shape (not the global DTensor shape). + Returns: list: Names of sharded parameters """ world_size = dist.get_world_size() sharded_params = [] + tp_plan = _get_active_tp_plan(model_tp) for (name, param), (_, param_full) in zip(model_tp.named_parameters(), model_ref.named_parameters()): - if param.shape != param_full.shape: + # For DTensor params, get the local shape for comparison + param_local = _to_local(param) + if param_local.shape != param_full.shape: sharded_params.append(name) if rank == 0: - print(f"[TP Test Debug] TP sharded: {name} - full: {param_full.shape} -> sharded: {param.shape}") + print(f"[TP Test Debug] TP sharded: {name} - full: {param_full.shape} -> sharded: {param_local.shape}") # Verify sharding is correct - for dim in range(param.ndim): - if param.size(dim) != param_full.size(dim): - param_plan = _get_parameter_tp_plan(name, model_tp.tp_plan, is_weight=True) - if param_plan == "packed_colwise": + for dim in range(param_local.ndim): + if param_local.size(dim) != param_full.size(dim): + param_plan = _get_parameter_tp_plan(name, tp_plan, is_weight=True) + if _is_packed_colwise_plan(param_plan): expected_size = param_full.size(dim) // world_size - assert param.size(dim) == expected_size, ( - f"Packed weight {name} sharding incorrect: expected {expected_size}, got {param.size(dim)}" + assert param_local.size(dim) == expected_size, ( + f"Packed weight {name} sharding incorrect: expected {expected_size}, got {param_local.size(dim)}" ) else: expected_size = (param_full.size(dim) + world_size - 1) // world_size - assert param.size(dim) <= expected_size, ( - f"Weight {name} sharding incorrect: expected <= {expected_size}, got {param.size(dim)}" + assert param_local.size(dim) <= expected_size, ( + f"Weight {name} sharding incorrect: expected <= {expected_size}, got {param_local.size(dim)}" ) break @@ -165,7 +205,7 @@ def _test_tp_forward_impl(_rank, model_path, model_class, atol, rtol): """Implementation for comparing TP and non-TP model outputs.""" set_seed(0) - model_tp, model, device = _load_tp_and_reference_models(model_path, model_class) + model_tp, model, device = _load_tp_and_reference_models(model_path, model_class, enable_sequence_parallel=True) _verify_tp_sharding(_rank, model_tp, model) @@ -178,7 +218,7 @@ def _test_tp_forward_impl(_rank, model_path, model_class, atol, rtol): with torch.no_grad(): logits = model(input_ids).logits - logits_tp = model_tp(input_ids).logits + logits_tp = _to_local(model_tp(input_ids).logits) diff = (logits - logits_tp).abs() assert torch.allclose(logits, logits_tp, atol=atol, rtol=rtol), ( @@ -192,7 +232,8 @@ def _test_tp_backward_impl(rank, model_path, model_class, atol, rtol): """Implementation for comparing TP and non-TP model backward passes.""" set_seed(0) - model_tp, model, device = _load_tp_and_reference_models(model_path, model_class) + model_tp, model, device = _load_tp_and_reference_models(model_path, model_class, enable_sequence_parallel=True) + tp_plan = _get_active_tp_plan(model_tp) model_tp.train() model.train() @@ -202,32 +243,35 @@ def _test_tp_backward_impl(rank, model_path, model_class, atol, rtol): set_seed(0) labels = torch.randint(0, vocab_size, (2, 64)).to(device) - loss = model(input_ids, labels=labels).loss + loss = model(input_ids, labels=labels, use_cache=False).loss loss.backward() - loss_tp = model_tp(input_ids, labels=labels).loss + loss_tp = model_tp(input_ids, labels=labels, use_cache=False).loss loss_tp.backward() - assert torch.allclose(loss, loss_tp, atol=atol, rtol=rtol), ( + loss_tp_local = _to_local(loss_tp) + assert torch.allclose(loss, loss_tp_local, atol=atol, rtol=rtol), ( f"TP and non-TP model losses differ. " - f"Non-TP loss: {loss.item()}, TP loss: {loss_tp.item()}, " - f"Diff: {(loss - loss_tp).abs().item()}" + f"Non-TP loss: {loss.item()}, TP loss: {loss_tp_local.item()}, " + f"Diff: {(loss - loss_tp_local).abs().item()}" ) # Compare gradients for matching parameters world_size = dist.get_world_size() + + # Debug: check tied weights and parameter alignment failed_grads = {} - for (name, param), (_, param_tp) in zip(model.named_parameters(), model_tp.named_parameters()): + for (name, param), (name_tp, param_tp) in zip(model.named_parameters(), model_tp.named_parameters()): if param.grad is not None and param_tp.grad is not None: grad = param.grad - grad_tp = param_tp.grad + grad_tp = _to_local(param_tp.grad) # Slice reference gradient to match local shard if parameter is sharded if grad.shape != grad_tp.shape: for dim in range(grad.ndim): if grad.size(dim) != grad_tp.size(dim): - param_plan = _get_parameter_tp_plan(name, model_tp.tp_plan, is_weight=True) - if param_plan == "packed_colwise": + param_plan = _get_parameter_tp_plan(name, tp_plan, is_weight=True) + if _is_packed_colwise_plan(param_plan): # interleaved slicing grad = get_packed_grad_shard(grad, world_size, rank, dim) else: @@ -238,10 +282,15 @@ def _test_tp_backward_impl(rank, model_path, model_class, atol, rtol): break if not torch.allclose(grad.cpu(), grad_tp.cpu(), atol=atol, rtol=rtol): - failed_grads[name] = (grad.cpu() - grad_tp.cpu()).abs().max().item() + max_diff = (grad.cpu() - grad_tp.cpu()).abs().max().item() + ref_abs_max = grad.cpu().abs().max().item() + tp_abs_max = grad_tp.cpu().abs().max().item() + ratio = tp_abs_max / ref_abs_max if ref_abs_max > 0 else float("inf") + failed_grads[name] = (max_diff, ref_abs_max, tp_abs_max, ratio) assert not failed_grads, f"Gradients differ for {len(failed_grads)} parameter(s):\n" + "\n".join( - f" {name}: max diff = {diff}" for name, diff in failed_grads.items() + f" {name}: max_diff={v[0]:.6f}, ref_max={v[1]:.6f}, tp_max={v[2]:.6f}, tp/ref ratio={v[3]:.4f}" + for name, v in failed_grads.items() ) dist.barrier() @@ -273,7 +322,7 @@ def _test_tp_generation_impl(_rank, model_path, model_class, atol, rtol, max_new # Compare logits/scores at each generation step scores = torch.stack(output.scores) - scores_tp = torch.stack(output_tp.scores) + scores_tp = torch.stack([_to_local(s) for s in output_tp.scores]) diff = (scores - scores_tp).abs() assert torch.allclose(scores, scores_tp, atol=atol, rtol=rtol), ( @@ -282,9 +331,10 @@ def _test_tp_generation_impl(_rank, model_path, model_class, atol, rtol, max_new ) # Compare generated token sequences - assert torch.equal(output.sequences, output_tp.sequences), ( + sequences_tp = _to_local(output_tp.sequences) + assert torch.equal(output.sequences, sequences_tp), ( f"TP and non-TP model generated different token sequences (direct load path). " - f"Non-TP: {output.sequences.tolist()} | TP: {output_tp.sequences.tolist()}" + f"Non-TP: {output.sequences.tolist()} | TP: {sequences_tp.tolist()}" ) dist.barrier() @@ -296,7 +346,9 @@ 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_plan="auto"), quantization_config=quantization_config + ) dist.barrier() device = model_tp.device @@ -359,8 +411,8 @@ class TensorParallelTesterMixin(ABC): # Configuration (can be overridden per model) # ============================================================ tensor_parallel_size: int = 2 - tensor_parallel_atol: float = 1e-5 - tensor_parallel_rtol: float = 1e-5 + tensor_parallel_atol: float = 5e-3 + tensor_parallel_rtol: float = 5e-3 @property @abstractmethod @@ -470,15 +522,4 @@ def test_tp_generation_quantized(self): if not is_torchao_available(): self.skipTest("Test requires torchao") - config = self.model_tester.get_config() - model_class = self._get_tp_model_class() - max_new_tokens = 25 - - with tempfile.TemporaryDirectory() as tmp_dir: - set_seed(42) - model = model_class(config) - model.save_pretrained(tmp_dir, save_original_format=True) - - _init_distributed(tp=self.tensor_parallel_size)(_test_tp_generation_quantized_impl)( - tmp_dir, model_class, max_new_tokens - ) + self.skipTest("Quantization is not currently supported with distributed training") From abfd57eeca2097931fa3a9ad201565aa7dc0e699 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Mon, 13 Apr 2026 15:00:06 +0000 Subject: [PATCH 005/116] revert some files --- .../modeling_flash_attention_utils.py | 19 ++-- tests/test_modeling_common.py | 90 ++----------------- 2 files changed, 15 insertions(+), 94 deletions(-) diff --git a/src/transformers/modeling_flash_attention_utils.py b/src/transformers/modeling_flash_attention_utils.py index 0454b6dfefae..9211ccb19a9e 100644 --- a/src/transformers/modeling_flash_attention_utils.py +++ b/src/transformers/modeling_flash_attention_utils.py @@ -74,10 +74,8 @@ def is_flash_attn_available(): 2: { "flash_attn_version": 2, "general_availability_check": is_flash_attn_2_available, - "pkg_availability_check": lambda *args, **kwargs: ( - importlib.util.find_spec("flash_attn") is not None - and "flash-attn" in [pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn"]] - ), + "pkg_availability_check": lambda *args, **kwargs: importlib.util.find_spec("flash_attn") is not None + and "flash-attn" in [pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn"]], "supported_devices": ( (is_torch_cuda_available, "cuda"), (is_torch_mlu_available, "mlu"), @@ -95,21 +93,16 @@ def is_flash_attn_available(): 3: { "flash_attn_version": 3, "general_availability_check": is_flash_attn_3_available, - "pkg_availability_check": lambda *args, **kwargs: ( - importlib.util.find_spec("flash_attn_interface") is not None - and "flash-attn-3" - in [pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn_interface"]] - ), + "pkg_availability_check": lambda *args, **kwargs: importlib.util.find_spec("flash_attn_interface") is not None + and "flash-attn-3" in [pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn_interface"]], "supported_devices": ((is_torch_cuda_available, "cuda"),), "cuda_min_major_version": 8, # Ampere }, 4: { "flash_attn_version": 4, "general_availability_check": is_flash_attn_4_available, - "pkg_availability_check": lambda *args, **kwargs: ( - importlib.util.find_spec("flash_attn") is not None - and "flash-attn-4" in [pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn"]] - ), + "pkg_availability_check": lambda *args, **kwargs: importlib.util.find_spec("flash_attn") is not None + and "flash-attn-4" in [pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn"]], "supported_devices": ((is_torch_cuda_available, "cuda"),), "cuda_min_major_version": 9, # Hopper }, diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py index f9b3e48a55fd..13b81855aaa6 100755 --- a/tests/test_modeling_common.py +++ b/tests/test_modeling_common.py @@ -434,10 +434,6 @@ def _can_output_attn(model): outputs_eager = outputs_eager["language_model_outputs"] outputs_sdpa = outputs_sdpa["language_model_outputs"] key = "hidden_states" if "hidden_states" in outputs_eager else "decoder_hidden_states" - elif "decoder_output" in outputs_eager and "clipseg" in model_class.__name__.lower(): - outputs_eager = outputs_eager["decoder_output"] - outputs_sdpa = outputs_sdpa["decoder_output"] - key = "hidden_states" if "hidden_states" in outputs_eager else "decoder_hidden_states" else: key = "hidden_states" @@ -1305,7 +1301,7 @@ def test_init_weights_can_init_buffers(self): config.scale = 0 for sub_key in config.sub_configs: subconfig = getattr(config, sub_key) - if subconfig is not None and hasattr(subconfig, "scale"): + if hasattr(subconfig, "scale"): subconfig.scale = 0 for model_class in self.all_model_classes: @@ -1754,10 +1750,7 @@ def _set_subconfig_attributes(self, config, attribute_name, value): """Helper function to recursively set a config attr to a given value""" for k in config.sub_configs: if ( - self._is_composite - and attribute_name == "output_attentions" - and k == "vision_config" - and "Timm" in getattr(config, k).__class__.__name__ + self._is_composite and attribute_name == "output_attentions" and k == "vision_config" ): # skip because it's not needed and causes errors e.g with Timm continue if getattr(config, k) is not None: @@ -2405,55 +2398,6 @@ def test_resize_embeddings_untied_with_deepspeed_multi_gpu(self): with _deepspeed_zero3(ds_config): self.test_resize_embeddings_untied() - def test_resize_embeddings_untied_no_reinit_on_post_init(self): - if not self.test_resize_embeddings: - self.skipTest(reason="test_resize_embeddings is set to `False`") - - original_config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() - original_config.tie_word_embeddings = False - try: - original_config.get_text_config().tie_word_embeddings = False - except Exception as e: - model_type = getattr(original_config, "model_type", "unknown") - print(f"Could not set text config's `tie_word_embeddings` for model type `{model_type}`: {e}") - - if original_config.tie_word_embeddings: - self.skipTest(reason="Model cannot untie embeddings") - - for model_class in self.all_model_classes: - with self.subTest(model_class): - config = copy.deepcopy(original_config) - model = model_class(config).to(torch_device) - model.eval() - - # The bug only affects nn.Linear LM heads created by _get_resized_lm_head - output_embeds = model.get_output_embeddings() - if not isinstance(output_embeds, nn.Linear): - continue - - model_vocab_size = config.get_text_config().vocab_size - try: - model.resize_token_embeddings(model_vocab_size + 10) - except (NotImplementedError, AttributeError): - continue - - output_embeds = model.get_output_embeddings() - weights_before = output_embeds.weight.data.clone() - bias_before = output_embeds.bias.data.clone() if output_embeds.bias is not None else None - - model.post_init() - - output_embeds_after = model.get_output_embeddings() - self.assertTrue( - torch.equal(weights_before, output_embeds_after.weight.data), - "Output embedding weights were reinitialized by post_init() after resize_token_embeddings()", - ) - if bias_before is not None: - self.assertTrue( - torch.equal(bias_before, output_embeds_after.bias.data), - "Output embedding bias was reinitialized by post_init() after resize_token_embeddings()", - ) - def test_model_get_set_embeddings(self): config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() @@ -2533,10 +2477,8 @@ def test_can_use_safetensors(self): torch.testing.assert_close( v, reloaded_state[k], - msg=lambda x: ( - f"{model_class.__name__}: Tensor {k}: {x}.\n{v}\nvs\n{reloaded_state[k]}\n" - "This probably means that it was not set with the correct value when tying." - ), + msg=lambda x: f"{model_class.__name__}: Tensor {k}: {x}.\n{v}\nvs\n{reloaded_state[k]}\n" + "This probably means that it was not set with the correct value when tying.", ) # Checking the tensor sharing are correct on the new model (weights are properly tied in both cases) @@ -2582,9 +2524,7 @@ def test_load_save_without_tied_weights(self): torch.testing.assert_close( v, reloaded_state[k], - msg=lambda x: ( - f"{model_class.__name__}: Tensor {k}: {x}. Key {k} was serialized: {k in serialized_keys}. If `False`, this means it was probably aliased and safetensors removed it. If `True` it means `_init_weights` overwrote that key" - ), + msg=lambda x: f"{model_class.__name__}: Tensor {k}: {x}. Key {k} was serialized: {k in serialized_keys}. If `False`, this means it was probably aliased and safetensors removed it. If `True` it means `_init_weights` overwrote that key", ) # Checking there was no complain of missing weights @@ -3672,7 +3612,6 @@ def test_sdpa_can_dispatch_on_flash(self): "PaliGemma-like models currently (transformers==4.41.0) requires an attention_mask input" ) if config.model_type in [ - "evolla", "modernbert", "gemma3", "t5gemma", @@ -3684,9 +3623,6 @@ def test_sdpa_can_dispatch_on_flash(self): "kosmos-2", "mllama", "lighton_ocr", - "parakeet_encoder", - "parakeet_ctc", - "pi0", "pixtral", "sam", "sam_hq", @@ -4747,7 +4683,7 @@ def test_tp_plan_matches_params(self): len(unused_entries) == 0, f"The following entries of the TP-plan are not valid: {unused_entries}" ) - def test_reverse_loading_mapping(self, check_keys_were_modified=True, skip_base_model=False): + def test_reverse_loading_mapping(self, check_keys_were_modified=True): """Make sure we can load and save correctly the models having any weight renaming mapping or weight conversion mapping. Note that this test would be better if we could start from the serialized keys, and check that the model @@ -4761,11 +4697,6 @@ def test_reverse_loading_mapping(self, check_keys_were_modified=True, skip_base_ check_keys_were_modified (`bool`, *optional*, defaults to `True`): Whether to expect keys being modified or not. In some cases, models do not change keys but their weights, e.g. via transpose, memory alignment, etc. - skip_base_model (`bool`, *optional*, defaults to `False`): - Sometimes, mappings are only visible when applied to the model with head, and not visible on the - base model. This allows to skip the check on the base model. See e.g. `llava` mapping where this - is the case. In practice, the mappings are still coherent and a base model can still be loaded from - the head model, thanks to the `base_model_prefix` which will remove the prefix automatically. """ config, _ = self.model_tester.prepare_config_and_inputs_for_common() @@ -4778,15 +4709,13 @@ def test_reverse_loading_mapping(self, check_keys_were_modified=True, skip_base_ config_to_set.num_dense_layers = 1 # lfm2_moe for model_class in self.all_model_classes: - if skip_base_model and "For" not in model_class.__name__: - continue # Each individual model is a subtest with self.subTest(model_class.__name__): model = model_class(copy.deepcopy(config)) # Skip if no conversions conversions = get_model_conversion_mapping(model, add_legacy=False) if len(conversions) == 0: - self.skipTest(f"No conversion found for {model_class}") + self.skipTest("No conversion found for this model") # Find the model keys, so the targets according to the conversions model_keys = list(model.state_dict().keys()) @@ -4825,7 +4754,7 @@ def test_reverse_loading_mapping(self, check_keys_were_modified=True, skip_base_ self.assertTrue( num_matches > 0, f"`{source_pattern}` in `{conversion}` did not match any of the source keys. " - "This indicates whether that the pattern is not properly written, or that it could not be reversed correctly", + "This indicates whether that the pattern is not properly written, ot that it could not be reversed correctly", ) # If everything is still good at this point, let's test that we perform the same operations both when @@ -4862,7 +4791,7 @@ def test_can_load_from_already_mapped_keys(self): # Skip if no conversions conversions = get_model_conversion_mapping(model, add_legacy=False) if len(conversions) == 0: - self.skipTest(f"No conversion found for {model_class}") + self.skipTest("No conversion found for this model") with tempfile.TemporaryDirectory() as tmpdirname: # Serialize without reverting the mapping @@ -4928,7 +4857,6 @@ def _audio_features_prepare_config_and_inputs(self): or "input_values" in key or "input_features" in key or key in ["padding_mask", "is_longer", "feature_attention_mask"] - or (config.model_type == "musicflamingo" and key == "input_ids") } return config, inputs_dict From c33873e71c3225d233a0415ad10aa2313e032627 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Mon, 13 Apr 2026 15:36:43 +0000 Subject: [PATCH 006/116] Add distributed training scripts - train_fsdp_tp.py: minimal FSDP+TP training example - train_fsdp_tp_torchtitan_style.py: torchtitan-style training example - verify_loading.py: save/load roundtrip verification - run_compare.sh: FSDP+TP vs FSDP-only comparison - run_verify_all.sh: run verification across all modes - tmp_generate.py: quick generation test --- run_compare.sh | 56 +++++++ run_verify_all.sh | 160 ++++++++++++++++++++ tmp_generate.py | 63 ++++++++ train_fsdp_tp.py | 125 ++++++++++++++++ train_fsdp_tp_torchtitan_style.py | 239 ++++++++++++++++++++++++++++++ verify_loading.py | 137 +++++++++++++++++ 6 files changed, 780 insertions(+) create mode 100644 run_compare.sh create mode 100644 run_verify_all.sh create mode 100644 tmp_generate.py create mode 100644 train_fsdp_tp.py create mode 100644 train_fsdp_tp_torchtitan_style.py create mode 100644 verify_loading.py diff --git a/run_compare.sh b/run_compare.sh new file mode 100644 index 000000000000..eb47e1841fa9 --- /dev/null +++ b/run_compare.sh @@ -0,0 +1,56 @@ +#!/bin/bash +set -euo pipefail + +SCRIPT="train_fsdp_tp.py" +LOG_FSDP_TP="log.txt" +LOG_FSDP_ONLY="ref.txt" + +MODEL_NAME="${MODEL_NAME:-hf-internal-testing/tiny-random-MixtralForCausalLM}" +COMMON_ARGS="--model_name $MODEL_NAME --lr 3e-4 --seed 42" + +rm -rf ./checkpoints_tp ./checkpoints_tp_resumed ./checkpoints_fsdp ./checkpoints_fsdp_resumed + +echo "=== Phase 1: Train steps 0-9, save checkpoint ===" +echo "--- Launching FSDP+TP and FSDP-only in parallel ---" + +CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nproc_per_node=4 --master_port=29500 \ + $SCRIPT $COMMON_ARGS --fsdp_size 2 --tp_size 2 --enable_sp \ + --num_steps 10 --save_dir ./checkpoints_tp > "${LOG_FSDP_TP}.phase1" 2>&1 & +PID1=$! + +CUDA_VISIBLE_DEVICES=4,5 torchrun --nproc_per_node=2 --master_port=29501 \ + $SCRIPT $COMMON_ARGS --fsdp_size 2 \ + --num_steps 10 --save_dir ./checkpoints_fsdp > "${LOG_FSDP_ONLY}.phase1" 2>&1 & +PID2=$! + +echo "FSDP+TP PID=$PID1 | FSDP-only PID=$PID2" +wait $PID1 && echo "Phase 1 FSDP+TP done" || { echo "Phase 1 FSDP+TP failed (exit $?)"; cat "${LOG_FSDP_TP}.phase1"; exit 1; } +wait $PID2 && echo "Phase 1 FSDP-only done" || { echo "Phase 1 FSDP-only failed (exit $?)"; cat "${LOG_FSDP_ONLY}.phase1"; exit 1; } + +echo "" +echo "=== Phase 2: Resume from checkpoint, train steps 10-19, save ===" +echo "--- Launching FSDP+TP and FSDP-only in parallel ---" + +CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nproc_per_node=4 --master_port=29500 \ + $SCRIPT $COMMON_ARGS --fsdp_size 2 --tp_size 2 --enable_sp \ + --num_steps 10 --start_step 10 \ + --resume_dir ./checkpoints_tp --save_dir ./checkpoints_tp_resumed > "${LOG_FSDP_TP}.phase2" 2>&1 & +PID1=$! + +CUDA_VISIBLE_DEVICES=4,5 torchrun --nproc_per_node=2 --master_port=29501 \ + $SCRIPT $COMMON_ARGS --fsdp_size 2 \ + --num_steps 10 --start_step 10 \ + --resume_dir ./checkpoints_fsdp --save_dir ./checkpoints_fsdp_resumed > "${LOG_FSDP_ONLY}.phase2" 2>&1 & +PID2=$! + +echo "FSDP+TP PID=$PID1 | FSDP-only PID=$PID2" +wait $PID1 && echo "Phase 2 FSDP+TP done" || { echo "Phase 2 FSDP+TP failed (exit $?)"; cat "${LOG_FSDP_TP}.phase2"; exit 1; } +wait $PID2 && echo "Phase 2 FSDP-only done" || { echo "Phase 2 FSDP-only failed (exit $?)"; cat "${LOG_FSDP_ONLY}.phase2"; exit 1; } + +# Combine phase logs +cat "${LOG_FSDP_TP}.phase1" "${LOG_FSDP_TP}.phase2" > "$LOG_FSDP_TP" +cat "${LOG_FSDP_ONLY}.phase1" "${LOG_FSDP_ONLY}.phase2" > "$LOG_FSDP_ONLY" + +echo "" +echo "=== Full Loss & Grad Diff (steps 0-19) ===" +git diff --no-index --color --word-diff=color "$LOG_FSDP_TP" "$LOG_FSDP_ONLY" || true diff --git a/run_verify_all.sh b/run_verify_all.sh new file mode 100644 index 000000000000..16aa3267fe9a --- /dev/null +++ b/run_verify_all.sh @@ -0,0 +1,160 @@ +#!/bin/bash + +GREEN='\033[0;32m' +RED='\033[0;31m' +CYAN='\033[0;36m' +YELLOW='\033[1;33m' +BOLD='\033[1m' +DIM='\033[0;90m' +NC='\033[0m' + +SCRIPT="verify_loading.py" +LOGDIR="$(dirname "$0")/verify_logs" +mkdir -p "$LOGDIR" + +NUM_GPUS=$(nvidia-smi -L | wc -l) + +# Job definitions: "mode nproc_per_node" +declare -a JOBS=( + "single_gpu 1" + "fsdp 2" + "tp 2" + "tp_sp 2" + "tp_fsdp 4" + "tp_sp_fsdp 4" +) +MODE_NAMES=(single_gpu fsdp tp tp_sp tp_fsdp tp_sp_fsdp) + +echo -e "${BOLD}==========================================" +echo -e " Verify Loading (${NUM_GPUS} GPUs available)" +echo -e " Modes: ${MODE_NAMES[*]}" +echo -e " Logs: $LOGDIR/" +echo -e "==========================================${NC}" +echo "" + +# ============================================================ +# Round-robin GPU scheduler +# ============================================================ +NEXT_GPU=0 +MASTER_PORT=29500 +PIDS=() +PID_MODES=() + +for job in "${JOBS[@]}"; do + mode=${job% *} + nproc=${job#* } + + # Wait if not enough GPUs left in this round + if [ $((NEXT_GPU + nproc)) -gt "$NUM_GPUS" ]; then + echo -e "${DIM} (waiting for current round to finish...)${NC}" + for pid in "${PIDS[@]}"; do + wait "$pid" 2>/dev/null + done + PIDS=() + NEXT_GPU=0 + fi + + # Build CUDA_VISIBLE_DEVICES range + GPU_END=$((NEXT_GPU + nproc - 1)) + GPUS="" + for g in $(seq "$NEXT_GPU" "$GPU_END"); do + [ -n "$GPUS" ] && GPUS="${GPUS}," + GPUS="${GPUS}${g}" + done + + echo -e " ${CYAN}[${mode}]${NC} GPUs ${NEXT_GPU}-${GPU_END} (nproc=${nproc})" + + if [ "$nproc" -eq 1 ]; then + CUDA_VISIBLE_DEVICES="$GPUS" python "$SCRIPT" --mode "$mode" \ + > "$LOGDIR/${mode}.log" 2>&1 & + else + CUDA_VISIBLE_DEVICES="$GPUS" torchrun \ + --nproc_per_node="$nproc" --master_port="$MASTER_PORT" \ + "$SCRIPT" --mode "$mode" \ + > "$LOGDIR/${mode}.log" 2>&1 & + ((MASTER_PORT++)) + fi + + PIDS+=($!) + PID_MODES+=("$mode") + NEXT_GPU=$((GPU_END + 1)) +done + +# Wait for remaining jobs +echo "" +echo -e "${BOLD}Waiting for all jobs to finish...${NC}" +for i in "${!PIDS[@]}"; do + mode="${PID_MODES[$i]}" + if wait "${PIDS[$i]}"; then + echo -e " ${GREEN}✓${NC} ${mode}" + else + echo -e " ${RED}✗${NC} ${mode} (exit $?)" + fi +done + +# ============================================================ +# Results +# ============================================================ +echo "" +echo -e "${BOLD}=== Results ===${NC}" +for mode in "${MODE_NAMES[@]}"; do + log="$LOGDIR/$mode.log" + loss_before=$(grep -oP 'loss_before = \K[0-9.]+' "$log" 2>/dev/null) + loss_after=$(grep -oP 'loss_after = \K[0-9.]+' "$log" 2>/dev/null) + if grep -q '^PASS' "$log" 2>/dev/null; then + printf " ${GREEN}%-12s PASS (before=%-10s after=%s)${NC}\n" "$mode" "$loss_before" "$loss_after" + elif [ -n "$loss_before" ]; then + diff=$(grep -oP 'diff = \K[0-9.e+-]+' "$log" 2>/dev/null) + printf " ${RED}%-12s FAIL (before=%-10s after=%-10s diff=%s)${NC}\n" "$mode" "$loss_before" "$loss_after" "$diff" + else + printf " ${RED}%-12s ERROR (see log)${NC}\n" "$mode" + fi +done + +# ============================================================ +# Cross-mode loss comparison +# ============================================================ +echo "" +echo -e "${BOLD}=== Cross-mode loss comparison (PASS modes only) ===${NC}" +REF_LOSS="" +ALL_MATCH=1 +for mode in "${MODE_NAMES[@]}"; do + log="$LOGDIR/$mode.log" + # Only include modes where save/load roundtrip passed + if ! grep -q '^PASS' "$log" 2>/dev/null; then + continue + fi + loss=$(grep -oP 'loss_before = \K[0-9.]+' "$log" 2>/dev/null) + if [ -z "$loss" ]; then + continue + fi + if [ -z "$REF_LOSS" ]; then + REF_LOSS="$loss" + printf " ${GREEN}%-12s %s (reference)${NC}\n" "$mode" "$loss" + elif [ "$loss" = "$REF_LOSS" ]; then + printf " ${GREEN}%-12s %s${NC}\n" "$mode" "$loss" + else + printf " ${YELLOW}%-12s %s (differs from %s)${NC}\n" "$mode" "$loss" "$REF_LOSS" + ALL_MATCH=0 + fi +done +if [ "$ALL_MATCH" -eq 1 ] && [ -n "$REF_LOSS" ]; then + echo -e " ${GREEN}All modes produce the same loss.${NC}" +fi + +# Hints for failures +HAS_FAIL=0 +for mode in "${MODE_NAMES[@]}"; do + if ! grep -q '^PASS' "$LOGDIR/$mode.log" 2>/dev/null; then + HAS_FAIL=1 + fi +done +if [ "$HAS_FAIL" -eq 1 ]; then + echo "" + echo -e "${YELLOW}Some modes failed. Check logs:${NC}" + for mode in "${MODE_NAMES[@]}"; do + if ! grep -q '^PASS' "$LOGDIR/$mode.log" 2>/dev/null; then + echo -e " ${YELLOW}cat $LOGDIR/$mode.log${NC}" + fi + done +fi diff --git a/tmp_generate.py b/tmp_generate.py new file mode 100644 index 000000000000..9685bed643ed --- /dev/null +++ b/tmp_generate.py @@ -0,0 +1,63 @@ +import argparse +import os + +import torch +from torch.distributed.elastic.multiprocessing.errors import record + +from transformers import AutoModelForCausalLM, AutoTokenizer +from transformers.distributed import DistributedConfig + +model_id = "mistralai/Mixtral-8x7B-Instruct-v0.1" +# model_id = "Qwen/Qwen3-14B" +# model_id = "Qwen/Qwen3-0.6B" +# model_id = "Qwen/Qwen1.5-MoE-A2.7B-Chat" +# model_id = "Qwen/Qwen3-30B-A3B-Instruct-2507" + +rank = int(os.environ["RANK"]) +world_size = int(os.environ["WORLD_SIZE"]) +device = torch.device(f"cuda:{rank}") +# Need to be initialized explicitly to use the `barrier` before loading +torch.distributed.init_process_group(backend="nccl", rank=rank, world_size=world_size, device_id=rank) + +@record +def main(args): + + distributed_config = DistributedConfig(tp_size=4, tp_plan="auto") + model = AutoModelForCausalLM.from_pretrained(model_id, distributed_config=distributed_config, dtype=torch.bfloat16) + # model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.bfloat16, device_map="auto") + tokenizer = AutoTokenizer.from_pretrained(model_id) + + messages = [ + {"role": "user", "content": "What do you think about life?"}, + ] + inputs = tokenizer.apply_chat_template(messages, return_tensors="pt").to(model.device) + input_size = inputs.input_ids.shape[-1] + + if args.profile: + # Warmup + with torch.no_grad(): + _ = model.generate(**inputs, max_new_tokens=5, do_sample=False) + + with torch.profiler.profile( + activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA], + record_shapes=True, + ) as prof: + output = model.generate(**inputs, max_new_tokens=2, do_sample=False) + + if rank == 0: + print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=30)) + prof.export_chrome_trace("trace.json") + else: + output = model.generate(**inputs, max_new_tokens=100, do_sample=False) + + text = tokenizer.batch_decode(output[:, input_size:])[0] + if rank == 0: + print(text) + +parser = argparse.ArgumentParser() +parser.add_argument("--profile", action="store_true") +args = parser.parse_args() + +main(args) + +torch.distributed.destroy_process_group() \ No newline at end of file diff --git a/train_fsdp_tp.py b/train_fsdp_tp.py new file mode 100644 index 000000000000..0232f8b3bc3d --- /dev/null +++ b/train_fsdp_tp.py @@ -0,0 +1,125 @@ +# torchrun --nproc_per_node=4 train_fsdp_tp.py + +import argparse +import os + +import torch +from datasets import load_dataset +from torch.distributed.tensor import DTensor +from torch.utils.data import DataLoader +from transformers import AutoModelForCausalLM, AutoTokenizer +from transformers.distributed import DistributedConfig +from transformers.distributed.utils import load_optimizer, save_optimizer + +def build_packed_dataset(dataset_name, tokenizer, seq_len, dp_rank, dp_world_size): + """Stream + tokenize + greedy-pack documents into fixed-length (input, label) windows.""" + ds = load_dataset(dataset_name, name="en", split="train", streaming=True) + ds = ds.shard(num_shards=dp_world_size, index=dp_rank) + buf, w = [], seq_len + 1 + + def pack(batch): + for t in batch["text"]: + buf.extend(tokenizer(t)["input_ids"]) + ids, lbls = [], [] + while len(buf) >= w: + ids.append(buf[:seq_len]); lbls.append(buf[1:w]); del buf[:w] + return {"input_ids": ids, "labels": lbls} + + ds = ds.map(pack, batched=True, remove_columns=ds.column_names) + return ds.with_format("torch") + +def build_fixed_batches(dp_rank): + """Load pre-generated fixed batches for a given DP rank.""" + return torch.load(f"fixed_batches_dp{dp_rank}.pt", weights_only=True) + +if __name__ == "__main__": + + parser = argparse.ArgumentParser() + parser.add_argument("--model_name", type=str, default="Qwen/Qwen3-0.6B") + parser.add_argument("--num_steps", type=int, default=20) + parser.add_argument("--lr", type=float, default=3e-4) + parser.add_argument("--seq_len", type=int, default=512) + parser.add_argument("--batch_size", type=int, default=1) + parser.add_argument("--save_dir", type=str, default="./checkpoints") + parser.add_argument("--tp_size", type=int, default=0, help="Tensor parallel size (0 = disabled)") + parser.add_argument("--fsdp_size", type=int, default=0, help="FSDP size (0 = disabled)") + parser.add_argument("--enable_sp", action="store_true", help="Enable sequence parallelism") + parser.add_argument("--seed", type=int, default=42, help="Random seed") + parser.add_argument("--fixed_batches", action="store_true", help="Use pre-generated fixed batches instead of C4") + parser.add_argument("--resume_dir", type=str, default=None, help="Resume from this checkpoint directory") + parser.add_argument("--start_step", type=int, default=0, help="Starting step number (for logging)") + args = parser.parse_args() + + torch.distributed.init_process_group(backend="nccl") + rank, local_rank = int(os.environ["RANK"]), int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + torch.manual_seed(args.seed) + + dc_kwargs = {} + if args.tp_size > 0: + dc_kwargs["tp_size"] = args.tp_size + dc_kwargs["tp_plan"] = "auto" + if args.fsdp_size > 0: + dc_kwargs["fsdp_size"] = args.fsdp_size + dc_kwargs["fsdp_plan"] = "auto" + if args.enable_sp: + dc_kwargs["enable_sequence_parallel"] = True + distributed_config = DistributedConfig(**dc_kwargs) + + load_path = args.resume_dir if args.resume_dir else args.model_name + model = AutoModelForCausalLM.from_pretrained( + load_path, + distributed_config=distributed_config, + torch_dtype=torch.bfloat16, + ) + + dp_rank = model.device_mesh["fsdp"].get_local_rank() if "fsdp" in model.device_mesh.mesh_dim_names else 0 + dp_size = model.device_mesh["fsdp"].size() if "fsdp" in model.device_mesh.mesh_dim_names else 1 + + if args.fixed_batches: + fixed = build_fixed_batches(dp_rank) + else: + tokenizer = AutoTokenizer.from_pretrained(args.model_name) + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + dataset = build_packed_dataset("allenai/c4", tokenizer, args.seq_len, dp_rank, dp_size) + dataloader = iter(DataLoader(dataset, batch_size=args.batch_size)) + + optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr) + + if args.resume_dir: + load_optimizer(optimizer, os.path.join(args.resume_dir, "optimizer")) + if rank == 0: + print(f"Resumed optimizer from {args.resume_dir}") + + model.train() + for step in range(args.start_step, args.start_step + args.num_steps): + if args.fixed_batches: + input_ids = fixed[step]["input_ids"].to(f"cuda:{local_rank}") + labels = fixed[step]["labels"].to(f"cuda:{local_rank}") + else: + batch = next(dataloader) + input_ids = batch["input_ids"].to(f"cuda:{local_rank}") + labels = batch["labels"].to(f"cuda:{local_rank}") + loss = model(input_ids, labels=labels).loss + loss.backward() + + # Custom grad clip: convert DTensor grads to local to avoid mixed-mesh torch.stack + grads = [p.grad for p in model.parameters() if p.grad is not None] + local_grads = [g.full_tensor() if isinstance(g, DTensor) else g for g in grads] + total_norm = torch.nn.utils.get_total_norm(local_grads, norm_type=2.0) + torch.nn.utils.clip_grads_with_norm_(grads, max_norm=1.0, total_norm=total_norm) + optimizer.step() + optimizer.zero_grad() + + if rank == 0: + print(f"Step {step:>4d} | Loss: {loss.item():.4f} | Grad norm: {total_norm.item():.4f}") + + # Save model (HF format) and optimizer (DCP) + model.save_pretrained(args.save_dir) + save_optimizer(optimizer, os.path.join(args.save_dir, "optimizer")) + + if rank == 0: + print(f"Saved to {args.save_dir}") + + torch.distributed.destroy_process_group() diff --git a/train_fsdp_tp_torchtitan_style.py b/train_fsdp_tp_torchtitan_style.py new file mode 100644 index 000000000000..325ee112c778 --- /dev/null +++ b/train_fsdp_tp_torchtitan_style.py @@ -0,0 +1,239 @@ +# torchrun --nproc_per_node=4 train_fsdp_tp_torchtitan_style.py +# LOAD_PRETRAINED=1 torchrun --nproc_per_node=4 train_fsdp_tp_torchtitan_style.py +# +# Minimal standalone training script that reuses torchtitan's components +# (model wrapper, parallelization, loss, optimizer, grad clipping) directly. +# This is the same code path as `./run_train.sh` but without the config system. + +import os + +import torch +import torch.distributed as dist +import torch.distributed.checkpoint as dcp +import torch.nn.functional as F +from huggingface_hub import snapshot_download +from torch.distributed.checkpoint import HuggingFaceStorageReader + +# ---------- torchtitan imports ---------- +from torchtitan.distributed import ParallelDims +from torchtitan.distributed import utils as dist_utils +from torchtitan.experiments.transformers_modeling_backend.infra.parallelize import ( + apply_fsdp, + apply_non_moe_tp, + disable_fsdp_gradient_division, +) +from torchtitan.experiments.transformers_modeling_backend.model.args import ( + HFTransformerModelArgs, + TitanDenseModelArgs, +) +from torchtitan.experiments.transformers_modeling_backend.model.model import ( + HFTransformerModel, +) + +# ---------- transformers imports ---------- +from transformers import AutoConfig, AutoTokenizer + +IGNORE_INDEX = -100 + + +def build_model_args(hf_model_name: str, seq_len: int) -> HFTransformerModelArgs: + """Build HFTransformerModelArgs from a HuggingFace model name.""" + hf_config = AutoConfig.from_pretrained( + hf_model_name, attn_implementation="sdpa", trust_remote_code=True + ) + hf_config_dict = hf_config.to_dict() + + model_args = HFTransformerModelArgs(titan_dense_args=TitanDenseModelArgs()) + + # Map TorchTitan attr names → HF attr names + for titan_name, hf_name in model_args._tt_to_hf_attribute_map.items(): + if hasattr(hf_config, hf_name): + setattr(model_args, titan_name, getattr(hf_config, hf_name)) + + # Copy all HF config attributes + for key, value in hf_config_dict.items(): + setattr(model_args, key, value) + + # Override with training-specific settings + model_args.max_seq_len = seq_len + model_args.deterministic = False + model_args.attention_bias = False + model_args.mlp_bias = False + model_args.use_cache = False + model_args.initializer_range = 1.0 + model_args.pruned_heads = getattr(hf_config, "pruned_heads", {}) + + if "head_dim" not in hf_config_dict: + model_args.head_dim = model_args.dim // model_args.num_attention_heads + + return model_args + + +if __name__ == "__main__": + # ── Config ────────────────────────────────────────────────────────── + model_name = "Qwen/Qwen3-0.6B" + seq_len = 512 + num_steps = 50 + lr = 3e-4 + max_norm = 1.0 + tp_degree = 2 + dp_degree = 2 # FSDP shard degree + batch_size = 4 + + # ── Distributed init ──────────────────────────────────────────────── + dist.init_process_group(backend="nccl") + rank = int(os.environ["RANK"]) + local_rank = int(os.environ["LOCAL_RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(local_rank) + device = torch.device(f"cuda:{local_rank}") + + parallel_dims = ParallelDims( + dp_shard=dp_degree, + dp_replicate=1, + tp=tp_degree, + pp=1, + ep=1, + etp=1, + cp=1, + world_size=world_size, + ) + world_mesh = parallel_dims.build_mesh() + + # ── C4 dataset (same as torchtitan) ───────────────────────────────── + from torchtitan.hf_datasets.text_datasets import build_text_dataloader + from torchtitan.components.tokenizer import build_hf_tokenizer + from torchtitan.config.job_config import JobConfig as TTJobConfig + from types import SimpleNamespace + + tt_tokenizer = build_hf_tokenizer( + SimpleNamespace( + model=SimpleNamespace( + hf_assets_path=snapshot_download(model_name), + name="transformers_modeling_backend", + tokenizer_path="", + ) + ) + ) + dp_rank = parallel_dims.get_mesh("fsdp").get_local_rank() + dp_world_size = parallel_dims.get_mesh("fsdp").size() + tt_job_config = TTJobConfig() + tt_job_config.training.dataset = "c4" + tt_job_config.training.dataset_path = None + tt_job_config.training.local_batch_size = batch_size + tt_job_config.training.seq_len = seq_len + dataloader = build_text_dataloader( + dp_world_size=dp_world_size, + dp_rank=dp_rank, + tokenizer=tt_tokenizer, + job_config=tt_job_config, + infinite=True, + ) + + # ── Model ─────────────────────────────────────────────────────────── + model_args = build_model_args(model_name, seq_len) + + with torch.device("meta"): + model = HFTransformerModel(model_args) + + # ── Parallelize (same as torchtitan's parallelize_hf_transformers) ── + tp_mesh = parallel_dims.get_mesh("tp") + apply_non_moe_tp( + model, + tp_mesh, + loss_parallel=True, # lm_head output → Shard(-1) + enable_float8_tensorwise_tp=False, + ) + + dp_mesh = parallel_dims.get_mesh("fsdp") + apply_fsdp( + model, + dp_mesh, + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + pp_enabled=False, + ) + disable_fsdp_gradient_division(model) + + # ── Materialize + init weights ────────────────────────────────────── + model.to_empty(device=device) + with torch.no_grad(): + model.init_weights() + model.train() + + # ── (Optional) Load pretrained weights via DCP ────────────────────── + # Set LOAD_PRETRAINED=1 to load HF weights. Default: train from random init + # (matching what torchtitan's run_train.sh does without a checkpoint). + if os.environ.get("LOAD_PRETRAINED", "0") == "1": + checkpoint_path = snapshot_download(model_name) + state_dict = model.state_dict() + PREFIX = "model." + hf_keyed = {k[len(PREFIX):]: v for k, v in state_dict.items() if k.startswith(PREFIX)} + dcp.load(hf_keyed, storage_reader=HuggingFaceStorageReader(checkpoint_path)) + model.load_state_dict({PREFIX + k: v for k, v in hf_keyed.items()}) + if rank == 0: + print("Pretrained weights loaded via DCP.") + else: + if rank == 0: + print("Training from random init (no pretrained weights).") + + # ── Optimizer ─────────────────────────────────────────────────────── + optimizer = torch.optim.AdamW( + model.parameters(), + lr=lr, + betas=(0.9, 0.95), + eps=1e-8, + weight_decay=0.1, + fused=True, + ) + + # ── loss_parallel context (logits are Shard(-1) on TP mesh) ───────── + loss_parallel_enabled = parallel_dims.tp_enabled + train_context = dist_utils.get_train_context(loss_parallel_enabled) + + # ── Training loop ─────────────────────────────────────────────────── + data_iterator = iter(dataloader) + for step in range(num_steps): + optimizer.zero_grad() + + # torchtitan dataloader yields ({"input": input_ids}, labels) + # both of shape (batch, seq_len) — already shifted, no padding. + input_dict, labels = next(data_iterator) + input_ids = input_dict["input"].to(device) + labels = labels.to(device) + + # No padding in C4 stream — all tokens are valid + local_valid_tokens = (labels != IGNORE_INDEX).sum().to(device) + global_valid_tokens = dist_utils.dist_sum( + local_valid_tokens, parallel_dims.get_mesh("batch") + ) + + # Forward + loss under train_context (enables loss_parallel if TP) + # input_ids and labels are same length (seq_len), already shifted by dataloader. + # pred aligns directly with labels — no slicing needed. + with train_context(): + pred = model(input_ids) # (batch, seq_len, vocab) as Shard(-1) DTensor + loss_sum = F.cross_entropy( + pred.flatten(0, 1).float(), + labels.flatten(0, 1), + reduction="sum", + ignore_index=IGNORE_INDEX, + ) + loss = loss_sum / global_valid_tokens + del pred + loss.backward() + + # Gradient clipping (torchtitan's implementation) + grad_norm = dist_utils.clip_grad_norm_( + list(model.parameters()), max_norm, foreach=True + ) + + optimizer.step() + + if rank == 0: + print( + f"Step {step:>4d} | Loss: {loss.item():.4f} | " + f"Grad norm: {grad_norm.item():.4f}" + ) + + dist.destroy_process_group() diff --git a/verify_loading.py b/verify_loading.py new file mode 100644 index 000000000000..ea008f9626f7 --- /dev/null +++ b/verify_loading.py @@ -0,0 +1,137 @@ +# Save/load roundtrip test for distributed models (TP, FSDP, TP+FSDP). +# +# Verifies that save_pretrained → from_pretrained preserves model weights by +# checking that the cross-entropy loss is identical before and after the roundtrip. +# This catches bugs in DTensor gather-on-save and shard-on-read paths. +# +# Usage: +# python verify_loading.py --mode single_gpu +# torchrun --nproc_per_node=2 verify_loading.py --mode fsdp +# torchrun --nproc_per_node=2 verify_loading.py --mode tp +# torchrun --nproc_per_node=4 verify_loading.py --mode tp_fsdp +# MODEL=Qwen/Qwen3-0.6B torchrun --nproc_per_node=2 verify_loading.py --mode tp +import argparse +import os +import shutil + +import torch +from torch.distributed.tensor import DTensor, Replicate + +from transformers import AutoModelForCausalLM, AutoTokenizer +from transformers.distributed import DistributedConfig + + +parser = argparse.ArgumentParser() +parser.add_argument("--mode", choices=["single_gpu", "fsdp", "tp", "tp_sp", "tp_fsdp", "tp_sp_fsdp"], required=True) +parser.add_argument("--model", type=str, default=None, help="Model ID (or set MODEL env var)") +args = parser.parse_args() + +model_id = args.model or os.environ.get("MODEL") or os.environ.get("MODEL_ID") or "hf-internal-testing/tiny-random-MixtralForCausalLM" + +if args.mode != "single_gpu": + torch.distributed.init_process_group(backend="nccl") + rank = int(os.environ["RANK"]) + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) +else: + rank = 0 + local_rank = 0 + torch.cuda.set_device(0) + +configs = { + "single_gpu": None, + "fsdp": DistributedConfig(fsdp_size=2, fsdp_plan="auto"), + "tp": DistributedConfig(tp_size=2, tp_plan="auto"), + "tp_sp": DistributedConfig(tp_size=2, tp_plan="auto", enable_sequence_parallel=True), + "tp_fsdp": DistributedConfig(tp_size=2, tp_plan="auto", fsdp_size=2, fsdp_plan="auto"), + "tp_sp_fsdp": DistributedConfig(tp_size=2, tp_plan="auto", fsdp_size=2, fsdp_plan="auto", enable_sequence_parallel=True), +} + +tokenizer = AutoTokenizer.from_pretrained(model_id) +text = "The capital of France is Paris. The largest ocean is the Pacific." + + +def materialize_full_logits(logits: torch.Tensor) -> torch.Tensor: + if isinstance(logits, DTensor): + with torch.no_grad(): + return logits.redistribute(placements=[Replicate()] * logits.device_mesh.ndim, async_op=False).to_local() + return logits + + +def compute_loss(model): + inputs = tokenizer(text, return_tensors="pt").to(f"cuda:{local_rank}") + input_ids = inputs["input_ids"] + # Pad sequence length to a multiple of tp_size so DTensor Shard(1) splits evenly + # across ranks in SP mode. Always pad (even for non-TP modes) so that all modes + # compute on the same input and losses are directly comparable. + max_tp = max((c.tp_size if c is not None else 1) for c in configs.values()) + seq_len = input_ids.shape[1] + if seq_len % max_tp != 0: + pad_len = max_tp - (seq_len % max_tp) + pad_token_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id + input_ids = torch.cat([input_ids, input_ids.new_full((1, pad_len), pad_token_id)], dim=1) + labels = input_ids.clone() + labels[:, seq_len:] = -100 # ignore padding in loss + position_ids = torch.arange(input_ids.shape[1], device=input_ids.device).unsqueeze(0) + + model.eval() + with torch.no_grad(): + logits = model(input_ids, position_ids=position_ids).logits + logits = materialize_full_logits(logits) + loss = torch.nn.functional.cross_entropy( + logits.flatten(0, 1).float(), + labels.flatten(0, 1), + reduction="mean", + ignore_index=-100, + ) + return loss.item() + + +# --- Step 1: Load original model and compute loss --- +model = AutoModelForCausalLM.from_pretrained(model_id, distributed_config=configs[args.mode], dtype=torch.float32) +if args.mode == "single_gpu": + model = model.to("cuda:0") + +loss_before = compute_loss(model) +if rank == 0: + print(f"{args.mode}: loss_before = {loss_before:.6f}") + +# --- Step 2: Save to local dir (shared path across ranks) --- +save_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), f"verify_ckpt_{args.mode}") +if rank == 0: + if os.path.exists(save_dir): + shutil.rmtree(save_dir) + os.makedirs(save_dir) +if args.mode != "single_gpu": + torch.distributed.barrier() +model.save_pretrained(save_dir, is_main_process=(rank == 0)) +if rank == 0: + print(f"{args.mode}: saved to {save_dir}") + +# Ensure all ranks see the saved files before reloading +if args.mode != "single_gpu": + torch.distributed.barrier() + +del model +torch.cuda.empty_cache() + +# --- Step 3: Reload from saved checkpoint and compute loss --- +model2 = AutoModelForCausalLM.from_pretrained(save_dir, distributed_config=configs[args.mode], dtype=torch.float32) +if args.mode == "single_gpu": + model2 = model2.to("cuda:0") + +loss_after = compute_loss(model2) +if rank == 0: + print(f"{args.mode}: loss_after = {loss_after:.6f}") + +# --- Step 4: Compare --- +if rank == 0: + diff = abs(loss_before - loss_after) + print(f"{args.mode}: diff = {diff:.2e}") + if diff < 1e-5: + print("PASS: save/load roundtrip is lossless") + else: + print("FAIL: loss mismatch after save/load roundtrip!") + +if args.mode != "single_gpu": + torch.distributed.destroy_process_group() From 34db8405496c10371997e4f0de1cee93feb54ca9 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Mon, 13 Apr 2026 15:40:05 +0000 Subject: [PATCH 007/116] Remove train_fsdp_tp_torchtitan_style.py --- train_fsdp_tp_torchtitan_style.py | 239 ------------------------------ 1 file changed, 239 deletions(-) delete mode 100644 train_fsdp_tp_torchtitan_style.py diff --git a/train_fsdp_tp_torchtitan_style.py b/train_fsdp_tp_torchtitan_style.py deleted file mode 100644 index 325ee112c778..000000000000 --- a/train_fsdp_tp_torchtitan_style.py +++ /dev/null @@ -1,239 +0,0 @@ -# torchrun --nproc_per_node=4 train_fsdp_tp_torchtitan_style.py -# LOAD_PRETRAINED=1 torchrun --nproc_per_node=4 train_fsdp_tp_torchtitan_style.py -# -# Minimal standalone training script that reuses torchtitan's components -# (model wrapper, parallelization, loss, optimizer, grad clipping) directly. -# This is the same code path as `./run_train.sh` but without the config system. - -import os - -import torch -import torch.distributed as dist -import torch.distributed.checkpoint as dcp -import torch.nn.functional as F -from huggingface_hub import snapshot_download -from torch.distributed.checkpoint import HuggingFaceStorageReader - -# ---------- torchtitan imports ---------- -from torchtitan.distributed import ParallelDims -from torchtitan.distributed import utils as dist_utils -from torchtitan.experiments.transformers_modeling_backend.infra.parallelize import ( - apply_fsdp, - apply_non_moe_tp, - disable_fsdp_gradient_division, -) -from torchtitan.experiments.transformers_modeling_backend.model.args import ( - HFTransformerModelArgs, - TitanDenseModelArgs, -) -from torchtitan.experiments.transformers_modeling_backend.model.model import ( - HFTransformerModel, -) - -# ---------- transformers imports ---------- -from transformers import AutoConfig, AutoTokenizer - -IGNORE_INDEX = -100 - - -def build_model_args(hf_model_name: str, seq_len: int) -> HFTransformerModelArgs: - """Build HFTransformerModelArgs from a HuggingFace model name.""" - hf_config = AutoConfig.from_pretrained( - hf_model_name, attn_implementation="sdpa", trust_remote_code=True - ) - hf_config_dict = hf_config.to_dict() - - model_args = HFTransformerModelArgs(titan_dense_args=TitanDenseModelArgs()) - - # Map TorchTitan attr names → HF attr names - for titan_name, hf_name in model_args._tt_to_hf_attribute_map.items(): - if hasattr(hf_config, hf_name): - setattr(model_args, titan_name, getattr(hf_config, hf_name)) - - # Copy all HF config attributes - for key, value in hf_config_dict.items(): - setattr(model_args, key, value) - - # Override with training-specific settings - model_args.max_seq_len = seq_len - model_args.deterministic = False - model_args.attention_bias = False - model_args.mlp_bias = False - model_args.use_cache = False - model_args.initializer_range = 1.0 - model_args.pruned_heads = getattr(hf_config, "pruned_heads", {}) - - if "head_dim" not in hf_config_dict: - model_args.head_dim = model_args.dim // model_args.num_attention_heads - - return model_args - - -if __name__ == "__main__": - # ── Config ────────────────────────────────────────────────────────── - model_name = "Qwen/Qwen3-0.6B" - seq_len = 512 - num_steps = 50 - lr = 3e-4 - max_norm = 1.0 - tp_degree = 2 - dp_degree = 2 # FSDP shard degree - batch_size = 4 - - # ── Distributed init ──────────────────────────────────────────────── - dist.init_process_group(backend="nccl") - rank = int(os.environ["RANK"]) - local_rank = int(os.environ["LOCAL_RANK"]) - world_size = int(os.environ["WORLD_SIZE"]) - torch.cuda.set_device(local_rank) - device = torch.device(f"cuda:{local_rank}") - - parallel_dims = ParallelDims( - dp_shard=dp_degree, - dp_replicate=1, - tp=tp_degree, - pp=1, - ep=1, - etp=1, - cp=1, - world_size=world_size, - ) - world_mesh = parallel_dims.build_mesh() - - # ── C4 dataset (same as torchtitan) ───────────────────────────────── - from torchtitan.hf_datasets.text_datasets import build_text_dataloader - from torchtitan.components.tokenizer import build_hf_tokenizer - from torchtitan.config.job_config import JobConfig as TTJobConfig - from types import SimpleNamespace - - tt_tokenizer = build_hf_tokenizer( - SimpleNamespace( - model=SimpleNamespace( - hf_assets_path=snapshot_download(model_name), - name="transformers_modeling_backend", - tokenizer_path="", - ) - ) - ) - dp_rank = parallel_dims.get_mesh("fsdp").get_local_rank() - dp_world_size = parallel_dims.get_mesh("fsdp").size() - tt_job_config = TTJobConfig() - tt_job_config.training.dataset = "c4" - tt_job_config.training.dataset_path = None - tt_job_config.training.local_batch_size = batch_size - tt_job_config.training.seq_len = seq_len - dataloader = build_text_dataloader( - dp_world_size=dp_world_size, - dp_rank=dp_rank, - tokenizer=tt_tokenizer, - job_config=tt_job_config, - infinite=True, - ) - - # ── Model ─────────────────────────────────────────────────────────── - model_args = build_model_args(model_name, seq_len) - - with torch.device("meta"): - model = HFTransformerModel(model_args) - - # ── Parallelize (same as torchtitan's parallelize_hf_transformers) ── - tp_mesh = parallel_dims.get_mesh("tp") - apply_non_moe_tp( - model, - tp_mesh, - loss_parallel=True, # lm_head output → Shard(-1) - enable_float8_tensorwise_tp=False, - ) - - dp_mesh = parallel_dims.get_mesh("fsdp") - apply_fsdp( - model, - dp_mesh, - param_dtype=torch.bfloat16, - reduce_dtype=torch.float32, - pp_enabled=False, - ) - disable_fsdp_gradient_division(model) - - # ── Materialize + init weights ────────────────────────────────────── - model.to_empty(device=device) - with torch.no_grad(): - model.init_weights() - model.train() - - # ── (Optional) Load pretrained weights via DCP ────────────────────── - # Set LOAD_PRETRAINED=1 to load HF weights. Default: train from random init - # (matching what torchtitan's run_train.sh does without a checkpoint). - if os.environ.get("LOAD_PRETRAINED", "0") == "1": - checkpoint_path = snapshot_download(model_name) - state_dict = model.state_dict() - PREFIX = "model." - hf_keyed = {k[len(PREFIX):]: v for k, v in state_dict.items() if k.startswith(PREFIX)} - dcp.load(hf_keyed, storage_reader=HuggingFaceStorageReader(checkpoint_path)) - model.load_state_dict({PREFIX + k: v for k, v in hf_keyed.items()}) - if rank == 0: - print("Pretrained weights loaded via DCP.") - else: - if rank == 0: - print("Training from random init (no pretrained weights).") - - # ── Optimizer ─────────────────────────────────────────────────────── - optimizer = torch.optim.AdamW( - model.parameters(), - lr=lr, - betas=(0.9, 0.95), - eps=1e-8, - weight_decay=0.1, - fused=True, - ) - - # ── loss_parallel context (logits are Shard(-1) on TP mesh) ───────── - loss_parallel_enabled = parallel_dims.tp_enabled - train_context = dist_utils.get_train_context(loss_parallel_enabled) - - # ── Training loop ─────────────────────────────────────────────────── - data_iterator = iter(dataloader) - for step in range(num_steps): - optimizer.zero_grad() - - # torchtitan dataloader yields ({"input": input_ids}, labels) - # both of shape (batch, seq_len) — already shifted, no padding. - input_dict, labels = next(data_iterator) - input_ids = input_dict["input"].to(device) - labels = labels.to(device) - - # No padding in C4 stream — all tokens are valid - local_valid_tokens = (labels != IGNORE_INDEX).sum().to(device) - global_valid_tokens = dist_utils.dist_sum( - local_valid_tokens, parallel_dims.get_mesh("batch") - ) - - # Forward + loss under train_context (enables loss_parallel if TP) - # input_ids and labels are same length (seq_len), already shifted by dataloader. - # pred aligns directly with labels — no slicing needed. - with train_context(): - pred = model(input_ids) # (batch, seq_len, vocab) as Shard(-1) DTensor - loss_sum = F.cross_entropy( - pred.flatten(0, 1).float(), - labels.flatten(0, 1), - reduction="sum", - ignore_index=IGNORE_INDEX, - ) - loss = loss_sum / global_valid_tokens - del pred - loss.backward() - - # Gradient clipping (torchtitan's implementation) - grad_norm = dist_utils.clip_grad_norm_( - list(model.parameters()), max_norm, foreach=True - ) - - optimizer.step() - - if rank == 0: - print( - f"Step {step:>4d} | Loss: {loss.item():.4f} | " - f"Grad norm: {grad_norm.item():.4f}" - ) - - dist.destroy_process_group() From 6f9e2b679bebca6635c349f14aad0c207bc34723 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Mon, 13 Apr 2026 16:45:59 +0000 Subject: [PATCH 008/116] unify the utils for fsdp --- src/transformers/distributed/utils.py | 9 ++++++++- src/transformers/integrations/__init__.py | 4 ++-- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/src/transformers/distributed/utils.py b/src/transformers/distributed/utils.py index 40278c8fcf2a..2cb9cdca9015 100644 --- a/src/transformers/distributed/utils.py +++ b/src/transformers/distributed/utils.py @@ -27,6 +27,7 @@ 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 @@ -39,12 +40,18 @@ def is_fsdp_enabled() -> bool: 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) or getattr(module, "_is_fsdp_managed_module", False) + return isinstance(module, FullyShardedDataParallel) diff --git a/src/transformers/integrations/__init__.py b/src/transformers/integrations/__init__.py index e3515eab24b1..336db3773f76 100755 --- a/src/transformers/integrations/__init__.py +++ b/src/transformers/integrations/__init__.py @@ -50,7 +50,7 @@ "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"], + "fsdp": ["is_fsdp_enabled", "is_fsdp_managed_module"], "ggml": [ "GGUF_CONFIG_DEFAULTS_MAPPING", "GGUF_CONFIG_MAPPING", @@ -209,7 +209,7 @@ 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 + from .fsdp import is_fsdp_enabled, is_fsdp_managed_module from .ggml import ( GGUF_CONFIG_DEFAULTS_MAPPING, GGUF_CONFIG_MAPPING, From 37dcc14d02e0b4e54e481710e0224ac25f7ec963 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 14 Apr 2026 13:43:58 +0000 Subject: [PATCH 009/116] Fix CI: re-export moved FSDP utils + remove stale type: ignore - Re-export is_fsdp_enabled and is_fsdp_managed_module from integrations/fsdp.py (moved to distributed/utils.py) - Remove unused # type: ignore comments in generation/utils.py --- src/transformers/generation/utils.py | 4 ++-- src/transformers/integrations/fsdp.py | 1 + 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py index ffb7266a5b2f..ec47642f5000 100644 --- a/src/transformers/generation/utils.py +++ b/src/transformers/generation/utils.py @@ -2111,7 +2111,7 @@ def _extract_generation_mode_kwargs( "assistant_model": assistant_model, "streamer": streamer, } - world_size = dist.get_world_size() if dist.is_available() and dist.is_initialized() else 1 # type: ignore + world_size = dist.get_world_size() if dist.is_available() and dist.is_initialized() else 1 generation_mode_kwargs["synced_gpus"] = ( (is_deepspeed_zero3_enabled() or is_fsdp_managed_module(self)) and world_size > 1 if synced_gpus is None @@ -2562,7 +2562,7 @@ def _has_unfinished_sequences(self, this_peer_finished: bool, synced_gpus: bool, # The following logic allows an early break if all peers finished generating their sequence this_peer_finished_flag = torch.tensor(0.0 if this_peer_finished else 1.0, device=device) # send 0.0 if we finished, 1.0 otherwise - dist.all_reduce(this_peer_finished_flag, op=dist.ReduceOp.SUM) # type: ignore + dist.all_reduce(this_peer_finished_flag, op=dist.ReduceOp.SUM) # did all peers finish? the reduced sum will be 0.0 then if this_peer_finished_flag.item() == 0.0: return False diff --git a/src/transformers/integrations/fsdp.py b/src/transformers/integrations/fsdp.py index 10936490ee03..526aa2780ab3 100644 --- a/src/transformers/integrations/fsdp.py +++ b/src/transformers/integrations/fsdp.py @@ -17,6 +17,7 @@ import os from typing import Any, Literal +from ..distributed.utils import is_fsdp_enabled, is_fsdp_managed_module # noqa: F401 from ..utils import is_torch_available, is_torch_greater_or_equal, logging from ..utils.quantization_config import QuantizationMethod From 21f05610b20caf50e6e017e23f94f8d027f5cefd Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 14 Apr 2026 14:22:30 +0000 Subject: [PATCH 010/116] Fix ruff formatting in core_model_loading.py --- src/transformers/core_model_loading.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 2b528bd0a829..63ff8fac7d7c 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -1004,6 +1004,7 @@ def get_parallel_materialization_context( return None + def dot_natural_key(s: str): """Sort key for state-dict names: split on ``"."`` and sort digits numerically and strings alphabetically. We emit a tuple at each point to sort ints From cd45107fc506798b41789fcacffd257320f7db8e Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 14 Apr 2026 14:24:02 +0000 Subject: [PATCH 011/116] Fix ruff linting and formatting --- src/transformers/core_model_loading.py | 1 + src/transformers/distributed/configuration_utils.py | 9 ++++++--- src/transformers/integrations/tensor_parallel.py | 2 +- src/transformers/models/qwen3/modeling_qwen3.py | 1 + tests/test_tensor_parallel_mixin.py | 5 ++++- 5 files changed, 13 insertions(+), 5 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 2b528bd0a829..63ff8fac7d7c 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -1004,6 +1004,7 @@ def get_parallel_materialization_context( return None + def dot_natural_key(s: str): """Sort key for state-dict names: split on ``"."`` and sort digits numerically and strings alphabetically. We emit a tuple at each point to sort ints diff --git a/src/transformers/distributed/configuration_utils.py b/src/transformers/distributed/configuration_utils.py index e40aed267bcd..89281b0a9a39 100644 --- a/src/transformers/distributed/configuration_utils.py +++ b/src/transformers/distributed/configuration_utils.py @@ -14,9 +14,10 @@ import json import os -import torch from dataclasses import asdict, dataclass +import torch + @dataclass class DistributedConfig: @@ -46,9 +47,11 @@ def __post_init__(self): self.tp_plan = "auto" if self.fsdp_size > 1 and self.fsdp_plan is None: self.fsdp_plan = "auto" - + world_size = torch.distributed.get_world_size() - assert self.tp_size * self.fsdp_size == world_size, f"tp_size ({self.tp_size}) * fsdp_size ({self.fsdp_size}) must be equal to world_size ({world_size})" + assert self.tp_size * self.fsdp_size == world_size, ( + f"tp_size ({self.tp_size}) * fsdp_size ({self.fsdp_size}) must be equal to world_size ({world_size})" + ) @classmethod def from_dict(cls, config_dict: dict, **kwargs) -> "DistributedConfig": diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index 41f4414c8d91..c7e4546017c7 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -17,7 +17,7 @@ from dataclasses import dataclass from typing import Literal -from torch.distributed.tensor import DTensor, Replicate, Shard +from torch.distributed.tensor import Replicate, Shard from torch.distributed.tensor.parallel import ( ColwiseParallel, RowwiseParallel, diff --git a/src/transformers/models/qwen3/modeling_qwen3.py b/src/transformers/models/qwen3/modeling_qwen3.py index 30e7f5d2d9e1..a9e4cd4b75bb 100644 --- a/src/transformers/models/qwen3/modeling_qwen3.py +++ b/src/transformers/models/qwen3/modeling_qwen3.py @@ -456,6 +456,7 @@ def __init__(self, config): # Initialize weights and apply final processing self.post_init() + @can_return_tuple @auto_docstring def forward( diff --git a/tests/test_tensor_parallel_mixin.py b/tests/test_tensor_parallel_mixin.py index c138be86ad68..e07d2beb539b 100644 --- a/tests/test_tensor_parallel_mixin.py +++ b/tests/test_tensor_parallel_mixin.py @@ -36,6 +36,7 @@ from torch.distributed.tensor import DTensor, Replicate 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"): @@ -139,7 +140,9 @@ def _load_tp_and_reference_models(model_path, model_class, enable_sequence_paral tuple: (model_tp, model_ref, device) """ tp_size = dist.get_world_size() - distributed_config = DistributedConfig(tp_size=tp_size, tp_plan="auto", enable_sequence_parallel=enable_sequence_parallel) + distributed_config = DistributedConfig( + tp_size=tp_size, tp_plan="auto", enable_sequence_parallel=enable_sequence_parallel + ) model_tp = model_class.from_pretrained( model_path, distributed_config=distributed_config, attn_implementation="sdpa" ) From ba3990fb960b88d293b2e9e295fa70b33024573d Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 14 Apr 2026 15:21:27 +0000 Subject: [PATCH 012/116] Backport new TP/FSDP API from orchestration-save-load branch --- src/transformers/core_model_loading.py | 474 +++++++++-------- src/transformers/distributed/utils.py | 89 +++- src/transformers/integrations/fsdp.py | 56 +- .../integrations/tensor_parallel.py | 479 +++++++++++++++++- src/transformers/modeling_utils.py | 267 +++++----- .../models/afmoe/modeling_afmoe.py | 5 + .../models/apertus/modeling_apertus.py | 5 + .../models/arcee/modeling_arcee.py | 5 + src/transformers/models/aria/modeling_aria.py | 5 + .../models/bitnet/modeling_bitnet.py | 5 + src/transformers/models/blt/modeling_blt.py | 5 + .../models/chameleon/modeling_chameleon.py | 4 + src/transformers/models/csm/modeling_csm.py | 5 + .../models/cwm/configuration_cwm.py | 15 +- src/transformers/models/cwm/modeling_cwm.py | 5 + src/transformers/models/dbrx/modeling_dbrx.py | 5 + .../deepseek_v3/modeling_deepseek_v3.py | 5 + src/transformers/models/dia/modeling_dia.py | 5 + .../models/diffllama/modeling_diffllama.py | 5 + src/transformers/models/doge/modeling_doge.py | 5 + .../models/dots1/modeling_dots1.py | 5 + src/transformers/models/emu3/modeling_emu3.py | 5 + .../models/eurobert/modeling_eurobert.py | 5 + .../models/exaone4/modeling_exaone4.py | 5 + .../models/exaone_moe/modeling_exaone_moe.py | 5 + .../models/falcon/modeling_falcon.py | 4 + .../models/falcon_h1/modeling_falcon_h1.py | 5 + .../models/gemma/modeling_gemma.py | 5 + .../models/gemma2/modeling_gemma2.py | 5 + .../models/gemma3/modeling_gemma3.py | 5 + .../glm4_moe_lite/modeling_glm4_moe_lite.py | 5 + .../modeling_gpt_neox_japanese.py | 4 + .../models/granite/modeling_granite.py | 5 + .../models/granitemoe/modeling_granitemoe.py | 5 + .../modeling_granitemoehybrid.py | 5 + .../modeling_granitemoeshared.py | 5 + .../configuration_higgs_audio_v2.py | 15 +- .../higgs_audio_v2/modeling_higgs_audio_v2.py | 5 + .../modeling_hunyuan_v1_dense.py | 5 + .../hunyuan_v1_moe/modeling_hunyuan_v1_moe.py | 5 + .../models/jais2/modeling_jais2.py | 5 + .../models/jamba/modeling_jamba.py | 5 + .../models/jetmoe/modeling_jetmoe.py | 5 + .../modeling_jina_embeddings_v3.py | 5 + .../modeling_kyutai_speech_to_text.py | 5 + src/transformers/models/lasr/modeling_lasr.py | 5 + src/transformers/models/lfm2/modeling_lfm2.py | 5 + .../models/lfm2_moe/modeling_lfm2_moe.py | 5 + src/transformers/models/mimi/modeling_mimi.py | 4 + .../models/minimax/modeling_minimax.py | 5 + .../ministral/configuration_ministral.py | 15 +- .../models/ministral/modeling_ministral.py | 5 + .../models/ministral3/modeling_ministral3.py | 5 + .../models/mistral4/modeling_mistral4.py | 5 + .../models/mixtral/modeling_mixtral.py | 5 + .../models/mllama/modeling_mllama.py | 4 + .../models/moshi/modeling_moshi.py | 4 + .../models/nanochat/modeling_nanochat.py | 5 + .../models/nemotron_h/modeling_nemotron_h.py | 5 + .../models/nomic_bert/modeling_nomic_bert.py | 5 + .../models/olmoe/modeling_olmoe.py | 5 + .../models/parakeet/modeling_parakeet.py | 5 + .../models/persimmon/modeling_persimmon.py | 4 + src/transformers/models/phi/modeling_phi.py | 5 + .../models/phimoe/modeling_phimoe.py | 5 + .../models/qwen2/modeling_qwen2.py | 5 + .../models/qwen2_moe/modeling_qwen2_moe.py | 5 + .../models/qwen3/modeling_qwen3.py | 1 + .../models/qwen3_moe/modeling_qwen3_moe.py | 5 + .../qwen3_omni_moe/modeling_qwen3_omni_moe.py | 5 + .../models/qwen3_vl/modeling_qwen3_vl.py | 5 + .../qwen3_vl_moe/modeling_qwen3_vl_moe.py | 5 + .../modeling_recurrent_gemma.py | 4 + .../models/seed_oss/modeling_seed_oss.py | 5 + .../models/smollm3/modeling_smollm3.py | 5 + .../models/solar_open/modeling_solar_open.py | 5 + .../models/stablelm/modeling_stablelm.py | 4 + .../models/starcoder2/modeling_starcoder2.py | 5 + .../models/t5gemma/modeling_t5gemma.py | 5 + .../models/t5gemma2/modeling_t5gemma2.py | 5 + .../models/timesfm2_5/modeling_timesfm2_5.py | 5 + .../models/vaultgemma/modeling_vaultgemma.py | 5 + .../modeling_voxtral_realtime.py | 5 + .../models/youtu/modeling_youtu.py | 5 + .../models/zamba2/modeling_zamba2.py | 5 + 85 files changed, 1356 insertions(+), 426 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 63ff8fac7d7c..cf06f62fd685 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -32,7 +32,6 @@ import torch from .integrations.accelerate import get_device, offload_weight -from .integrations.tensor_parallel import ALL_PARALLEL_STYLES, get_tensor_shard from .utils import is_env_variable_true from .utils.loading_report import LoadStateDictInfo from .utils.logging import get_logger, tqdm @@ -46,6 +45,7 @@ elif _torch_distributed_available: from torch.distributed.tensor import DTensor from torch.distributed.tensor._utils import compute_local_shape_and_global_offset + from torch.distributed.tensor.placement_types import Shard logger = get_logger(__name__) @@ -118,17 +118,6 @@ def reverse_op(self) -> ConversionOps: raise NotImplementedError -class _IdentityOp(ConversionOps): - """Pass-through reverse op for dequantize operations. - - Dequantized weights are already in their target dtype and should be - saved as-is without any conversion. - """ - - def convert(self, input_dict: dict[str, Any], **kwargs) -> dict[str, Any]: - return input_dict - - class Chunk(ConversionOps): """Split a tensor along ``dim`` into equally sized chunks.""" @@ -535,16 +524,11 @@ class WeightTransform: target_patterns: str | list[str] = field(init=True) compiled_sources: re.Pattern = field(init=False) - distributed_operation: Any | None = None quantization_operation: ConversionOps | None = None collected_tensors: dict[str, list[Future]] = field(default_factory=lambda: defaultdict(list), init=False) layer_targets: dict[str, set[str]] = field(default_factory=lambda: defaultdict(set), init=False) - # Those are needed to be able to reverse correctly the transform, as the patterns may be processed - _original_source_patterns: list[str] = field(init=False) - _original_target_patterns: list[str] = field(init=False) - def __setattr__(self, name, value): if name in ("source_patterns", "target_patterns"): # We do not allow to re-set the patterns, as they are linked between each other and changing one @@ -561,13 +545,10 @@ def __post_init__(self): # when instantiating the reverse mapping (i.e. the targets become sources, and sources become targets) # The issues lie in the sources usually, so here we need to check the targets for the reversed mapping - # We need to copy the exact original patterns to later reverse (before processing may change them) - self._original_source_patterns = self.source_patterns.copy() - self._original_target_patterns = self.target_patterns.copy() - # Process target_patterns: detect capturing groups and replace with \1 # Store the original capturing group patterns for reverse mapping target_capturing_groups: list[str] = [] + unprocess_targets = self.target_patterns.copy() for i, pattern in enumerate(self.target_patterns): self.target_patterns[i], captured_group = process_target_pattern(pattern) if captured_group is not None: @@ -596,7 +577,7 @@ def __post_init__(self): pattern = pattern.replace(r"\1", unique_capturing_group, 1) # Potentially process a bit more for consistency - only if they are consistent pairs, i.e. the length is the same if len(self.source_patterns) == len(self.target_patterns): - pattern = process_source_pattern(pattern, self._original_target_patterns[i]) + pattern = process_source_pattern(pattern, unprocess_targets[i]) self.source_patterns[i] = pattern # Construct the regex we will use to rename keys from the sources to the targets @@ -635,7 +616,7 @@ def rename_source_key(self, source_key: str) -> tuple[str, str | None]: # inside that matched named group replaced_group_idx = self.compiled_sources.groupindex[matching_group_name] + 1 replacement = replacement.replace(r"\1", match_object.group(replaced_group_idx)) - renamed_key = source_key.replace(match_object.group(0), replacement, 1) + renamed_key = source_key.replace(match_object.group(0), replacement) return renamed_key, source_pattern_that_matched def reverse_transform(self) -> WeightTransform: @@ -651,7 +632,7 @@ 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.target_patterns, target_patterns=self.source_patterns, **kwargs ) return reverse_transform @@ -840,169 +821,246 @@ 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_parallel_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 active parallelism strategy if `thread_pool` is provided, or - return a Callable that will load the tensor synchronously when called.""" + When ``sharding_op`` is given the tensor is sharded according to the DTensor + placement strategy; otherwise it is simply copied to *device*/*dtype*. + Without a thread pool a deferred callable is returned instead of a Future. + """ def _job(): - return sharding_method.shard_tensor(tensor, tensor_idx=tensor_idx, device=device, dtype=dtype) + if sharding_op is not None: + return sharding_op.shard_tensor(tensor, tensor_idx=tensor_idx, device=device, dtype=dtype) + return _materialize_copy(tensor, device, dtype) if thread_pool is not None: return thread_pool.submit(_job) - else: - # Return the Callable here, not the Tensor itself, so we actually delay loading to avoid saturating cpu - # memory during Conversion - return _job - - -@dataclass(slots=True) -class ParallelMaterializationContext: - distributed_operation: Any - tensor_idx: int | None - device: Any + # Return the Callable here, not the Tensor itself, so we actually delay loading + # to avoid saturating cpu memory during Conversion + return _job -def is_dtensor_like(value: Any) -> bool: - return all(hasattr(value, attr) for attr in ("device_mesh", "placements", "to_local")) +class DtensorShardOperation: + """Extracts the local shard from a full checkpoint tensor based on this rank's DTensor placements.""" + # tensor_dim -> list of (start, end) index ranges this rank owns + DimRanges = dict[int, list[tuple[int, int]]] -@dataclass(slots=True) -class FSDPShardOperation: - device_mesh: Any - rank: int - empty_param: Any - placements: tuple[Any, ...] - shard_placement: Any | None = field(init=False, default=None) - local_shape: tuple[int, ...] = field(init=False) - - def __post_init__(self): - shard_placements = [placement for placement in self.placements if placement.is_shard()] - if len(shard_placements) > 1: - raise NotImplementedError( - f"FSDP shard-on-read does not support multiple shard placements yet: {self.placements}" - ) - self.shard_placement = shard_placements[0] if shard_placements else None - if self.shard_placement is not None and len(self.placements) != 1: - raise NotImplementedError( - f"FSDP shard-on-read only supports a single placement today. Got placements={self.placements}." - ) - self.local_shape = self.get_expected_sharded_shape(self.empty_param.shape) - - @classmethod - def from_param(cls, param: Any) -> FSDPShardOperation: - return cls( - device_mesh=param.device_mesh, - rank=param.device_mesh.get_local_rank(), - empty_param=param, - placements=tuple(param.placements), - ) + def __init__(self, param: DTensor): + self.device_mesh = param.device_mesh + self.param = param + self.placements = tuple(param.placements) + local_shape, _ = compute_local_shape_and_global_offset(param.shape, self.device_mesh, self.placements) + self.local_shape = tuple(local_shape) def shard_tensor( self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None ) -> torch.Tensor | None: - if self.shard_placement is None: - local_tensor = param[...] - else: - param_shape = list(param.shape) if isinstance(param, torch.Tensor) else param.get_shape() - # Mixtral-style converted expert weights first stack individual expert tensors along dim 0 before - # concatenating. Only materialize the experts owned by this rank. - if ( - tensor_idx is not None - and len(self.empty_param.shape) == len(param_shape) + 1 - and self.shard_placement.dim == 0 - ): - local_expert_count = self.local_shape[0] - expert_offset = compute_local_shape_and_global_offset( - self.empty_param.shape, self.device_mesh, self.placements - )[1][0] - if tensor_idx < expert_offset or tensor_idx >= expert_offset + local_expert_count: + """Return the local shard of ``param`` for this rank, dispatching to the appropriate strategy.""" + # Find which placements actually shard data. + # _StridedShard.is_shard() returns False in PyTorch, so we also check for + # the ``dim`` attribute that both Shard and _StridedShard have. + sharding_placements = [ + (i, p) + for i, p in enumerate(self.placements) + if p.is_shard() or (hasattr(p, "dim") and not p.is_replicate()) + ] + param_shape = list(param.shape) if isinstance(param, torch.Tensor) else param.get_shape() + + if not sharding_placements: + return param[...].to(device=device, dtype=dtype) + + if tensor_idx is not None and len(self.param.shape) == len(param_shape) + 1: + # Expert parallelism: dim 0 (expert dimension) is sharded across ranks. + # When dim 0 is the only sharding placement, return the full expert or + # skip it. When TP also shards an inner dim, keep applying the remaining + # placements to the owned expert tensor. + has_expert_sharding = any(self._normalize_param_dim(p.dim) == 0 for _, p in sharding_placements) + if has_expert_sharding: + if not self._owns_local_expert(tensor_idx): return None - local_tensor = param[...] - else: - local_tensor = get_tensor_shard( - param, - self.empty_param, - self.device_mesh, - self.rank, - self.shard_placement.dim, - tensor_idx=tensor_idx, + inner_placements = [(i, p) for i, p in sharding_placements if self._normalize_param_dim(p.dim) != 0] + if not inner_placements: + return param[...].to(device=device, dtype=dtype) + return self._shard_nd(param, inner_placements, param_shape, device, dtype) + + return self._shard_nd(param, sharding_placements, param_shape, device, dtype) + + def _shard_nd(self, param, sharding_placements, param_shape, device, dtype): + """Handle multi-dimensional sharding, choosing the best strategy.""" + if not self._can_shard_on_read(sharding_placements): + return self._materialize_and_split(param, sharding_placements, device, dtype) + + # All placements are plain Shard on different dims. + # compute_local_shape_and_global_offset gives us one contiguous range per dim directly. + has_strided = any(not p.is_shard() for _, p in sharding_placements) + if not has_strided: + local_shape, global_offset = compute_local_shape_and_global_offset( + self.param.shape, self.device_mesh, self.placements + ) + slices = [slice(None)] * len(param_shape) + for _, placement in sharding_placements: + dim = self._checkpoint_dim(placement.dim, param_shape) + offset = global_offset[placement.dim] + slices[dim] = slice(offset, offset + local_shape[placement.dim]) + return param[tuple(slices)].to(device=device, dtype=dtype) + + dim_ranges = self._compute_dim_ranges(sharding_placements, param_shape) + return self._slice_and_read(param, param_shape, dim_ranges, device, dtype) + + def _can_shard_on_read(self, sharding_placements) -> bool: + """Check whether range-based shard-on-read is feasible. + + Returns ``False`` when a ``_StridedShard`` and another placement share the + same tensor dimension — the strided reorder can't be composed via range + arithmetic because ``Shard`` would need to cut across the concatenated + result of ``_StridedShard``'s disjoint ranges. + """ + dims_seen: dict[int, bool] = {} # dim -> has_strided + for _, placement in sharding_placements: + dim = placement.dim + is_strided = not placement.is_shard() + if dim in dims_seen and (is_strided or dims_seen[dim]): + logger.debug( + "Cannot shard-on-read: dim %d has both Shard and _StridedShard placements, " + "falling back to materialize-then-split.", + dim, ) - if local_tensor is None: - return None - return local_tensor.to(device=device, dtype=dtype) - - def get_expected_sharded_shape(self, full_shape: tuple[int, ...] | torch.Size) -> tuple[int, ...]: - local_shape, _ = compute_local_shape_and_global_offset(full_shape, self.device_mesh, self.placements) - return tuple(local_shape) - - def update_module_attributes(self, module: torch.nn.Module): - return None - - -def get_parallel_materialization_context( - mapping: WeightTransform, - renamed_key: str, - source_pattern: str, - empty_param: Any, - device_mesh: Any, - parallel_plan: dict[str, Any], - parallel_pattern_matcher: re.Pattern | None, - parallel_pattern_by_group_name: dict[str, str] | None, - device_map: dict[str, Any], -) -> ParallelMaterializationContext | None: - tensor_idx = ( - len(mapping.collected_tensors.get(source_pattern, [])) - if isinstance(mapping, WeightConverter) and isinstance(mapping.operations[0], MergeModulelist) - else None - ) + return False + dims_seen[dim] = is_strided + return True + + def _materialize_and_split(self, param, sharding_placements, device, dtype): + """Fallback: load the full tensor, then iteratively split per mesh dim.""" + tensor = param[...] if not isinstance(param, torch.Tensor) else param + for mesh_dim_idx, placement in sharding_placements: + sub_mesh = self._get_sub_mesh(mesh_dim_idx) + rank = sub_mesh.get_local_rank() + shards, _ = placement._split_tensor(tensor, sub_mesh.size(), with_padding=False, contiguous=True) + tensor = shards[rank] + return tensor.to(device=device, dtype=dtype) + + def _compute_dim_ranges(self, sharding_placements, param_shape) -> DtensorShardOperation.DimRanges: + """Compute per-dimension index ranges for this rank. + + Each sharding placement narrows the ranges on its tensor dimension: + - ``Shard``: one contiguous sub-range per previous range. + - ``_StridedShard``: multiple disjoint sub-ranges (one per split-factor group). + """ + dim_ranges: DtensorShardOperation.DimRanges = {} + for mesh_dim_idx, placement in sharding_placements: + sub_mesh = self._get_sub_mesh(mesh_dim_idx) + rank = sub_mesh.get_local_rank() + world_size = sub_mesh.size() + dim = self._checkpoint_dim(placement.dim, param_shape) + prev_ranges = dim_ranges.get(dim, [(0, param_shape[dim])]) + + if placement.is_shard(): + new_ranges = self._contiguous_ranges(prev_ranges, rank, world_size) + elif self._source_tensor_needs_packing(param_shape): + # _StridedShard only makes sense once the packed axis exists. While + # loading pre-packed source tensors (e.g. w1/w3 before gate_up_proj + # concatenation), take the contiguous chunk for this rank and let the + # WeightConverter recreate the packed layout afterward. + new_ranges = self._contiguous_ranges(prev_ranges, rank, world_size) + else: + new_ranges = self._strided_ranges(prev_ranges, rank, world_size, placement.split_factor) + dim_ranges[dim] = new_ranges + return dim_ranges - if ( - device_mesh - and parallel_plan - and parallel_pattern_matcher is not None - and parallel_pattern_by_group_name is not None - ): - if matched_parallel_pattern := parallel_pattern_matcher.search(renamed_key): - matched_parallel_pattern = parallel_pattern_by_group_name[matched_parallel_pattern.lastgroup] - if getattr(mapping, "distributed_operation", None) is None: - parallel_layer = ALL_PARALLEL_STYLES[parallel_plan[matched_parallel_pattern]].__class__ - mapping.distributed_operation = parallel_layer( - device_mesh=device_mesh, rank=device_mesh.get_local_rank(), empty_param=empty_param.clone() - ) - return ParallelMaterializationContext(mapping.distributed_operation, tensor_idx, device_map[""]) - - if is_dtensor_like(empty_param): - if getattr(mapping, "distributed_operation", None) is None: - mapping.distributed_operation = FSDPShardOperation.from_param(empty_param) - return ParallelMaterializationContext( - mapping.distributed_operation, - tensor_idx, - get_device(device_map, renamed_key, valid_torch_device=True), - ) + def _slice_and_read(self, param, param_shape, dim_ranges: DtensorShardOperation.DimRanges, device, dtype): + """Build slices from computed ranges and read from the tensor. - return None + At most one dim can have multiple disjoint ranges (from ``_StridedShard``). + If so, read each disjoint range separately and concatenate. + """ + concat_dim = None + concat_ranges = None + base_slices = [slice(None)] * len(param_shape) + for dim, ranges in dim_ranges.items(): + if len(ranges) == 1: + base_slices[dim] = slice(ranges[0][0], ranges[0][1]) + elif len(ranges) > 1: + if concat_dim is not None: + raise ValueError("Shard-on-read only supports disjoint ranges on a single checkpoint dimension.") + concat_dim = dim + concat_ranges = ranges + + if concat_dim is None: + return param[tuple(base_slices)].to(device=device, dtype=dtype) + + pieces = [] + for start, end in concat_ranges: + slices = list(base_slices) + slices[concat_dim] = slice(start, end) + pieces.append(param[tuple(slices)]) + return torch.cat(pieces, dim=concat_dim).to(device=device, dtype=dtype) + + # ------------------------------------------------------------------ + # Utilities + # ------------------------------------------------------------------ + + def _contiguous_ranges( + self, prev_ranges: list[tuple[int, int]], rank: int, world_size: int + ) -> list[tuple[int, int]]: + """Narrow each range by picking the ``rank``-th contiguous chunk (``Shard`` semantics).""" + new_ranges = [] + for prev_start, prev_end in prev_ranges: + shard_size, offset = Shard.local_shard_size_and_offset(prev_end - prev_start, world_size, rank) + if shard_size > 0: + new_ranges.append((prev_start + offset, prev_start + offset + shard_size)) + return new_ranges + + def _strided_ranges( + self, prev_ranges: list[tuple[int, int]], rank: int, world_size: int, split_factor: int + ) -> list[tuple[int, int]]: + """Narrow each range using ``_StridedShard`` semantics. + + Divides each range into ``split_factor`` groups, then within each group + picks the ``rank``-th chunk of ``world_size`` equal pieces. + """ + new_ranges = [] + for prev_start, prev_end in prev_ranges: + group_size = math.ceil((prev_end - prev_start) / split_factor) + for g in range(split_factor): + g_start = prev_start + g * group_size + g_end = min(g_start + group_size, prev_end) + if g_end <= g_start: + continue + shard_size, offset = Shard.local_shard_size_and_offset(g_end - g_start, world_size, rank) + if shard_size > 0: + new_ranges.append((g_start + offset, g_start + offset + shard_size)) + return new_ranges + + def _get_sub_mesh(self, mesh_dim_idx: int): + """Return the 1-D sub-mesh for ``mesh_dim_idx``.""" + if self.device_mesh.ndim > 1: + return self.device_mesh[self.device_mesh.mesh_dim_names[mesh_dim_idx]] + return self.device_mesh + + def _normalize_param_dim(self, dim: int) -> int: + return dim if dim >= 0 else self.param.ndim + dim + + def _checkpoint_dim(self, placement_dim: int, param_shape) -> int: + """Map a placement dim from the DTensor shape to the checkpoint tensor shape.""" + dim = self._normalize_param_dim(placement_dim) + ndim_diff = self.param.ndim - len(param_shape) + if ndim_diff > 0 and dim >= ndim_diff: + dim -= ndim_diff + return dim + + def _owns_local_expert(self, tensor_idx: int) -> bool: + _, offsets = compute_local_shape_and_global_offset(self.param.shape, self.device_mesh, self.placements) + return offsets[0] <= tensor_idx < offsets[0] + self.local_shape[0] + + def _source_tensor_needs_packing(self, param_shape) -> bool: + # A single source tensor still missing the leading expert axis is being + # converted into a packed expert parameter. In that case _StridedShard's + # split groups do not exist yet. + return self.param.ndim == len(param_shape) + 1 def dot_natural_key(s: str): @@ -1073,7 +1131,6 @@ def set_param_for_module( target_name: str, param_value: torch.Tensor, loading_info: LoadStateDictInfo, - distributed_operation: Any | None, hf_quantizer: HfQuantizer, ): module_path, _, param_name = target_name.rpartition(".") @@ -1088,26 +1145,30 @@ 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) and not is_dtensor_like(ref): + 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/FSDP shard-on-read, use the local shard shape; otherwise, use full shape - if is_dtensor_like(ref): - local_shape, _ = compute_local_shape_and_global_offset(ref.shape, ref.device_mesh, ref.placements) + if isinstance(ref, DTensor): + local_shape, global_offset = compute_local_shape_and_global_offset( + ref.shape, ref.device_mesh, ref.placements + ) expected_shape = torch.Size(local_shape) - elif distributed_operation is not None: - expected_shape = torch.Size(distributed_operation.get_expected_sharded_shape(ref.shape)) else: expected_shape = ref.shape + # When a WeightConverter produces the full global tensor, slice it to the local DTensor shard. + if isinstance(ref, DTensor) and param_value.shape == ref.shape and param_value.shape != expected_shape: + slices = [slice(global_offset[d], global_offset[d] + local_shape[d]) for d in range(param_value.ndim)] + param_value = param_value[tuple(slices)].contiguous() + 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 is_dtensor_like(ref): + if isinstance(ref, DTensor): local_param = param_value.detach() if isinstance(param_value, torch.nn.Parameter) else param_value fsdp_param = DTensor.from_local( local_param.contiguous(), @@ -1128,8 +1189,6 @@ def set_param_for_module( # 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( @@ -1315,10 +1374,21 @@ def convert_and_load_state_dict_in_model( """ prefix = model.base_model_prefix tp_plan = tp_plan or {} - device_map = load_config.device_map or {"": "cpu"} hf_quantizer = load_config.hf_quantizer dtype = load_config.dtype device_mesh = load_config.device_mesh + + if load_config.device_map is not None: + device_map = load_config.device_map + elif device_mesh is not None: + if device_mesh.device_type == "cpu": + device_map = {"": torch.device("cpu")} + else: + device_map = { + "": torch.device(device_mesh.device_type, getattr(torch, device_mesh.device_type).current_device()) + } + else: + device_map = {"": "cpu"} disk_offload_folder = load_config.disk_offload_folder offload_buffers = load_config.offload_buffers dtype_plan = load_config.dtype_plan or {} @@ -1353,10 +1423,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())) @@ -1433,30 +1499,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 parallel shard-on-read or device_map placement - future_or_tensor = None - if parallel_context := get_parallel_materialization_context( - mapping=mapping, - renamed_key=renamed_key, - source_pattern=source_pattern, - empty_param=empty_param, - device_mesh=device_mesh, - parallel_plan=tp_plan, - parallel_pattern_matcher=tp_plan_alt if tp_plan else None, - parallel_pattern_by_group_name=tp_plan_by_group_name if tp_plan else None, - device_map=device_map, - ): - future_or_tensor = spawn_parallel_materialize( + # 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, - parallel_context.distributed_operation, - parallel_context.tensor_idx, - parallel_context.device, + param_device, _dtype, + sharding_op=DtensorShardOperation(empty_param), + tensor_idx=tensor_idx, ) - - if future_or_tensor is None: - param_device = get_device(device_map, renamed_key, valid_torch_device=True) + else: future_or_tensor = spawn_materialize(thread_pool, tensor, param_device, _dtype) mapping.add_tensor(renamed_key, original_key, source_pattern, future_or_tensor) @@ -1486,14 +1545,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/utils.py b/src/transformers/distributed/utils.py index 2cb9cdca9015..fa4d554f130d 100644 --- a/src/transformers/distributed/utils.py +++ b/src/transformers/distributed/utils.py @@ -16,14 +16,19 @@ import os from typing import TYPE_CHECKING -from ..utils import is_torch_available, strtobool +from ..utils import is_torch_available, is_torch_greater_or_equal, strtobool 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 ..integrations.tensor_parallel import convert_strided_to_shard, restore_strided_from_shard def is_fsdp_enabled() -> bool: @@ -55,3 +60,85 @@ def is_fsdp_managed_module(module: nn.Module) -> bool: except ImportError: return False return isinstance(module, FullyShardedDataParallel) + + +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"} + backend = backend_map.get(device_type) + + torch.distributed.init_process_group(backend=backend, rank=rank, world_size=world_size) + current_device = getattr(torch, device_type) + if device_type != "cpu": + current_device.set_device(local_rank) + except Exception as e: + raise OSError( + "We tried to initialize torch.distributed for you, but it failed. Make " + "sure you init torch distributed in your script to use distributed training." + ) from e + + +def init_device_mesh(distributed_config: DistributedConfig) -> torch.distributed.device_mesh.DeviceMesh: + if not is_torch_greater_or_equal("2.5"): + raise OSError("Distributed training with DistributedConfig requires `torch>=2.5`.") + + device_type = torch._C._get_accelerator().type + _ensure_torch_distributed(device_type) + + world_size = torch.distributed.get_world_size() + if device_type != "cpu": + getattr(torch, device_type).set_device(int(os.environ.get("LOCAL_RANK", 0))) + + tp_size = distributed_config.tp_size + fsdp_size = distributed_config.fsdp_size + + assert world_size == tp_size * fsdp_size, ( + f"world_size ({world_size}) must be equal to tp_size ({tp_size}) * fsdp_size ({fsdp_size})" + ) + + dims, names = [], [] + if fsdp_size > 1: + dims.append(fsdp_size) + names.append("fsdp") + if tp_size > 1: + dims.append(tp_size) + names.append("tp") + + # Build from a 1D world mesh via _unflatten so that PyTorch can flatten + # sub-dimensions back when needed (e.g. for single all_reduce across + # [fsdp, tp] during grad norm computation instead of 2 sequential ones). + world_mesh = torch.distributed.init_device_mesh(device_type, (world_size,), mesh_dim_names=("world",)) + mesh = world_mesh._unflatten(0, tuple(dims), tuple(names)) + + # Pre-create flattened sub-mesh for multi-dimensional meshes so DTensor + # can use a single collective instead of sequential per-dimension ones. + if len(dims) > 1: + mesh._flatten("_".join(names)) + + return mesh + + +def save_optimizer(optimizer, checkpoint_dir: str) -> None: + # Save optimizer state via DCP, handling _StridedShard placements transparently. + osd = optimizer.state_dict() + placement_map = convert_strided_to_shard(osd) + dcp.save({"optimizer": osd}, checkpoint_id=checkpoint_dir) + if placement_map and torch.distributed.get_rank() == 0: + torch.save(placement_map, os.path.join(checkpoint_dir, "placement_map.pt")) + + +def load_optimizer(optimizer, checkpoint_dir: str) -> None: + # Load optimizer state via DCP, restoring _StridedShard placements transparently. + osd = optimizer.state_dict() + dcp.load({"optimizer": osd}, checkpoint_id=checkpoint_dir) + pmap_path = os.path.join(checkpoint_dir, "placement_map.pt") + if os.path.exists(pmap_path): + placement_map = torch.load(pmap_path, weights_only=False) + restore_strided_from_shard(osd, placement_map) + optimizer.load_state_dict(osd) diff --git a/src/transformers/integrations/fsdp.py b/src/transformers/integrations/fsdp.py index 526aa2780ab3..128cba7d253f 100644 --- a/src/transformers/integrations/fsdp.py +++ b/src/transformers/integrations/fsdp.py @@ -25,12 +25,8 @@ if is_torch_available() and is_torch_greater_or_equal("2.5"): import torch import torch.distributed as dist - import torch.distributed.checkpoint as dcp from torch.distributed._composable.fsdp import fully_shard - from torch.distributed.checkpoint.hf_storage import HuggingFaceStorageWriter - from torch.distributed.checkpoint.state_dict import get_model_state_dict from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy, OffloadPolicy - from torch.distributed.tensor import DTensor logger = logging.get_logger(__name__) @@ -360,9 +356,9 @@ def _get_manual_plan_modules(fsdp_plan: dict[str, Any]) -> dict[str, list[str]]: return modules -def apply_fsdp2( +def apply_fully_shard_data_parallel( model, - device_mesh, + fsdp_mesh, fsdp_plan: dict[str, Any] | str | None, ): """ @@ -398,10 +394,10 @@ def apply_fsdp2( if not is_torch_greater_or_equal("2.5"): raise OSError("FSDP2 requires torch>=2.5") - if device_mesh is None: - raise ValueError("device_mesh is required for FSDP2") + if fsdp_plan is None: + return model - if isinstance(fsdp_plan, str): + if fsdp_plan == "auto": fsdp_plan = {"mode": fsdp_plan} input_embed = getattr(model, "get_input_embeddings", lambda: None)() @@ -428,17 +424,17 @@ def apply_fsdp2( "Could not auto-detect transformer block classes for FSDP. Applying FSDP only to root module." ) else: - _auto_shard_input_embedding(input_embed, is_weights_tied, device_mesh, auto_policy_kwargs) + _auto_shard_input_embedding(input_embed, is_weights_tied, fsdp_mesh, auto_policy_kwargs) - _auto_shard_transformer_blocks(model, block_classes, device_mesh, auto_policy_kwargs) + _auto_shard_transformer_blocks(model, block_classes, fsdp_mesh, auto_policy_kwargs) tail_modules = _auto_get_tail_modules( model, decoder_layer_names, input_embed, output_embed, is_weights_tied ) - _auto_shard_tail_modules(tail_modules, device_mesh, auto_policy_kwargs) + _auto_shard_tail_modules(tail_modules, fsdp_mesh, auto_policy_kwargs) # Shard root model - fully_shard(model, mesh=device_mesh, **auto_policy_kwargs) + fully_shard(model, mesh=fsdp_mesh, **auto_policy_kwargs) logger.info( f"FSDP2 applied to model: {len(block_classes)} block type(s), {len(decoder_layer_names)} decoder layers" @@ -469,7 +465,7 @@ def apply_fsdp2( for name, module in _iter_manual_plan_targets(model, pattern, name_to_module, already_sharded_names): if name in already_sharded_names: continue - shard_kwargs = {"mesh": device_mesh, "reshard_after_forward": reshard} + shard_kwargs = {"mesh": fsdp_mesh, "reshard_after_forward": reshard} if mp_policy is not None: shard_kwargs["mp_policy"] = mp_policy if offload_policy is not None: @@ -481,7 +477,7 @@ def apply_fsdp2( # Shard root model with the same policies as sub-modules. # MixedPrecisionPolicy.output_dtype casting happens in post_forward # for every fully_shard-wrapped module, even with no direct parameters. - fully_shard(model, mesh=device_mesh, mp_policy=root_mp_policy, offload_policy=root_offload_policy) + fully_shard(model, mesh=fsdp_mesh, mp_policy=root_mp_policy, offload_policy=root_offload_policy) # Used by generation code to detect FSDP and enable synced_gpus. model._is_fsdp_managed_module = True @@ -497,36 +493,6 @@ def apply_fsdp2( return model -# TODO(3outeille): probably remove this function. Will be handled when someone tackle PEFT + FSDP. -def save_fsdp_model(model, save_directory): - """Save FSDP2 model weights as HF safetensors via DCP distributed save + consolidation. - - Each rank saves its DTensor shard in parallel, then rank 0 consolidates - into standard HF-compatible safetensors files. - """ - model_sd = get_model_state_dict(model) - - # Clone tensors sharing storage (tied weights) — safetensors refuses aliased tensors - seen_data_ptrs = {} - for key in list(model_sd.keys()): - tensor = model_sd[key] - t = tensor._local_tensor if isinstance(tensor, DTensor) else tensor - ptr = t.data_ptr() - if ptr in seen_data_ptrs: - model_sd[key] = tensor.clone() - else: - seen_data_ptrs[ptr] = key - - dcp.save( - model_sd, - storage_writer=HuggingFaceStorageWriter( - path=save_directory, - save_distributed=True, - enable_consolidation=True, - ), - ) - - # ========================= PEFT compatibility ========================= # TODO(3outeille): make sure new FSDP works with PEFT def get_fsdp_ckpt_kwargs(): diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index c7e4546017c7..759280defe8f 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -17,13 +17,16 @@ from dataclasses import dataclass from typing import Literal -from torch.distributed.tensor import Replicate, Shard +from torch.distributed.tensor import DTensor, Partial, Replicate, Shard, distribute_tensor from torch.distributed.tensor.parallel import ( ColwiseParallel, + PrepareModuleInput, RowwiseParallel, + SequenceParallel, parallelize_module, ) from torch.distributed.tensor.parallel.style import ParallelStyle +from torch.distributed.tensor.placement_types import _StridedShard from ..utils import logging from ..utils.import_utils import is_torch_available @@ -31,6 +34,7 @@ if is_torch_available(): import torch + import torch.distributed as dist # Cache this result has it's a C FFI call which can be pretty time-consuming _torch_distributed_available = torch.distributed.is_available() @@ -72,6 +76,11 @@ def _get_parameter_tp_plan(parameter_name: str, tp_plan: dict[str, str], is_weig # ============================================================================= +# ============================================================================= +# High-Level API Functions +# ============================================================================= + + def _to_cpu_fresh(tensor: torch.Tensor) -> torch.Tensor: """Plain tensor → contiguous CPU tensor with fresh storage for safetensors.""" if tensor.device.type == "meta": @@ -84,33 +93,116 @@ def _to_cpu_fresh(tensor: torch.Tensor) -> torch.Tensor: return out.contiguous() -# ============================================================================= -# High-Level API Functions -# ============================================================================= +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 ``{}``. + """ + tp_size = model.tp_size + is_rank0 = dist.get_rank() == 0 + + # Get state dict — FSDP unshard if needed (returns DTensors, not full tensors) + if getattr(model, "_is_fsdp_managed_module", False): + from torch.distributed.checkpoint.state_dict import get_model_state_dict + + state_dict = get_model_state_dict(model) + else: + state_dict = model.state_dict() + + # No TP — materialize on rank 0 only + if tp_size is None: + if is_rank0: + return {k: _to_cpu_fresh(v) for k, v in state_dict.items()} + return {} + + # Stream: gather one param at a time, only rank 0 keeps the CPU copy + result = {} + for key, tensor in state_dict.items(): + if isinstance(tensor, DTensor): + # All ranks participate in the collective, only rank 0 keeps the result + with torch.no_grad(): + full = tensor.redistribute( + placements=[Replicate()] * tensor.device_mesh.ndim, async_op=False + ).to_local() + if is_rank0: + result[key] = _to_cpu_fresh(full) + del full + elif is_rank0: + result[key] = _to_cpu_fresh(tensor) + + return result + + +def _redistribute_dtensor(tensor: DTensor, target_placements: tuple) -> DTensor: + """Redistribute a DTensor via Replicate as an intermediate step. + + PyTorch doesn't implement all placement conversions (e.g. _StridedShard↔Shard). + Going through Replicate first is always supported. + """ + with torch.no_grad(): + replicated = tensor.redistribute(placements=[Replicate()] * tensor.device_mesh.ndim) + return replicated.redistribute(placements=target_placements) + + +def convert_strided_to_shard(state_dict: dict) -> dict[str, tuple]: + # Convert _StridedShard DTensors in a state dict to plain Shard for DCP compatibility. + placement_map: dict[str, tuple] = {} + for key, value in state_dict.items(): + if isinstance(value, dict): + nested = convert_strided_to_shard(value) + for nk, nv in nested.items(): + placement_map[f"{key}.{nk}"] = nv + elif isinstance(value, DTensor) and any(isinstance(p, _StridedShard) for p in value.placements): + placement_map[key] = tuple(value.placements) + shard_placements = tuple(Shard(p.dim) if isinstance(p, _StridedShard) else p for p in value.placements) + state_dict[key] = _redistribute_dtensor(value, shard_placements) + return placement_map + + +def restore_strided_from_shard(state_dict: dict, placement_map: dict[str, tuple]) -> None: + # Restore _StridedShard placements after dcp.load. + def _resolve(d, dotted_key): + parts = dotted_key.split(".", 1) + if len(parts) == 2 and parts[0] in d and isinstance(d[parts[0]], dict): + return _resolve(d[parts[0]], parts[1]) + return d, dotted_key + + for key, original_placements in placement_map.items(): + container, leaf_key = _resolve(state_dict, key) + if leaf_key in container and isinstance(container[leaf_key], DTensor): + container[leaf_key] = _redistribute_dtensor(container[leaf_key], original_placements) def verify_tp_plan(expected_keys: list[str], tp_plan: dict[str, str | TPStyle] | None): """ Verify the TP plan of the model, log a warning if the layers that were not sharded and the rules that were not applied. - Only weight-sharding rules (colwise, rowwise, vocab) are checked. + Only weight-sharding rules (colwise, rowwise, vocab, moe_experts) are checked. + Module/activation entries (e.g. PrepareModuleInput, SequenceParallel) set up + communication hooks on modules, not weight sharding, so they are excluded. """ if tp_plan is None: return + # Filter out module-level comm hooks — they don't shard weights + _NON_WEIGHT_KINDS = {"activation", "module"} + weight_plan = {k: v for k, v in tp_plan.items() if not isinstance(v, TPStyle) or v.kind not in _NON_WEIGHT_KINDS} + generic_keys = {replace_layer_number_by_wildcard(key) for key in expected_keys} unsharded_layers = set(generic_keys) - unused_rules = tp_plan.copy() + unused_rules = weight_plan.copy() for key in generic_keys: param_name = key.rsplit(".", 1)[0] if "." in key else key generic_param_name = re.sub(r"\d+", "*", param_name) - if generic_param_name in tp_plan: + if generic_param_name in weight_plan: unused_rules.pop(generic_param_name, None) unsharded_layers.discard(key) - elif "." in generic_param_name and (parent_param_name := generic_param_name.rsplit(".", 1)[0]) in tp_plan: + elif "." in generic_param_name and (parent_param_name := generic_param_name.rsplit(".", 1)[0]) in weight_plan: unused_rules.pop(parent_param_name, None) unsharded_layers.discard(key) @@ -120,12 +212,292 @@ def verify_tp_plan(expected_keys: list[str], tp_plan: dict[str, str | TPStyle] | logger.warning(f"The following layers were not sharded: {', '.join(unsharded_layers)}") +class PrepareModuleInputOutput(ParallelStyle): + """Allgather input (Shard(1) → Replicate) + local split output (Replicate → Shard(1)). + + Used for MoE blocks with SP: the input sequence is gathered before routing, + and the output (after expert allreduce) is split back to match the residual. + Forward output split is a local op (no comm). Backward creates the all-gather. + """ + + def __init__(self, use_local_output=True): + super().__init__() + self.use_local_output = use_local_output + + def _apply(self, module, device_mesh): + def input_hook(mod, inputs): + x = inputs[0] if isinstance(inputs, tuple) else inputs + if not isinstance(x, DTensor): + x = DTensor.from_local(x, device_mesh, [Shard(1)], run_check=False) + x = x.redistribute(placements=[Replicate()]) + x = x.to_local() + return (x,) + (inputs[1:] if isinstance(inputs, tuple) else ()) + + def output_hook(mod, inputs, output): + if not isinstance(output, DTensor): + output = DTensor.from_local(output, device_mesh, [Replicate()], run_check=False) + output = output.redistribute(placements=[Shard(1)]) + return output.to_local() + + module.register_forward_pre_hook(input_hook) + module.register_forward_hook(output_hook) + return module + + +class PackedColwiseParallel(ParallelStyle): + """Column-wise parallel style for fused linear weights packed along the output dimension.""" + + def __init__( + self, + *, + input_layouts=None, + use_local_output: bool = True, + split_factor: int = 2, + ): + super().__init__() + self.input_layouts = (input_layouts or Replicate(),) + self.use_local_output = use_local_output + self.split_factor = split_factor + + def _partition_linear_fn(self, module, device_mesh): + if getattr(module, "weight", None) is None: + return + + packed_shard = _StridedShard(dim=0, split_factor=self.split_factor) + module.register_parameter( + "weight", + torch.nn.Parameter( + distribute_tensor(module.weight, device_mesh, [packed_shard], src_data_rank=self.src_data_rank), + requires_grad=module.weight.requires_grad, + ), + ) + + if getattr(module, "bias", None) is not None: + module.register_parameter( + "bias", + torch.nn.Parameter( + distribute_tensor(module.bias, device_mesh, [packed_shard], src_data_rank=self.src_data_rank), + requires_grad=module.bias.requires_grad, + ), + ) + + def _prepare_input_fn(self, mod, inputs, device_mesh): + input_tensor = inputs[0] + if not isinstance(input_tensor, DTensor): + input_tensor = DTensor.from_local(input_tensor, device_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() + + local_param_shadows = {} + for param_name, param in list(mod.named_parameters(recurse=False)): + if isinstance(param, DTensor): + local_param_shadows[param_name] = param + mod._parameters.pop(param_name) + setattr(mod, param_name, param.to_local()) + if local_param_shadows: + shadow_stack = getattr(mod, "_packed_local_param_shadows", None) + if shadow_stack is None: + shadow_stack = [] + mod._packed_local_param_shadows = shadow_stack + shadow_stack.append(local_param_shadows) + return (input_tensor,) + inputs[1:] + + def _prepare_output_fn(self, mod, outputs, device_mesh): + shadow_stack = getattr(mod, "_packed_local_param_shadows", None) + if shadow_stack: + for param_name, param in shadow_stack.pop().items(): + if hasattr(mod, param_name): + delattr(mod, param_name) + mod.register_parameter(param_name, param) + + if outputs is None or self.use_local_output: + return outputs + return DTensor.from_local( + outputs, device_mesh, (_StridedShard(dim=-1, split_factor=self.split_factor),), run_check=False + ) + + def _apply(self, module, device_mesh): + if not isinstance(module, torch.nn.Linear): + raise NotImplementedError("PackedColwiseParallel currently only supports nn.Linear!") + + self._partition_linear_fn(module, device_mesh) + module.register_forward_pre_hook(lambda mod, inputs: self._prepare_input_fn(mod, inputs, device_mesh)) + module.register_forward_hook( + lambda mod, inputs, outputs: self._prepare_output_fn(mod, outputs, device_mesh), + always_call=True, + ) + return module + + def __repr__(self) -> str: + tmpstr = self.__class__.__name__ + "(" + tmpstr += f"input_layouts={self.input_layouts}, " + tmpstr += f"use_local_output={self.use_local_output}, " + tmpstr += f"split_factor={self.split_factor}" + tmpstr += ")" + return tmpstr + + +# Maps string tp_plan entries for MoE experts to DTensor placements. +# Used by MoEExpertsParallel._partition_fn to create DTensors from the config plan. +_STRING_TO_PLACEMENT = { + "packed_colwise": lambda: _StridedShard(dim=-2, split_factor=2), + "colwise": lambda: Shard(-2), + "rowwise": lambda: Shard(-1), +} + + +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(ParallelStyle): + """Hybrid parallel style for MoE expert modules. + + Converts expert weights to DTensors based on the ``shard_plan`` (e.g. + ``{"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}``). + Communication uses DTensor ``from_local``/``to_local`` on activations only — + compatible with ``grouped_mm``. + """ + + def __init__(self, output_layouts=None): + super().__init__() + self.output_layouts = output_layouts or Replicate() + self._moe_shard_plan: dict[str, str] = {} + + @staticmethod + def _partition_fn(name, module, device_mesh, shard_plan): + for param_name, param in module.named_parameters(recurse=False): + plan_str = shard_plan.get(param_name) + if plan_str is None: + continue + placement_fn = _STRING_TO_PLACEMENT.get(plan_str) + if placement_fn is None: + continue + placement = placement_fn() + dtensor = distribute_tensor(param.data, device_mesh, [placement]) + module._parameters[param_name] = torch.nn.Parameter(dtensor, requires_grad=param.requires_grad) + + @staticmethod + def _uses_partial_outputs(mod) -> bool: + cached = getattr(mod, "_moe_outputs_are_partial", None) + if cached is not None: + return cached + + # Under TP-only the expert MLP dimension is sharded, so each rank emits a + # partial hidden-state contribution that must be reduced. Under TP+FSDP, + # FSDP can swap in full gathered expert weights for the current rank's + # forward, in which case the local output is already complete. + if hasattr(mod, "gate_up_proj"): + gate_up_proj = mod.gate_up_proj.to_local() if isinstance(mod.gate_up_proj, DTensor) else mod.gate_up_proj + full_expert_out = 2 * mod.intermediate_dim + sharded_dim = -1 if getattr(mod, "is_transposed", False) else -2 + cached = gate_up_proj.shape[sharded_dim] != full_expert_out + elif hasattr(mod, "up_proj"): + up_proj = mod.up_proj.to_local() if isinstance(mod.up_proj, DTensor) else mod.up_proj + full_expert_out = mod.intermediate_dim + sharded_dim = -1 if getattr(mod, "is_transposed", False) else -2 + cached = up_proj.shape[sharded_dim] != full_expert_out + else: + cached = True + + mod._moe_outputs_are_partial = cached + return cached + + @staticmethod + def _prepare_input_fn(mod, inputs, device_mesh): + hidden_states, top_k_index, top_k_weights = inputs[0], inputs[1], inputs[2] + # from_local([Replicate()]).to_local(): forward sees plain tensor, + # backward graph goes through DTensor all-reduce on gradient. + if not isinstance(hidden_states, DTensor): + hidden_states = DTensor.from_local(hidden_states, device_mesh, [Replicate()], run_check=False) + hidden_states = hidden_states.to_local() + # Route weights are replicated (same on all ranks), but their backward + # gradient is partial (each rank's contribution from its expert shard). + # Use allreduce-sum (not Replicate's allreduce-then-divide) to aggregate. + tp_group = device_mesh.get_group() if device_mesh.ndim == 1 else device_mesh.get_group("tp") + if isinstance(top_k_weights, DTensor): + top_k_weights = top_k_weights.to_local() + top_k_weights = _AllReduceBackward.apply(top_k_weights, tp_group) + local_param_shadows = {} + for param_name, param in list(mod.named_parameters(recurse=False)): + if isinstance(param, DTensor): + # grouped_mm expects plain tensors, but we must restore the + # original DTensor params after the forward so save_pretrained + # still sees the canonical sharded weights. + local_param_shadows[param_name] = param + mod._parameters.pop(param_name) + setattr(mod, param_name, param.to_local()) + if local_param_shadows: + shadow_stack = getattr(mod, "_moe_local_param_shadows", None) + if shadow_stack is None: + shadow_stack = [] + mod._moe_local_param_shadows = shadow_stack + shadow_stack.append(local_param_shadows) + return (hidden_states, top_k_index, top_k_weights) + + @staticmethod + def _prepare_output_fn(output_layouts, mod, outputs, device_mesh): + shadow_stack = getattr(mod, "_moe_local_param_shadows", None) + if shadow_stack: + for param_name, param in shadow_stack.pop().items(): + if hasattr(mod, param_name): + delattr(mod, param_name) + mod.register_parameter(param_name, param) + if outputs is None: + return None + # Plain TP expert weights produce partial outputs that need an all-reduce. + # TP+FSDP can leave experts replicated across TP and sharded only across + # experts/FSDP, in which case the local output is already complete. + source_layout = Partial() if MoEExpertsParallel._uses_partial_outputs(mod) else Replicate() + if not isinstance(outputs, DTensor): + outputs = DTensor.from_local(outputs, device_mesh, [source_layout], run_check=False) + # MoE experts output 2D [num_tokens, hidden]. For SP reduce-scatter, + # Shard(1) means sequence dim in 3D, but in 2D the token dim is 0. + actual_layouts = output_layouts + if outputs.dim() == 2 and isinstance(output_layouts, Shard) and output_layouts.dim == 1: + actual_layouts = Shard(0) + if outputs.placements != (actual_layouts,): + outputs = outputs.redistribute(placements=(actual_layouts,)) + return outputs.to_local() + + def _apply(self, module, device_mesh): + # Don't use PyTorch's distribute_module — it would auto-convert all + # params to Replicate DTensors. We create DTensors with proper Shard + # placements in _partition_fn instead, and register hooks manually. + self._partition_fn(module.__class__.__name__, module, device_mesh, self._moe_shard_plan) + module.register_forward_pre_hook(lambda mod, inputs: self._prepare_input_fn(mod, inputs, device_mesh)) + module.register_forward_hook( + lambda mod, inputs, outputs: self._prepare_output_fn(self.output_layouts, mod, outputs, device_mesh), + always_call=True, + ) + return module + + @dataclass(frozen=True) class TPStyle: - kind: Literal["colwise", "rowwise", "vocab"] - comm: Literal["none", "allreduce", "reduce_scatter"] + kind: Literal["colwise", "packed_colwise", "rowwise", "vocab", "activation", "module", "moe_experts"] + comm: Literal["none", "allreduce", "reduce_scatter", "allgather", "allgather_split", "loss_parallel"] sequence_dim: int = 1 use_local_output: bool = True + input_key: str | None = None + shard_plan: dict[str, str] | None = None def to_dtensor_style(self) -> ParallelStyle: """Convert to the corresponding PyTorch DTensor ParallelStyle.""" @@ -135,6 +507,18 @@ def to_dtensor_style(self) -> ParallelStyle: return ColwiseParallel( input_layouts=Replicate(), output_layouts=Shard(-1), use_local_output=self.use_local_output ) + case "allgather": + return ColwiseParallel( + input_layouts=Replicate(), + output_layouts=Replicate(), + use_local_output=self.use_local_output, + ) + case "loss_parallel": + return ColwiseParallel(input_layouts=Shard(1), output_layouts=Shard(-1), use_local_output=False) + elif self.kind == "packed_colwise": + match self.comm: + case "none": + return PackedColwiseParallel(input_layouts=Replicate(), use_local_output=self.use_local_output) elif self.kind == "rowwise": match self.comm: case "allreduce": @@ -161,11 +545,41 @@ def to_dtensor_style(self) -> ParallelStyle: return RowwiseParallel( input_layouts=Replicate(), output_layouts=Shard(1), use_local_output=self.use_local_output ) + elif self.kind == "activation": + match self.comm: + case "none": + return SequenceParallel(sequence_dim=self.sequence_dim, use_local_output=self.use_local_output) + elif self.kind == "module": + match self.comm: + case "allgather": + if self.input_key is not None: + return PrepareModuleInput( + input_kwarg_layouts={self.input_key: Shard(1)}, + desired_input_kwarg_layouts={self.input_key: Replicate()}, + use_local_output=self.use_local_output, + ) + return PrepareModuleInput( + input_layouts=(Shard(1),), + desired_input_layouts=(Replicate(),), + use_local_output=self.use_local_output, + ) + case "allgather_split": + return PrepareModuleInputOutput(use_local_output=self.use_local_output) + elif self.kind == "moe_experts": + match self.comm: + case "allreduce": + return MoEExpertsParallel(output_layouts=Replicate()) + case "reduce_scatter": + return MoEExpertsParallel(output_layouts=Shard(1)) raise ValueError( f"Invalid TPStyle({self.kind!r}, {self.comm!r}). Valid combinations:\n" - f" colwise: none\n" + f" colwise: none, allgather, loss_parallel\n" + f" packed_colwise: none\n" f" rowwise: allreduce, reduce_scatter\n" - f" vocab: allreduce, reduce_scatter" + f" vocab: allreduce, reduce_scatter\n" + f" activation: none\n" + f" module: allgather, allgather_split\n" + f" moe_experts: allreduce, reduce_scatter" ) def __str__(self): @@ -178,13 +592,17 @@ def apply_tensor_parallel(model, tp_mesh, tp_plan): """Apply tensor parallelism using PyTorch's parallelize_module. Converts the wildcard tp_plan from model config into a concrete plan - for ``parallelize_module``. Plan values are ``TPStyle`` instances. + for ``parallelize_module``. Plan values is a `TPStyle`` instances """ if tp_plan is None: return model if tp_plan == "auto": - base_plan = model.config.base_model_tp_plan or {} + enable_sp = getattr(getattr(model.config, "distributed_config", None), "enable_sequence_parallel", False) + if enable_sp and hasattr(model.config, "base_model_sp_plan"): + base_plan = model.config.base_model_sp_plan + else: + base_plan = model.config.base_model_tp_plan or {} # Prefix base model keys (e.g. "layers.*.q_proj" → "model.layers.*.q_proj") # Top-level keys like "lm_head" are kept as-is. @@ -204,9 +622,42 @@ def apply_tensor_parallel(model, tp_mesh, tp_plan): if isinstance(style_value, TPStyle): dtensor_style = style_value.to_dtensor_style() parallelize_plan[name] = dtensor_style + # For MoE modules, attach the per-parameter shard plan from TPStyle + # so _partition_fn can create DTensors with the correct placements. + if isinstance(dtensor_style, MoEExpertsParallel) and style_value.shard_plan: + dtensor_style._moe_shard_plan = style_value.shard_plan else: parallelize_plan[name] = style_value parallelize_module(model, tp_mesh, parallelize_plan) + # 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(isinstance(v, TPStyle) and v.comm == "loss_parallel" for v in tp_plan.values()) + if has_loss_parallel: + from torch.distributed.tensor.parallel import loss_parallel + + model._loss_parallel_ctx = loss_parallel() + model._loss_parallel_ctx.__enter__() + return model diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index b6feec6f7ba6..082b294fb41f 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -51,7 +51,7 @@ revert_weight_conversion, ) from .distributed import DistributedConfig -from .distributed.utils import is_fsdp_enabled +from .distributed.utils import init_device_mesh, is_fsdp_enabled 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 @@ -69,18 +69,15 @@ from .integrations.flash_attention import flash_attention_forward from .integrations.flash_paged import paged_attention_forward from .integrations.flex_attention import flex_attention_forward -from .integrations.fsdp import initialize_fsdp +from .integrations.fsdp import apply_fully_shard_data_parallel from .integrations.hub_kernels import allow_all_hub_kernels, is_kernel 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, + apply_tensor_parallel, + gather_full_state_dict, verify_tp_plan, ) from .loss.loss_utils import LOSS_MAPPING @@ -139,8 +136,6 @@ from accelerate.utils import extract_model_from_parallel -_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 @@ -1352,14 +1347,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()] @@ -1928,11 +1915,10 @@ def get_correct_attn_implementation(self, requested_attention: str | None, is_in def get_correct_experts_implementation(self, requested_experts: str | None) -> str: applicable_experts = "grouped_mm" if requested_experts is None else requested_experts - if applicable_experts not in ["eager", "grouped_mm", "batched_mm", "deepgemm"]: + if applicable_experts not in ["eager", "grouped_mm", "batched_mm"]: message = ( f'Specified `experts_implementation="{applicable_experts}"` is not supported. The only possible arguments are ' - '`experts_implementation="eager"`, `"experts_implementation=grouped_mm"`, `"experts_implementation=batched_mm"` ' - 'and `"experts_implementation=deepgemm"`.' + '`experts_implementation="eager"`, `"experts_implementation=grouped_mm"` and `"experts_implementation=batched_mm"`.' ) raise ValueError(message) @@ -2987,7 +2973,6 @@ def _get_resized_lm_head( new_lm_head, old_lm_head, num_tokens_to_copy, transposed, has_new_lm_head_bias ) - new_lm_head._is_hf_initialized = True return new_lm_head def _init_added_embeddings_weights_with_mean( @@ -3246,7 +3231,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." ) @@ -3285,64 +3270,58 @@ def save_pretrained( if self._auto_class is not None: 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 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) - 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] - - current_peft_config = self.peft_config[active_adapter] - current_peft_config.save_pretrained(save_directory) - - # FSDP2 models: use DCP distributed save + consolidation for safetensors. - # All ranks must call this collectively. Config/generation_config are - # already saved above (guarded by is_main_process). - if getattr(self, "_is_fsdp_managed_module", False): - from .integrations.fsdp import save_fsdp_model - - save_fsdp_model(model_to_save, save_directory) - - if push_to_hub: - 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, - ) - return + active_adapter = self.active_adapters() + + if len(active_adapter) > 1: + raise ValueError( + "Multiple active adapters detected, saving multiple active adapters is not supported yet. You can save adapters separately one by one " + "by iteratively calling `model.set_adapter(adapter_name)` then `model.save_pretrained(...)`" + ) + active_adapter = active_adapter[0] + + current_peft_config = self.peft_config[active_adapter] + current_peft_config.save_pretrained(save_directory) + finally: + if distributed_config is not None: + model_to_save.config.distributed_config = distributed_config - # Get the model state_dict + # Get the model state_dict (handles FSDP unshard + TP gather in one call) 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) + else: + state_dict = model_to_save.state_dict() # if any model parameters are offloaded, we need to know it for later is_offloaded = False @@ -3368,10 +3347,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) @@ -3642,10 +3617,7 @@ def get_init_context( elif is_quantized: init_contexts.extend([torch.device("meta"), set_quantized_state()]) else: - # meta_device_safe_creation_ops patches torch.linspace to default to CPU - # so that custom models calling .item() during __init__ (e.g. drop-path - # schedules) don't crash on meta tensors. - init_contexts.extend([torch.device("meta"), init.meta_device_safe_creation_ops()]) + init_contexts.append(torch.device("meta")) return init_contexts @@ -3867,13 +3839,22 @@ def from_pretrained( max_memory (`Dict`, *optional*): A dictionary device identifier to maximum memory if using `device_map`. Will default to the maximum memory available for each GPU and the available CPU RAM if unset. - tp_plan (`Optional[Union[dict, str]]`, *optional*): - A torch tensor parallel plan, see [here](https://pytorch.org/tutorials/intermediate/TP_tutorial.html). Use `tp_plan="auto"` to - use the predefined plan based on the model. If it's a dict, then it should match between module names and desired layout. - Note that if you use it, you should launch your script accordingly with `torchrun [args] script.py`. This will be much - faster than using a `device_map`, but has limitations. - tp_size (`str`, *optional*): - A torch tensor parallel degree. If not provided would default to world size. + distributed_config ([`DistributedConfig`], *optional*): + Configuration for native distributed training (FSDP2 + TP) via `torch.distributed`. Mutually + exclusive with `quantization_config` (for now) and `device_map`. When set, accelerate is not used for + device placement or dispatch. Launch with `torchrun --nproc_per_node=N script.py`. + + Accepts `tp_size`, `tp_plan`, `fsdp_size`, `fsdp_plan`. When a size is specified without a + plan, the plan defaults to `"auto"`. `tp_plan="auto"` uses the model's predefined tensor + parallel sharding plan. `fsdp_plan="auto"` wraps each transformer layer individually with + FSDP2 (`fully_shard`). Both plans also accept a `dict` for manual control: `tp_plan` maps + parameter names to parallel styles (e.g. `{"model.layers.*.self_attn.q_proj": "colwise"}`), + `fsdp_plan` maps module names to wrap (e.g. `{"model.layers.0": {}, "model.layers.1": {}}`). + + Examples: + - TP-only: `DistributedConfig(tp_size=4)` + - FSDP-only: `DistributedConfig(fsdp_size=4)` + - 2D parallel: `DistributedConfig(tp_size=2, fsdp_size=2)` on 4 GPUs device_mesh (`torch.distributed.DeviceMesh`, *optional*): A torch device mesh. If not provided would default to world size. Used only for tensor parallel for now. If provided, it has to contain dimension named `"tp"` in case it's > 1 dimensional, this dimension will be used for tensor parallelism @@ -3954,9 +3935,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) - fsdp_plan = kwargs.pop("fsdp_plan", 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) @@ -3965,8 +3943,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"]: @@ -3997,27 +3984,18 @@ def from_pretrained( "`state_dict` cannot be passed together with a model name or a `gguf_file`. Use one of the two loading strategies." ) - if device_map == "auto" and int(os.environ.get("WORLD_SIZE", "0")): - logger.info( - "You've set device_map=`auto` while triggering a distributed run with torchrun. This might lead to unexpected behavior. " - "If your plan is to load the model on each device, you should set device_map={" - ": PartialState().process_index} where PartialState comes from accelerate library" - ) - - if fsdp_plan is not None and (tp_plan is not None or tp_size is not None): - raise ValueError("Combining `fsdp_plan` with tensor parallel loading is not supported yet.") + if distributed_config is not None: + device_mesh = init_device_mesh(distributed_config) + else: + # Accelerate path + if device_map == "auto" and int(os.environ.get("WORLD_SIZE", "0")): + logger.info( + "You've set device_map=`auto` while triggering a distributed run with torchrun. This might lead to unexpected behavior. " + "If your plan is to load the model on each device, you should set device_map={" + ": PartialState().process_index} where PartialState comes from accelerate library" + ) - if tp_plan is not None or tp_size is not None: # TP warnings, and setup - device_map, device_mesh, tp_size = initialize_tensor_parallelism( - tp_plan, tp_size=tp_size, device_mesh=device_mesh, device_map=device_map - ) - - if fsdp_plan is not None: - device_map, device_mesh, _ = initialize_fsdp( - fsdp_plan=fsdp_plan, - 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`.") @@ -4030,7 +4008,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: @@ -4140,20 +4117,28 @@ 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 and (tp_plan is not None or fsdp_plan is not None): - model = distribute_model(model, tp_plan, distributed_config, device_mesh, tp_size, fsdp_plan=fsdp_plan) + if distributed_config is not None: + model.config.distributed_config = distributed_config + model.device_mesh = device_mesh + + def sub_mesh(name): + return device_mesh[name] if device_mesh.ndim > 1 else device_mesh - # Prepare the full device map - if isinstance(device_map, dict): - device_map = _get_device_map(model, device_map, max_memory, hf_quantizer) - elif device_map is not None: - device_map = {"": device_map} + mesh_dim_names = device_mesh.mesh_dim_names or () + if "tp" in mesh_dim_names: + model = apply_tensor_parallel(model, sub_mesh("tp"), distributed_config.tp_plan) + if "fsdp" in mesh_dim_names: + model = apply_fully_shard_data_parallel(model, sub_mesh("fsdp"), distributed_config.fsdp_plan) + else: + # Accelerate path: auto device mapping + if device_map is not None: + device_map = _get_device_map(model, device_map, max_memory, hf_quantizer) # Finalize model weight initialization - active_tp_plan = getattr(model, "_tp_plan", None) if tp_size is not None else None + active_tp_plan = getattr(model, "_tp_plan", None) if getattr(distributed_config, "tp_plan", None) else None load_config = LoadStateDictConfig( pretrained_model_name_or_path=pretrained_model_name_or_path, ignore_mismatched_sizes=ignore_mismatched_sizes, @@ -4447,8 +4432,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): @@ -4555,10 +4540,10 @@ def _move_missing_keys_from_meta_to_device( # In this case we need to move everything back if is_fsdp_enabled() and not is_local_dist_rank_0() and not is_quantized: for key, param in self.named_parameters(): - value = torch.zeros_like(param, device="cpu") + value = torch.empty_like(param, device="cpu") _load_parameter_into_model(self, key, value) for key, buffer in self.named_buffers(): - value = torch.zeros_like(buffer, device="cpu") + value = torch.empty_like(buffer, device="cpu") _load_parameter_into_model(self, key, value) return @@ -4567,15 +4552,29 @@ 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 + from torch.distributed.tensor import DTensor + + 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()), + ) + new_dtensor = DTensor.from_local( + local_value, + param.device_mesh, + param.placements, + run_check=False, + shape=param.shape, + stride=tuple(param.stride()), ) - # Otherwise, just move it to device + 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(): @@ -4654,7 +4653,7 @@ def mark_tied_weights_as_initialized(self, loading_info): later as they will be tied (overwritten) anyway. This is very important as most embeddings are tied, and they are huge params (vocabularies are often 256k), so running inits on them is very costly.""" - for tied_param in getattr(self, "all_tied_weights_keys", {}).keys(): + for tied_param in self.all_tied_weights_keys.keys(): param = self.get_parameter(tied_param) param._is_hf_initialized = True diff --git a/src/transformers/models/afmoe/modeling_afmoe.py b/src/transformers/models/afmoe/modeling_afmoe.py index 421119b33deb..819c30446755 100644 --- a/src/transformers/models/afmoe/modeling_afmoe.py +++ b/src/transformers/models/afmoe/modeling_afmoe.py @@ -23,6 +23,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -277,6 +278,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/apertus/modeling_apertus.py b/src/transformers/models/apertus/modeling_apertus.py index 7d14dd3d14c8..e616764f8fc8 100644 --- a/src/transformers/models/apertus/modeling_apertus.py +++ b/src/transformers/models/apertus/modeling_apertus.py @@ -23,6 +23,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2CLS, ACT2FN from ...cache_utils import Cache, DynamicCache @@ -170,6 +171,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/arcee/modeling_arcee.py b/src/transformers/models/arcee/modeling_arcee.py index 8d2d05bf2952..a30a2e03642b 100644 --- a/src/transformers/models/arcee/modeling_arcee.py +++ b/src/transformers/models/arcee/modeling_arcee.py @@ -23,6 +23,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from transformers.utils import auto_docstring @@ -175,6 +176,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/aria/modeling_aria.py b/src/transformers/models/aria/modeling_aria.py index e66b12438940..715c721b0b39 100644 --- a/src/transformers/models/aria/modeling_aria.py +++ b/src/transformers/models/aria/modeling_aria.py @@ -23,6 +23,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -404,6 +405,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/bitnet/modeling_bitnet.py b/src/transformers/models/bitnet/modeling_bitnet.py index 14c1581b250f..78ae2a49b77b 100644 --- a/src/transformers/models/bitnet/modeling_bitnet.py +++ b/src/transformers/models/bitnet/modeling_bitnet.py @@ -22,6 +22,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache @@ -106,6 +107,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/blt/modeling_blt.py b/src/transformers/models/blt/modeling_blt.py index 778f7ba80cf6..5fa5262c78c4 100644 --- a/src/transformers/models/blt/modeling_blt.py +++ b/src/transformers/models/blt/modeling_blt.py @@ -25,6 +25,7 @@ import torch.distributions import torch.nn as nn import torch.nn.functional as F +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -286,6 +287,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/chameleon/modeling_chameleon.py b/src/transformers/models/chameleon/modeling_chameleon.py index af69779959e4..cdc56f7e0458 100644 --- a/src/transformers/models/chameleon/modeling_chameleon.py +++ b/src/transformers/models/chameleon/modeling_chameleon.py @@ -180,6 +180,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/csm/modeling_csm.py b/src/transformers/models/csm/modeling_csm.py index eb78dca8faf5..2f03aeaf66ea 100644 --- a/src/transformers/models/csm/modeling_csm.py +++ b/src/transformers/models/csm/modeling_csm.py @@ -24,6 +24,7 @@ import torch import torch.nn as nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -226,6 +227,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/cwm/configuration_cwm.py b/src/transformers/models/cwm/configuration_cwm.py index ecc3743da19d..a8ea587eb41f 100644 --- a/src/transformers/models/cwm/configuration_cwm.py +++ b/src/transformers/models/cwm/configuration_cwm.py @@ -22,6 +22,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...utils import auto_docstring @@ -46,13 +47,13 @@ class CwmConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `CwmModel` 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", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/cwm/modeling_cwm.py b/src/transformers/models/cwm/modeling_cwm.py index 3e0eb0504be0..e6f98a0ce250 100644 --- a/src/transformers/models/cwm/modeling_cwm.py +++ b/src/transformers/models/cwm/modeling_cwm.py @@ -24,6 +24,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache @@ -135,6 +136,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/dbrx/modeling_dbrx.py b/src/transformers/models/dbrx/modeling_dbrx.py index 58735fb55c0b..0efbc1db2f9f 100644 --- a/src/transformers/models/dbrx/modeling_dbrx.py +++ b/src/transformers/models/dbrx/modeling_dbrx.py @@ -23,6 +23,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -134,6 +135,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index ab998cc99c21..96c958b7d552 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -11,6 +11,7 @@ import torch import torch.nn.functional as F from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -275,6 +276,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/dia/modeling_dia.py b/src/transformers/models/dia/modeling_dia.py index 629dfd4cdb35..4eba95828b53 100644 --- a/src/transformers/models/dia/modeling_dia.py +++ b/src/transformers/models/dia/modeling_dia.py @@ -23,6 +23,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -229,6 +230,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/diffllama/modeling_diffllama.py b/src/transformers/models/diffllama/modeling_diffllama.py index d80ccd572dc3..8e5c98875c38 100644 --- a/src/transformers/models/diffllama/modeling_diffllama.py +++ b/src/transformers/models/diffllama/modeling_diffllama.py @@ -26,6 +26,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -162,6 +163,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/doge/modeling_doge.py b/src/transformers/models/doge/modeling_doge.py index 4aad59b52a9a..31157deba6d5 100644 --- a/src/transformers/models/doge/modeling_doge.py +++ b/src/transformers/models/doge/modeling_doge.py @@ -27,6 +27,7 @@ import torch import torch.nn.functional as F from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -164,6 +165,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/dots1/modeling_dots1.py b/src/transformers/models/dots1/modeling_dots1.py index 399194648663..89561bfa4ed8 100644 --- a/src/transformers/models/dots1/modeling_dots1.py +++ b/src/transformers/models/dots1/modeling_dots1.py @@ -23,6 +23,7 @@ import torch import torch.nn.functional as F from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -161,6 +162,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/emu3/modeling_emu3.py b/src/transformers/models/emu3/modeling_emu3.py index 2481decd7aeb..584e32e505a5 100644 --- a/src/transformers/models/emu3/modeling_emu3.py +++ b/src/transformers/models/emu3/modeling_emu3.py @@ -28,6 +28,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -85,6 +86,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/eurobert/modeling_eurobert.py b/src/transformers/models/eurobert/modeling_eurobert.py index b93dd0649f14..9fef764976ca 100644 --- a/src/transformers/models/eurobert/modeling_eurobert.py +++ b/src/transformers/models/eurobert/modeling_eurobert.py @@ -24,6 +24,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss from ...activations import ACT2FN @@ -90,6 +91,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/exaone4/modeling_exaone4.py b/src/transformers/models/exaone4/modeling_exaone4.py index fab10b9b6937..9100d0ffefec 100644 --- a/src/transformers/models/exaone4/modeling_exaone4.py +++ b/src/transformers/models/exaone4/modeling_exaone4.py @@ -24,6 +24,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache @@ -160,6 +161,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/exaone_moe/modeling_exaone_moe.py b/src/transformers/models/exaone_moe/modeling_exaone_moe.py index 2836a3c2245d..cad1ee6be2d2 100644 --- a/src/transformers/models/exaone_moe/modeling_exaone_moe.py +++ b/src/transformers/models/exaone_moe/modeling_exaone_moe.py @@ -25,6 +25,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -92,6 +93,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/falcon/modeling_falcon.py b/src/transformers/models/falcon/modeling_falcon.py index 016b3209b6b1..26336cc7e674 100644 --- a/src/transformers/models/falcon/modeling_falcon.py +++ b/src/transformers/models/falcon/modeling_falcon.py @@ -94,6 +94,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/falcon_h1/modeling_falcon_h1.py b/src/transformers/models/falcon_h1/modeling_falcon_h1.py index 37b5da9df4b3..33a8de337953 100644 --- a/src/transformers/models/falcon_h1/modeling_falcon_h1.py +++ b/src/transformers/models/falcon_h1/modeling_falcon_h1.py @@ -29,6 +29,7 @@ import torch import torch.nn.functional as F from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -146,6 +147,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/gemma/modeling_gemma.py b/src/transformers/models/gemma/modeling_gemma.py index c6c5a55b8790..54be043b2eef 100644 --- a/src/transformers/models/gemma/modeling_gemma.py +++ b/src/transformers/models/gemma/modeling_gemma.py @@ -25,6 +25,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -190,6 +191,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/gemma2/modeling_gemma2.py b/src/transformers/models/gemma2/modeling_gemma2.py index 20673571b2d2..7c99443931c6 100644 --- a/src/transformers/models/gemma2/modeling_gemma2.py +++ b/src/transformers/models/gemma2/modeling_gemma2.py @@ -23,6 +23,7 @@ import torch import torch.nn as nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -175,6 +176,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/gemma3/modeling_gemma3.py b/src/transformers/models/gemma3/modeling_gemma3.py index 3ecd6344dc07..ef8d1c885bca 100644 --- a/src/transformers/models/gemma3/modeling_gemma3.py +++ b/src/transformers/models/gemma3/modeling_gemma3.py @@ -24,6 +24,7 @@ import torch import torch.nn as nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -258,6 +259,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed 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 d59fd2ab996e..b4786fb7f23e 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 @@ -26,6 +26,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -140,6 +141,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/gpt_neox_japanese/modeling_gpt_neox_japanese.py b/src/transformers/models/gpt_neox_japanese/modeling_gpt_neox_japanese.py index e334ce023d67..be0008400097 100755 --- a/src/transformers/models/gpt_neox_japanese/modeling_gpt_neox_japanese.py +++ b/src/transformers/models/gpt_neox_japanese/modeling_gpt_neox_japanese.py @@ -147,6 +147,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/granite/modeling_granite.py b/src/transformers/models/granite/modeling_granite.py index 934345fe6723..c0c7765dcb20 100644 --- a/src/transformers/models/granite/modeling_granite.py +++ b/src/transformers/models/granite/modeling_granite.py @@ -24,6 +24,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache @@ -69,6 +70,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/granitemoe/modeling_granitemoe.py b/src/transformers/models/granitemoe/modeling_granitemoe.py index 5fb53d6afe49..ce188272d5c0 100644 --- a/src/transformers/models/granitemoe/modeling_granitemoe.py +++ b/src/transformers/models/granitemoe/modeling_granitemoe.py @@ -24,6 +24,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from torch.nn import functional as F from ... import initialization as init @@ -294,6 +295,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py index dadffaea0072..a55da168936b 100644 --- a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py +++ b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py @@ -23,6 +23,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from torch.nn import functional as F from ... import initialization as init @@ -75,6 +76,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py b/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py index 71f8c6eaff7d..8f152f18f99f 100644 --- a/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py +++ b/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py @@ -23,6 +23,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from torch.nn import functional as F from ... import initialization as init @@ -282,6 +283,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed 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..ca3d81d225be 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 @@ -22,6 +22,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring from ...utils.type_validators import interval @@ -59,13 +60,13 @@ class HiggsAudioV2Config(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `HiggsAudioV2Model` 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", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/higgs_audio_v2/modeling_higgs_audio_v2.py b/src/transformers/models/higgs_audio_v2/modeling_higgs_audio_v2.py index a0f106167721..03534ee50ad5 100644 --- a/src/transformers/models/higgs_audio_v2/modeling_higgs_audio_v2.py +++ b/src/transformers/models/higgs_audio_v2/modeling_higgs_audio_v2.py @@ -24,6 +24,7 @@ import torch import torch.nn as nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -110,6 +111,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed 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..bf924c636482 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 @@ -23,6 +23,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from transformers.cache_utils import Cache @@ -109,6 +110,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed 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..0c61827781c9 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 @@ -24,6 +24,7 @@ import torch import torch.nn.functional as F from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -112,6 +113,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/jais2/modeling_jais2.py b/src/transformers/models/jais2/modeling_jais2.py index 5e6a37c0172d..4f303e45be9d 100644 --- a/src/transformers/models/jais2/modeling_jais2.py +++ b/src/transformers/models/jais2/modeling_jais2.py @@ -24,6 +24,7 @@ import torch import torch.nn as nn +from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache @@ -83,6 +84,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/jamba/modeling_jamba.py b/src/transformers/models/jamba/modeling_jamba.py index ae618fb4a2b3..b98efc90c175 100755 --- a/src/transformers/models/jamba/modeling_jamba.py +++ b/src/transformers/models/jamba/modeling_jamba.py @@ -26,6 +26,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -102,6 +103,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/jetmoe/modeling_jetmoe.py b/src/transformers/models/jetmoe/modeling_jetmoe.py index d3ee0bb14875..10de88e75f26 100644 --- a/src/transformers/models/jetmoe/modeling_jetmoe.py +++ b/src/transformers/models/jetmoe/modeling_jetmoe.py @@ -23,6 +23,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from torch.nn import functional as F from ... import initialization as init @@ -387,6 +388,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/jina_embeddings_v3/modeling_jina_embeddings_v3.py b/src/transformers/models/jina_embeddings_v3/modeling_jina_embeddings_v3.py index a55ffe0151c3..163fc9e157e9 100644 --- a/src/transformers/models/jina_embeddings_v3/modeling_jina_embeddings_v3.py +++ b/src/transformers/models/jina_embeddings_v3/modeling_jina_embeddings_v3.py @@ -24,6 +24,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss from ... import initialization as init @@ -190,6 +191,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed 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..b60706bacbdd 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 @@ -25,6 +25,7 @@ import torch import torch.nn as nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -386,6 +387,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/lasr/modeling_lasr.py b/src/transformers/models/lasr/modeling_lasr.py index 7ecea9099410..eb9a2742b62f 100644 --- a/src/transformers/models/lasr/modeling_lasr.py +++ b/src/transformers/models/lasr/modeling_lasr.py @@ -24,6 +24,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...integrations import use_kernel_func_from_hub, use_kernelized_func @@ -160,6 +161,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/lfm2/modeling_lfm2.py b/src/transformers/models/lfm2/modeling_lfm2.py index ef753e3b2893..c58329ead347 100644 --- a/src/transformers/models/lfm2/modeling_lfm2.py +++ b/src/transformers/models/lfm2/modeling_lfm2.py @@ -23,6 +23,7 @@ import torch import torch.nn.functional as F from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin @@ -180,6 +181,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py b/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py index 0369ae31b8ae..3a881abbee90 100644 --- a/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py +++ b/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py @@ -24,6 +24,7 @@ import torch import torch.nn.functional as F from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...cache_utils import Cache, DynamicCache @@ -256,6 +257,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/mimi/modeling_mimi.py b/src/transformers/models/mimi/modeling_mimi.py index 30480d1d1c03..1ef7ed2061f9 100644 --- a/src/transformers/models/mimi/modeling_mimi.py +++ b/src/transformers/models/mimi/modeling_mimi.py @@ -606,6 +606,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/minimax/modeling_minimax.py b/src/transformers/models/minimax/modeling_minimax.py index 69497f83cad8..9f35283377db 100644 --- a/src/transformers/models/minimax/modeling_minimax.py +++ b/src/transformers/models/minimax/modeling_minimax.py @@ -25,6 +25,7 @@ import torch import torch.nn.functional as F from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -351,6 +352,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/ministral/configuration_ministral.py b/src/transformers/models/ministral/configuration_ministral.py index 4ff445e6808a..cdd2074230b7 100644 --- a/src/transformers/models/ministral/configuration_ministral.py +++ b/src/transformers/models/ministral/configuration_ministral.py @@ -22,6 +22,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -49,13 +50,13 @@ class MinistralConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `MinistralModel` 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", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/ministral/modeling_ministral.py b/src/transformers/models/ministral/modeling_ministral.py index af4f7fbeae59..e45e9f3616e5 100644 --- a/src/transformers/models/ministral/modeling_ministral.py +++ b/src/transformers/models/ministral/modeling_ministral.py @@ -24,6 +24,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache @@ -91,6 +92,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/ministral3/modeling_ministral3.py b/src/transformers/models/ministral3/modeling_ministral3.py index 6aacf4c8ce3a..ce88bde32cd2 100644 --- a/src/transformers/models/ministral3/modeling_ministral3.py +++ b/src/transformers/models/ministral3/modeling_ministral3.py @@ -9,6 +9,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache @@ -60,6 +61,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/mistral4/modeling_mistral4.py b/src/transformers/models/mistral4/modeling_mistral4.py index 006ddad187bf..1ce9b9cd3e4e 100644 --- a/src/transformers/models/mistral4/modeling_mistral4.py +++ b/src/transformers/models/mistral4/modeling_mistral4.py @@ -23,6 +23,7 @@ import torch import torch.nn.functional as F from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -283,6 +284,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/mixtral/modeling_mixtral.py b/src/transformers/models/mixtral/modeling_mixtral.py index 991851dbadd3..a5f805083a46 100644 --- a/src/transformers/models/mixtral/modeling_mixtral.py +++ b/src/transformers/models/mixtral/modeling_mixtral.py @@ -29,6 +29,7 @@ import torch import torch.nn.functional as F from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -249,6 +250,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/mllama/modeling_mllama.py b/src/transformers/models/mllama/modeling_mllama.py index 3b9d12b9a225..8d9c7982e72a 100644 --- a/src/transformers/models/mllama/modeling_mllama.py +++ b/src/transformers/models/mllama/modeling_mllama.py @@ -495,6 +495,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/moshi/modeling_moshi.py b/src/transformers/models/moshi/modeling_moshi.py index a967445c18ec..8bb53bc0699d 100644 --- a/src/transformers/models/moshi/modeling_moshi.py +++ b/src/transformers/models/moshi/modeling_moshi.py @@ -360,6 +360,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/nanochat/modeling_nanochat.py b/src/transformers/models/nanochat/modeling_nanochat.py index 9205b89cd360..dda6e4fbb51e 100644 --- a/src/transformers/models/nanochat/modeling_nanochat.py +++ b/src/transformers/models/nanochat/modeling_nanochat.py @@ -24,6 +24,7 @@ import torch import torch.nn as nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -143,6 +144,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/nemotron_h/modeling_nemotron_h.py b/src/transformers/models/nemotron_h/modeling_nemotron_h.py index 9e264e5cfdcc..1cb01c08e1f8 100644 --- a/src/transformers/models/nemotron_h/modeling_nemotron_h.py +++ b/src/transformers/models/nemotron_h/modeling_nemotron_h.py @@ -26,6 +26,7 @@ import torch import torch.nn.functional as F from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -779,6 +780,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/nomic_bert/modeling_nomic_bert.py b/src/transformers/models/nomic_bert/modeling_nomic_bert.py index f2836a1ec0f6..26ffe02fcfb7 100644 --- a/src/transformers/models/nomic_bert/modeling_nomic_bert.py +++ b/src/transformers/models/nomic_bert/modeling_nomic_bert.py @@ -24,6 +24,7 @@ import torch import torch.nn as nn +from torch.distributed.tensor import DTensor, Replicate from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss from ... import initialization as init @@ -189,6 +190,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/olmoe/modeling_olmoe.py b/src/transformers/models/olmoe/modeling_olmoe.py index 5d89ec741529..bf20e0a58c1e 100644 --- a/src/transformers/models/olmoe/modeling_olmoe.py +++ b/src/transformers/models/olmoe/modeling_olmoe.py @@ -22,6 +22,7 @@ import torch import torch.nn.functional as F from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -175,6 +176,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/parakeet/modeling_parakeet.py b/src/transformers/models/parakeet/modeling_parakeet.py index 501a573f8494..ff759581dab6 100644 --- a/src/transformers/models/parakeet/modeling_parakeet.py +++ b/src/transformers/models/parakeet/modeling_parakeet.py @@ -24,6 +24,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -213,6 +214,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/persimmon/modeling_persimmon.py b/src/transformers/models/persimmon/modeling_persimmon.py index e0516ed7da9a..26d49f56fa34 100644 --- a/src/transformers/models/persimmon/modeling_persimmon.py +++ b/src/transformers/models/persimmon/modeling_persimmon.py @@ -151,6 +151,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/phi/modeling_phi.py b/src/transformers/models/phi/modeling_phi.py index e3f97a01ee4c..23d69f11775d 100644 --- a/src/transformers/models/phi/modeling_phi.py +++ b/src/transformers/models/phi/modeling_phi.py @@ -9,6 +9,7 @@ import torch import torch.nn as nn +from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache @@ -125,6 +126,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/phimoe/modeling_phimoe.py b/src/transformers/models/phimoe/modeling_phimoe.py index 23bc944c522a..fafdbff37c3e 100644 --- a/src/transformers/models/phimoe/modeling_phimoe.py +++ b/src/transformers/models/phimoe/modeling_phimoe.py @@ -24,6 +24,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -149,6 +150,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/qwen2/modeling_qwen2.py b/src/transformers/models/qwen2/modeling_qwen2.py index 9263e1d42937..021543db455f 100644 --- a/src/transformers/models/qwen2/modeling_qwen2.py +++ b/src/transformers/models/qwen2/modeling_qwen2.py @@ -9,6 +9,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache @@ -141,6 +142,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py b/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py index d4150d0a74d7..eb7585373a61 100644 --- a/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py +++ b/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py @@ -29,6 +29,7 @@ import torch import torch.nn.functional as F from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -187,6 +188,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/qwen3/modeling_qwen3.py b/src/transformers/models/qwen3/modeling_qwen3.py index a9e4cd4b75bb..7f483e22373d 100644 --- a/src/transformers/models/qwen3/modeling_qwen3.py +++ b/src/transformers/models/qwen3/modeling_qwen3.py @@ -446,6 +446,7 @@ 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"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py b/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py index ddf84fc575b7..2e1af14c437c 100644 --- a/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py +++ b/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py @@ -24,6 +24,7 @@ import torch import torch.nn.functional as F from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -81,6 +82,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed 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 7b6c8b5b1bd4..650c63de4a87 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 @@ -27,6 +27,7 @@ import numpy as np import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from torch.nn import Parameter from torch.nn import functional as F @@ -1465,6 +1466,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py b/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py index 73678ee8c736..b772ac6c49f0 100644 --- a/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py +++ b/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py @@ -26,6 +26,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -427,6 +428,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py b/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py index ce405683fc94..0c6c59f37c79 100644 --- a/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py +++ b/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py @@ -26,6 +26,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -209,6 +210,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py b/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py index 6e9c072b8860..a7c4702e1497 100644 --- a/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py +++ b/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py @@ -160,6 +160,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/seed_oss/modeling_seed_oss.py b/src/transformers/models/seed_oss/modeling_seed_oss.py index 1ebc8f10a272..861789644d6d 100644 --- a/src/transformers/models/seed_oss/modeling_seed_oss.py +++ b/src/transformers/models/seed_oss/modeling_seed_oss.py @@ -23,6 +23,7 @@ import torch import torch.nn as nn +from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache @@ -112,6 +113,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/smollm3/modeling_smollm3.py b/src/transformers/models/smollm3/modeling_smollm3.py index 8d911e414b0f..2cf346b08d8d 100644 --- a/src/transformers/models/smollm3/modeling_smollm3.py +++ b/src/transformers/models/smollm3/modeling_smollm3.py @@ -23,6 +23,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache @@ -139,6 +140,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/solar_open/modeling_solar_open.py b/src/transformers/models/solar_open/modeling_solar_open.py index dfa30292455f..68fcb6a64dde 100644 --- a/src/transformers/models/solar_open/modeling_solar_open.py +++ b/src/transformers/models/solar_open/modeling_solar_open.py @@ -23,6 +23,7 @@ import torch import torch.nn.functional as F from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -248,6 +249,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/stablelm/modeling_stablelm.py b/src/transformers/models/stablelm/modeling_stablelm.py index 9b9e0430e985..5e9c8eee71e7 100755 --- a/src/transformers/models/stablelm/modeling_stablelm.py +++ b/src/transformers/models/stablelm/modeling_stablelm.py @@ -150,6 +150,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/starcoder2/modeling_starcoder2.py b/src/transformers/models/starcoder2/modeling_starcoder2.py index 8b89a1d1745c..d944e3afa422 100644 --- a/src/transformers/models/starcoder2/modeling_starcoder2.py +++ b/src/transformers/models/starcoder2/modeling_starcoder2.py @@ -28,6 +28,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache @@ -95,6 +96,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/t5gemma/modeling_t5gemma.py b/src/transformers/models/t5gemma/modeling_t5gemma.py index a6b9b5392194..95eec7fb9f29 100644 --- a/src/transformers/models/t5gemma/modeling_t5gemma.py +++ b/src/transformers/models/t5gemma/modeling_t5gemma.py @@ -23,6 +23,7 @@ import torch import torch.nn as nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -189,6 +190,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/t5gemma2/modeling_t5gemma2.py b/src/transformers/models/t5gemma2/modeling_t5gemma2.py index 2e0dddc17876..6e587b1a0e30 100644 --- a/src/transformers/models/t5gemma2/modeling_t5gemma2.py +++ b/src/transformers/models/t5gemma2/modeling_t5gemma2.py @@ -24,6 +24,7 @@ import torch import torch.nn as nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -200,6 +201,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/timesfm2_5/modeling_timesfm2_5.py b/src/transformers/models/timesfm2_5/modeling_timesfm2_5.py index e7b4e799d20b..6ae3b994b47b 100644 --- a/src/transformers/models/timesfm2_5/modeling_timesfm2_5.py +++ b/src/transformers/models/timesfm2_5/modeling_timesfm2_5.py @@ -26,6 +26,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -229,6 +230,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/vaultgemma/modeling_vaultgemma.py b/src/transformers/models/vaultgemma/modeling_vaultgemma.py index f0a2e48d20b8..46e152403f71 100644 --- a/src/transformers/models/vaultgemma/modeling_vaultgemma.py +++ b/src/transformers/models/vaultgemma/modeling_vaultgemma.py @@ -24,6 +24,7 @@ import torch import torch.nn as nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -107,6 +108,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py b/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py index 07325b0ea559..624e2ca7971d 100644 --- a/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py +++ b/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py @@ -26,6 +26,7 @@ import torch import torch.nn as nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -279,6 +280,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/youtu/modeling_youtu.py b/src/transformers/models/youtu/modeling_youtu.py index d40bef358da6..6ed3bc109921 100644 --- a/src/transformers/models/youtu/modeling_youtu.py +++ b/src/transformers/models/youtu/modeling_youtu.py @@ -31,6 +31,7 @@ import torch import torch.nn.functional as F from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -180,6 +181,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/zamba2/modeling_zamba2.py b/src/transformers/models/zamba2/modeling_zamba2.py index 6e4ea7dcf2d8..b9eb255c399f 100644 --- a/src/transformers/models/zamba2/modeling_zamba2.py +++ b/src/transformers/models/zamba2/modeling_zamba2.py @@ -25,6 +25,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -219,6 +220,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed From 92a3491683cceab5ae3f8bad0ca0bf6739a289d3 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 14 Apr 2026 15:41:33 +0000 Subject: [PATCH 013/116] Fix DTensor imports in Copied-from model files --- src/transformers/models/chameleon/modeling_chameleon.py | 1 + src/transformers/models/falcon/modeling_falcon.py | 1 + .../models/gpt_neox_japanese/modeling_gpt_neox_japanese.py | 1 + src/transformers/models/mimi/modeling_mimi.py | 1 + src/transformers/models/mllama/modeling_mllama.py | 1 + src/transformers/models/moshi/modeling_moshi.py | 1 + src/transformers/models/persimmon/modeling_persimmon.py | 1 + .../models/recurrent_gemma/modeling_recurrent_gemma.py | 1 + src/transformers/models/stablelm/modeling_stablelm.py | 1 + 9 files changed, 9 insertions(+) diff --git a/src/transformers/models/chameleon/modeling_chameleon.py b/src/transformers/models/chameleon/modeling_chameleon.py index cdc56f7e0458..69e6cc834336 100644 --- a/src/transformers/models/chameleon/modeling_chameleon.py +++ b/src/transformers/models/chameleon/modeling_chameleon.py @@ -21,6 +21,7 @@ import torch import torch.nn.functional as F from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache diff --git a/src/transformers/models/falcon/modeling_falcon.py b/src/transformers/models/falcon/modeling_falcon.py index 26336cc7e674..f9582748823b 100644 --- a/src/transformers/models/falcon/modeling_falcon.py +++ b/src/transformers/models/falcon/modeling_falcon.py @@ -19,6 +19,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, LayerNorm, MSELoss from torch.nn import functional as F diff --git a/src/transformers/models/gpt_neox_japanese/modeling_gpt_neox_japanese.py b/src/transformers/models/gpt_neox_japanese/modeling_gpt_neox_japanese.py index be0008400097..a2f423a94578 100755 --- a/src/transformers/models/gpt_neox_japanese/modeling_gpt_neox_japanese.py +++ b/src/transformers/models/gpt_neox_japanese/modeling_gpt_neox_japanese.py @@ -19,6 +19,7 @@ import torch from torch import Tensor, nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN diff --git a/src/transformers/models/mimi/modeling_mimi.py b/src/transformers/models/mimi/modeling_mimi.py index 1ef7ed2061f9..92e5e22b8754 100644 --- a/src/transformers/models/mimi/modeling_mimi.py +++ b/src/transformers/models/mimi/modeling_mimi.py @@ -20,6 +20,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN diff --git a/src/transformers/models/mllama/modeling_mllama.py b/src/transformers/models/mllama/modeling_mllama.py index 8d9c7982e72a..d2373afd68d8 100644 --- a/src/transformers/models/mllama/modeling_mllama.py +++ b/src/transformers/models/mllama/modeling_mllama.py @@ -20,6 +20,7 @@ import torch import torch.nn.functional as F from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN diff --git a/src/transformers/models/moshi/modeling_moshi.py b/src/transformers/models/moshi/modeling_moshi.py index 8bb53bc0699d..51aeaaa50ffd 100644 --- a/src/transformers/models/moshi/modeling_moshi.py +++ b/src/transformers/models/moshi/modeling_moshi.py @@ -20,6 +20,7 @@ import torch import torch.nn as nn +from torch.distributed.tensor import DTensor, Replicate from torch.nn import CrossEntropyLoss from ... import initialization as init diff --git a/src/transformers/models/persimmon/modeling_persimmon.py b/src/transformers/models/persimmon/modeling_persimmon.py index 26d49f56fa34..ae75f666cd18 100644 --- a/src/transformers/models/persimmon/modeling_persimmon.py +++ b/src/transformers/models/persimmon/modeling_persimmon.py @@ -23,6 +23,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache diff --git a/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py b/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py index a7c4702e1497..3f45e4ce4cf3 100644 --- a/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py +++ b/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py @@ -20,6 +20,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN diff --git a/src/transformers/models/stablelm/modeling_stablelm.py b/src/transformers/models/stablelm/modeling_stablelm.py index 5e9c8eee71e7..45f0b03e574c 100755 --- a/src/transformers/models/stablelm/modeling_stablelm.py +++ b/src/transformers/models/stablelm/modeling_stablelm.py @@ -23,6 +23,7 @@ import torch from torch import nn +from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache From 7ca7911bc57222f8f52a5df3362db31783402af3 Mon Sep 17 00:00:00 2001 From: Ferdinand Mom <47445085+3outeille@users.noreply.github.com> Date: Tue, 14 Apr 2026 18:11:59 +0200 Subject: [PATCH 014/116] MoE expert parallelism + sequence parallelism (#45408) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * MoE expert parallelism + sequence parallelism - Add PackedColwiseParallel for fused gate_up_proj weights - Add MoEExpertsParallel with per-expert DTensor sharding - Add PrepareModuleInputOutput for SP allgather/split hooks - Add _AllReduceBackward for MoE routing weight gradients - Extend TPStyle with moe_experts, packed_colwise, activation, module kinds - _StridedShard handling in core_model_loading for interleaved weights - MoE model configs: mixtral, deepseek_v3, qwen3 with SP plans - DTensor rotary_pos_emb guard for mixtral * Fix ruff linting and formatting * Fix ruff formatting in core_model_loading.py * Restore _IdentityOp accidentally removed in 25a1f4808e The _IdentityOp class (added by PR #44983) was accidentally deleted during the MoE expert parallelism work. It is needed by finegrained_fp8.py and metal_quantization.py as a pass-through reverse_op for dequantize operations. Co-Authored-By: Claude Opus 4.6 (1M context) * Backport new TP/FSDP API + fix DTensor imports in Copied-from models * from_pretrained orchestration + distributed save/load (#45409) * from_pretrained orchestration + save/load - Add gather_full_state_dict() for DTensor→full tensor saving - Add convert_strided_to_shard() / restore_strided_from_shard() for DCP - Add _redistribute_dtensor() helper - Full distributed_config integration in from_pretrained/save_pretrained - Rename apply_fsdp2 → apply_fully_shard_data_parallel - save_optimizer() / load_optimizer() in distributed/utils - Trainer integration with distributed_config - Updated FSDP and TP tests for new orchestration API - DTensor shard-on-read test updates * revert distributed utils * eaaea * all tests for core modeling are passing * populate import from init for tp * ruff * ruff --------- Co-authored-by: Claude Opus 4.6 (1M context) --- src/transformers/core_model_loading.py | 11 + src/transformers/integrations/__init__.py | 23 +- src/transformers/modeling_utils.py | 13 +- .../deepseek_v3/configuration_deepseek_v3.py | 39 ++- .../models/dots1/modeling_dots1.py | 62 ++-- .../models/mixtral/configuration_mixtral.py | 44 ++- .../models/nanochat/modeling_nanochat.py | 60 ++-- .../models/qwen3/configuration_qwen3.py | 46 ++- .../models/qwen3/modeling_qwen3.py | 57 ++-- .../models/qwen3/modular_qwen3.py | 22 +- .../models/qwen3_5/modeling_qwen3_5.py | 74 ++--- .../qwen3_5_moe/modeling_qwen3_5_moe.py | 74 ++--- .../models/qwen3_moe/modeling_qwen3_moe.py | 62 ++-- .../models/qwen3_next/modeling_qwen3_next.py | 74 ++--- .../configuration_qwen3_omni_moe.py | 46 ++- .../qwen3_omni_moe/modeling_qwen3_omni_moe.py | 26 +- .../models/qwen3_vl/modeling_qwen3_vl.py | 21 +- .../qwen3_vl_moe/modeling_qwen3_vl_moe.py | 26 +- .../models/youtu/configuration_youtu.py | 15 + src/transformers/trainer.py | 7 +- tests/test_fsdp_mixin.py | 24 +- tests/utils/test_core_model_loading.py | 292 +++++++++++++++++- 22 files changed, 697 insertions(+), 421 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index cf06f62fd685..33a07251c0f1 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -118,6 +118,17 @@ def reverse_op(self) -> ConversionOps: raise NotImplementedError +class _IdentityOp(ConversionOps): + """Pass-through reverse op for dequantize operations. + + Dequantized weights are already in their target dtype and should be + saved as-is without any conversion. + """ + + def convert(self, input_dict: dict[str, Any], **kwargs) -> dict[str, Any]: + return input_dict + + class Chunk(ConversionOps): """Split a tensor along ``dim`` into equally sized chunks.""" diff --git a/src/transformers/integrations/__init__.py b/src/transformers/integrations/__init__.py index 336db3773f76..d274b31837e2 100755 --- a/src/transformers/integrations/__init__.py +++ b/src/transformers/integrations/__init__.py @@ -161,9 +161,12 @@ ] _import_structure["tensor_parallel"] = [ - "shard_and_distribute_module", - "ALL_PARALLEL_STYLES", - "translate_to_torch_parallel_style", + "TPStyle", + "apply_tensor_parallel", + "convert_strided_to_shard", + "gather_full_state_dict", + "restore_strided_from_shard", + "verify_tp_plan", ] try: if not is_torch_greater_or_equal("2.5"): @@ -295,6 +298,14 @@ from .quanto import replace_with_quanto_layers from .sinq import SinqDeserialize, SinqQuantize from .spqr import replace_with_spqr_linear + from .tensor_parallel import ( + TPStyle, + apply_tensor_parallel, + convert_strided_to_shard, + gather_full_state_dict, + restore_strided_from_shard, + verify_tp_plan, + ) from .vptq import replace_with_vptq_linear try: @@ -305,12 +316,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/modeling_utils.py b/src/transformers/modeling_utils.py index 082b294fb41f..1629c4ca4d9b 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -38,6 +38,7 @@ from safetensors import safe_open 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 @@ -4123,15 +4124,13 @@ def from_pretrained( if distributed_config is not None: model.config.distributed_config = distributed_config model.device_mesh = device_mesh - - def sub_mesh(name): - return device_mesh[name] if device_mesh.ndim > 1 else device_mesh - mesh_dim_names = device_mesh.mesh_dim_names or () if "tp" in mesh_dim_names: - model = apply_tensor_parallel(model, sub_mesh("tp"), distributed_config.tp_plan) + tp_mesh = device_mesh["tp"] if device_mesh.ndim > 1 else device_mesh + model = apply_tensor_parallel(model, tp_mesh, distributed_config.tp_plan) if "fsdp" in mesh_dim_names: - model = apply_fully_shard_data_parallel(model, sub_mesh("fsdp"), distributed_config.fsdp_plan) + fsdp_mesh = device_mesh["fsdp"] if device_mesh.ndim > 1 else device_mesh + model = apply_fully_shard_data_parallel(model, fsdp_mesh, distributed_config.fsdp_plan) else: # Accelerate path: auto device mapping if device_map is not None: @@ -4552,8 +4551,6 @@ 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) - from torch.distributed.tensor import DTensor - if isinstance(param, DTensor): # DTensor from parallelize_module on meta — materialize on actual device local_value = torch.empty( diff --git a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py index 4178547a5ff2..a9216f4db2ab 100644 --- a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py @@ -18,6 +18,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -49,15 +50,20 @@ 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.shared_experts.gate_proj": "colwise", - "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={ + "gate_up_proj": "packed_colwise", + "down_proj": "rowwise", + }, + ), + "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -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/dots1/modeling_dots1.py b/src/transformers/models/dots1/modeling_dots1.py index 89561bfa4ed8..b8a364b760c4 100644 --- a/src/transformers/models/dots1/modeling_dots1.py +++ b/src/transformers/models/dots1/modeling_dots1.py @@ -29,12 +29,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import ( - use_experts_implementation, - use_kernel_forward_from_hub, - use_kernel_func_from_hub, - use_kernelized_func, -) +from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -134,43 +129,6 @@ def forward(self, x, position_ids): return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) -def rotate_half(x): - """Rotates half the hidden dims of the input.""" - x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] - return torch.cat((-x2, x1), dim=-1) - - -@use_kernel_func_from_hub("rotary_pos_emb") -def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): - """Applies Rotary Position Embedding to the query and key tensors. - - Args: - q (`torch.Tensor`): The query tensor. - k (`torch.Tensor`): The key tensor. - cos (`torch.Tensor`): The cosine part of the rotary embedding. - sin (`torch.Tensor`): The sine part of the rotary embedding. - unsqueeze_dim (`int`, *optional*, defaults to 1): - The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and - sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note - that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and - k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes - cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have - the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. - Returns: - `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. - """ - cos = cos.unsqueeze(unsqueeze_dim) - sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) - q_embed = (q * cos) + (rotate_half(q) * sin) - k_embed = (k * cos) + (rotate_half(k) * sin) - return q_embed, k_embed - - def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, @@ -208,6 +166,24 @@ def eager_attention_forward( return attn_output, attn_weights +def rotate_half(x): + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + @use_kernelized_func(apply_rotary_pos_emb) class Dots1Attention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" diff --git a/src/transformers/models/mixtral/configuration_mixtral.py b/src/transformers/models/mixtral/configuration_mixtral.py index 240f24411031..7b6ab7aaa974 100644 --- a/src/transformers/models/mixtral/configuration_mixtral.py +++ b/src/transformers/models/mixtral/configuration_mixtral.py @@ -16,6 +16,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -42,14 +43,43 @@ 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.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={ + "gate_up_proj": "packed_colwise", + "down_proj": "rowwise", + }, + ), + } + + # TP + Sequence Parallelism plan (for training). + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather_split"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={ + "gate_up_proj": "packed_colwise", + "down_proj": "rowwise", + }, + ), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/nanochat/modeling_nanochat.py b/src/transformers/models/nanochat/modeling_nanochat.py index dda6e4fbb51e..7d351cf56306 100644 --- a/src/transformers/models/nanochat/modeling_nanochat.py +++ b/src/transformers/models/nanochat/modeling_nanochat.py @@ -123,36 +123,6 @@ def forward(self, x, position_ids): return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) -@use_kernel_func_from_hub("rotary_pos_emb") -def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): - """Applies Rotary Position Embedding to the query and key tensors. - - Args: - q (`torch.Tensor`): The query tensor. - k (`torch.Tensor`): The key tensor. - cos (`torch.Tensor`): The cosine part of the rotary embedding. - sin (`torch.Tensor`): The sine part of the rotary embedding. - unsqueeze_dim (`int`, *optional*, defaults to 1): - The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and - sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note - that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and - k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes - cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have - the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. - Returns: - `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. - """ - cos = cos.unsqueeze(unsqueeze_dim) - sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) - q_embed = (q * cos) + (rotate_half(q) * sin) - k_embed = (k * cos) + (rotate_half(k) * sin) - return q_embed, k_embed - - def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, @@ -190,6 +160,36 @@ def eager_attention_forward( return attn_output, attn_weights +@use_kernel_func_from_hub("rotary_pos_emb") +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + def rotate_half(x): """Rotates half the hidden dims of the input with flipped signs for NanoChat.""" x1 = x[..., : x.shape[-1] // 2] diff --git a/src/transformers/models/qwen3/configuration_qwen3.py b/src/transformers/models/qwen3/configuration_qwen3.py index 07ad0bb24b33..372922204d3e 100644 --- a/src/transformers/models/qwen3/configuration_qwen3.py +++ b/src/transformers/models/qwen3/configuration_qwen3.py @@ -16,6 +16,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -41,17 +42,42 @@ 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.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("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": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/qwen3/modeling_qwen3.py b/src/transformers/models/qwen3/modeling_qwen3.py index 7f483e22373d..550debc37eec 100644 --- a/src/transformers/models/qwen3/modeling_qwen3.py +++ b/src/transformers/models/qwen3/modeling_qwen3.py @@ -28,7 +28,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func +from ...integrations import use_kernel_forward_from_hub, use_kernelized_func from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -149,43 +149,6 @@ def forward(self, x, position_ids): return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) -def rotate_half(x): - """Rotates half the hidden dims of the input.""" - x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] - return torch.cat((-x2, x1), dim=-1) - - -@use_kernel_func_from_hub("rotary_pos_emb") -def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): - """Applies Rotary Position Embedding to the query and key tensors. - - Args: - q (`torch.Tensor`): The query tensor. - k (`torch.Tensor`): The key tensor. - cos (`torch.Tensor`): The cosine part of the rotary embedding. - sin (`torch.Tensor`): The sine part of the rotary embedding. - unsqueeze_dim (`int`, *optional*, defaults to 1): - The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and - sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note - that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and - k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes - cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have - the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. - Returns: - `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. - """ - cos = cos.unsqueeze(unsqueeze_dim) - sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) - q_embed = (q * cos) + (rotate_half(q) * sin) - k_embed = (k * cos) + (rotate_half(k) * sin) - return q_embed, k_embed - - def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, @@ -223,6 +186,24 @@ def eager_attention_forward( return attn_output, attn_weights +def rotate_half(x): + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + @use_kernelized_func(apply_rotary_pos_emb) class Qwen3Attention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" diff --git a/src/transformers/models/qwen3/modular_qwen3.py b/src/transformers/models/qwen3/modular_qwen3.py index 73cde6d89a7a..74a8447d1bb7 100644 --- a/src/transformers/models/qwen3/modular_qwen3.py +++ b/src/transformers/models/qwen3/modular_qwen3.py @@ -16,6 +16,7 @@ from collections.abc import Callable import torch +from torch.distributed.tensor import DTensor, Replicate from ...cache_utils import Cache from ...modeling_flash_attention_utils import FlashAttentionKwargs @@ -34,7 +35,6 @@ Qwen2ForTokenClassification, Qwen2RMSNorm, Qwen2RotaryEmbedding, - apply_rotary_pos_emb, eager_attention_forward, ) from .configuration_qwen3 import Qwen3Config @@ -45,6 +45,24 @@ _CHECKPOINT_FOR_DOC = "Qwen/Qwen3-8B" +def rotate_half(x): + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + class Qwen3RMSNorm(Qwen2RMSNorm): pass @@ -108,6 +126,8 @@ def forward( class Qwen3ForCausalLM(Qwen2ForCausalLM): + _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + def forward( self, **super_kwargs: Unpack[TransformersKwargs], diff --git a/src/transformers/models/qwen3_5/modeling_qwen3_5.py b/src/transformers/models/qwen3_5/modeling_qwen3_5.py index 4dd3dfbaaf60..81fc2a40ea60 100644 --- a/src/transformers/models/qwen3_5/modeling_qwen3_5.py +++ b/src/transformers/models/qwen3_5/modeling_qwen3_5.py @@ -533,6 +533,43 @@ def forward( return output +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] @@ -579,43 +616,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): return q_embed, k_embed -def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: - """ - This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, - num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) - """ - batch, num_key_value_heads, slen, head_dim = hidden_states.shape - if n_rep == 1: - return hidden_states - hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) - return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) - - -def eager_attention_forward( - module: nn.Module, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attention_mask: torch.Tensor | None, - scaling: float, - dropout: float = 0.0, - **kwargs: Unpack[TransformersKwargs], -): - key_states = repeat_kv(key, module.num_key_value_groups) - value_states = repeat_kv(value, module.num_key_value_groups) - - attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling - if attention_mask is not None: - attn_weights = attn_weights + attention_mask - - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) - attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) - attn_output = torch.matmul(attn_weights, value_states) - attn_output = attn_output.transpose(1, 2).contiguous() - - return attn_output, attn_weights - - @use_kernelized_func(apply_rotary_pos_emb) class Qwen3_5Attention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" 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 125ded124cf7..0e80ba9b5d5d 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 @@ -534,6 +534,43 @@ def forward( return output +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] @@ -580,43 +617,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): return q_embed, k_embed -def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: - """ - This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, - num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) - """ - batch, num_key_value_heads, slen, head_dim = hidden_states.shape - if n_rep == 1: - return hidden_states - hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) - return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) - - -def eager_attention_forward( - module: nn.Module, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attention_mask: torch.Tensor | None, - scaling: float, - dropout: float = 0.0, - **kwargs: Unpack[TransformersKwargs], -): - key_states = repeat_kv(key, module.num_key_value_groups) - value_states = repeat_kv(value, module.num_key_value_groups) - - attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling - if attention_mask is not None: - attn_weights = attn_weights + attention_mask - - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) - attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) - attn_output = torch.matmul(attn_weights, value_states) - attn_output = attn_output.transpose(1, 2).contiguous() - - return attn_output, attn_weights - - @use_kernelized_func(apply_rotary_pos_emb) class Qwen3_5MoeAttention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" diff --git a/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py b/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py index 2e1af14c437c..cb827352bd57 100644 --- a/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py +++ b/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py @@ -30,12 +30,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import ( - use_experts_implementation, - use_kernel_forward_from_hub, - use_kernel_func_from_hub, - use_kernelized_func, -) +from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -54,43 +49,6 @@ from .configuration_qwen3_moe import Qwen3MoeConfig -def rotate_half(x): - """Rotates half the hidden dims of the input.""" - x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] - return torch.cat((-x2, x1), dim=-1) - - -@use_kernel_func_from_hub("rotary_pos_emb") -def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): - """Applies Rotary Position Embedding to the query and key tensors. - - Args: - q (`torch.Tensor`): The query tensor. - k (`torch.Tensor`): The key tensor. - cos (`torch.Tensor`): The cosine part of the rotary embedding. - sin (`torch.Tensor`): The sine part of the rotary embedding. - unsqueeze_dim (`int`, *optional*, defaults to 1): - The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and - sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note - that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and - k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes - cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have - the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. - Returns: - `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. - """ - cos = cos.unsqueeze(unsqueeze_dim) - sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) - q_embed = (q * cos) + (rotate_half(q) * sin) - k_embed = (k * cos) + (rotate_half(k) * sin) - return q_embed, k_embed - - def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, @@ -128,6 +86,24 @@ def eager_attention_forward( return attn_output, attn_weights +def rotate_half(x): + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + if isinstance(q, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + @use_kernelized_func(apply_rotary_pos_emb) class Qwen3MoeAttention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" diff --git a/src/transformers/models/qwen3_next/modeling_qwen3_next.py b/src/transformers/models/qwen3_next/modeling_qwen3_next.py index cd152e3d3e59..af9239618f95 100644 --- a/src/transformers/models/qwen3_next/modeling_qwen3_next.py +++ b/src/transformers/models/qwen3_next/modeling_qwen3_next.py @@ -169,6 +169,43 @@ def extra_repr(self): return f"{tuple(self.weight.shape)}, eps={self.eps}" +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] @@ -215,43 +252,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): return q_embed, k_embed -def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: - """ - This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, - num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) - """ - batch, num_key_value_heads, slen, head_dim = hidden_states.shape - if n_rep == 1: - return hidden_states - hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) - return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) - - -def eager_attention_forward( - module: nn.Module, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attention_mask: torch.Tensor | None, - scaling: float, - dropout: float = 0.0, - **kwargs: Unpack[TransformersKwargs], -): - key_states = repeat_kv(key, module.num_key_value_groups) - value_states = repeat_kv(value, module.num_key_value_groups) - - attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling - if attention_mask is not None: - attn_weights = attn_weights + attention_mask - - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) - attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) - attn_output = torch.matmul(attn_weights, value_states) - attn_output = attn_output.transpose(1, 2).contiguous() - - return attn_output, attn_weights - - @use_kernelized_func(apply_rotary_pos_emb) class Qwen3NextAttention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" 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 70da6ca9a360..04534187d73b 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 @@ -21,6 +21,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring, logging @@ -258,17 +259,42 @@ 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.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("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": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), 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 650c63de4a87..7f1c1655771a 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 @@ -35,12 +35,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import ( - use_experts_implementation, - use_kernel_forward_from_hub, - use_kernel_func_from_hub, - use_kernelized_func, -) +from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -815,7 +810,6 @@ def _get_feat_extract_output_lengths(self, input_lengths: torch.LongTensor): def rotate_half(x): - """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) @@ -1445,25 +1439,7 @@ def extra_repr(self): return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" -@use_kernel_func_from_hub("rotary_pos_emb") def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): - """Applies Rotary Position Embedding to the query and key tensors. - - Args: - q (`torch.Tensor`): The query tensor. - k (`torch.Tensor`): The key tensor. - cos (`torch.Tensor`): The cosine part of the rotary embedding. - sin (`torch.Tensor`): The sine part of the rotary embedding. - unsqueeze_dim (`int`, *optional*, defaults to 1): - The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and - sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note - that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and - k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes - cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have - the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. - Returns: - `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. - """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) if isinstance(q, DTensor): diff --git a/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py b/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py index b772ac6c49f0..62e97f14b514 100644 --- a/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py +++ b/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py @@ -32,7 +32,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func +from ...integrations import use_kernel_forward_from_hub, use_kernelized_func from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -123,7 +123,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: def rotate_half(x): - """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) @@ -407,25 +406,7 @@ def extra_repr(self): return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" -@use_kernel_func_from_hub("rotary_pos_emb") def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): - """Applies Rotary Position Embedding to the query and key tensors. - - Args: - q (`torch.Tensor`): The query tensor. - k (`torch.Tensor`): The key tensor. - cos (`torch.Tensor`): The cosine part of the rotary embedding. - sin (`torch.Tensor`): The sine part of the rotary embedding. - unsqueeze_dim (`int`, *optional*, defaults to 1): - The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and - sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note - that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and - k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes - cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have - the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. - Returns: - `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. - """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) if isinstance(q, DTensor): diff --git a/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py b/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py index 0c6c59f37c79..4abf05fd8316 100644 --- a/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py +++ b/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py @@ -32,12 +32,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import ( - use_experts_implementation, - use_kernel_forward_from_hub, - use_kernel_func_from_hub, - use_kernelized_func, -) +from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -146,7 +141,6 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: def rotate_half(x): - """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) @@ -189,25 +183,7 @@ def eager_attention_forward( return attn_output, attn_weights -@use_kernel_func_from_hub("rotary_pos_emb") def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): - """Applies Rotary Position Embedding to the query and key tensors. - - Args: - q (`torch.Tensor`): The query tensor. - k (`torch.Tensor`): The key tensor. - cos (`torch.Tensor`): The cosine part of the rotary embedding. - sin (`torch.Tensor`): The sine part of the rotary embedding. - unsqueeze_dim (`int`, *optional*, defaults to 1): - The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and - sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note - that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and - k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes - cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have - the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. - Returns: - `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. - """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) if isinstance(q, DTensor): diff --git a/src/transformers/models/youtu/configuration_youtu.py b/src/transformers/models/youtu/configuration_youtu.py index 6d9f2cef1f96..61019d70391b 100644 --- a/src/transformers/models/youtu/configuration_youtu.py +++ b/src/transformers/models/youtu/configuration_youtu.py @@ -103,5 +103,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/trainer.py b/src/transformers/trainer.py index 1b8dacb632cc..57042d1aad92 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -2392,9 +2392,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/tests/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py index d7a0a4ca3340..f6f7f7e6e2ab 100644 --- a/tests/test_fsdp_mixin.py +++ b/tests/test_fsdp_mixin.py @@ -49,9 +49,10 @@ from torch.distributed.tensor import DTensor from torch.nn.parallel import DistributedDataParallel as DDP + from transformers.distributed import DistributedConfig from transformers.integrations.fsdp import ( _find_final_norm, - apply_fsdp2, + apply_fully_shard_data_parallel, get_transformer_block_classes, initialize_fsdp, ) @@ -373,12 +374,11 @@ def train_fsdp2( ): # -- Phase 1: Pre-checkpoint run -- train only the first `checkpoint_step` steps, then save _set_determinism(SEED) - _, device_mesh, _ = initialize_fsdp(fsdp_plan=fsdp_plan) + distributed_config = DistributedConfig(fsdp_plan=fsdp_plan) pre_ckpt_model = AutoModelForCausalLM.from_pretrained( init_model_dir, torch_dtype=dtype, - fsdp_plan=fsdp_plan, - device_mesh=device_mesh, + distributed_config=distributed_config, attn_implementation="eager", ) pre_ckpt_model.train() @@ -415,8 +415,7 @@ def train_fsdp2( resumed_model = AutoModelForCausalLM.from_pretrained( model_dir, torch_dtype=dtype, - fsdp_plan=fsdp_plan, - device_mesh=device_mesh, + distributed_config=distributed_config, attn_implementation="eager", ) resumed_model.train() @@ -461,16 +460,14 @@ def _test_fsdp2_save_load_impl(rank, config_class, config_dict): batches = _build_repeated_training_batches(config, device, 3) - auto_plan = {"mode": "auto"} + distributed_config = DistributedConfig(fsdp_plan="auto") init_tmpdir, init_tmpdir_obj = _save_init_pretrained(rank, config, torch.float32) try: - _, device_mesh, _ = initialize_fsdp(fsdp_plan=auto_plan) _set_determinism(SEED) model = AutoModelForCausalLM.from_pretrained( init_tmpdir, - fsdp_plan=auto_plan, - device_mesh=device_mesh, + distributed_config=distributed_config, attn_implementation="eager", ) dist.barrier() @@ -495,8 +492,7 @@ def _test_fsdp2_save_load_impl(rank, config_class, config_dict): new_model = AutoModelForCausalLM.from_pretrained( tmpdir, - fsdp_plan=auto_plan, - device_mesh=device_mesh, + distributed_config=distributed_config, attn_implementation="eager", ) dist.barrier() @@ -522,7 +518,7 @@ def _test_fsdp2_save_load_impl(rank, config_class, config_dict): def _test_fsdp2_sharding_structure_impl(rank, config_class, config_dict, tie_word_embeddings): """ - Verify that apply_fsdp2(fsdp_plan={"mode": "auto"}) wraps exactly the right modules. + Verify that apply_fully_shard_data_parallel(fsdp_plan={"mode": "auto"}) wraps exactly the right modules. Expected FSDP targets: UNTIED TIED @@ -570,7 +566,7 @@ def _test_fsdp2_sharding_structure_impl(rank, config_class, config_dict, tie_wor if not weights_tied: expected_targets |= {output_name} - model = apply_fsdp2(model, device_mesh, fsdp_plan=auto_plan) + model = apply_fully_shard_data_parallel(model, device_mesh, fsdp_plan=auto_plan) actual_targets = {name for name, module in model.named_modules() if type(module).__name__.startswith("FSDP")} diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py index 942dcdc99b11..787cf7b903ad 100644 --- a/tests/utils/test_core_model_loading.py +++ b/tests/utils/test_core_model_loading.py @@ -16,14 +16,15 @@ import torch import torch.nn as nn +from torch.distributed.tensor.placement_types import Replicate, Shard, _StridedShard from transformers import PretrainedConfig from transformers.conversion_mapping import get_checkpoint_conversion_mapping, register_checkpoint_conversion_mapping from transformers.core_model_loading import ( Chunk, Concatenate, + DtensorShardOperation, ErnieFuseAndSplitTextVisionExperts, - FSDPShardOperation, MergeModulelist, PermuteForRope, WeightConverter, @@ -32,7 +33,7 @@ convert_and_load_state_dict_in_model, rename_source_key, revert_weight_conversion, - spawn_parallel_materialize, + spawn_materialize, ) from transformers.modeling_utils import LoadStateDictConfig from transformers.utils.import_utils import is_triton_available @@ -217,24 +218,113 @@ def __init__(self, add_extra_moe=False): class FakeMesh: - def __init__(self, world_size: int, rank: int): - self.shape = (world_size,) - self._rank = rank + """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._rank + return self._coord[0] def get_coordinate(self): - return (self._rank,) + 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.""" + op = object.__new__(DtensorShardOperation) + op.device_mesh = mesh + op.placements = tuple(placements) + ns = SimpleNamespace(shape=torch.Size(param_shape), ndim=len(param_shape)) + ns.dim = lambda: len(param_shape) + op.param = ns + op.local_shape = tuple(local_shape) + return op class TestConvertAndLoadStateDict(unittest.TestCase): - def test_fsdp_shard_aware_mixtral_conversion_uses_only_local_experts(self): - shard_op = FSDPShardOperation( - device_mesh=FakeMesh(world_size=2, rank=0), - rank=0, - empty_param=torch.empty((2, 4, 2)), - placements=(torch.distributed.tensor.placement_types.Shard(0),), + 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"], @@ -252,7 +342,7 @@ def test_fsdp_shard_aware_mixtral_conversion_uses_only_local_experts(self): "model.layers.0.experts.gate_up_proj.weight", f"model.layers.0.experts.{idx}.w1.weight", "experts.*.w1.weight", - spawn_parallel_materialize(None, tensor, shard_op, idx, device="cpu", dtype=None), + spawn_materialize(None, tensor, device="cpu", dtype=None, sharding_op=shard_op, tensor_idx=idx), ) for idx, tensor in enumerate( @@ -265,7 +355,7 @@ def test_fsdp_shard_aware_mixtral_conversion_uses_only_local_experts(self): "model.layers.0.experts.gate_up_proj.weight", f"model.layers.0.experts.{idx}.w3.weight", "experts.*.w3.weight", - spawn_parallel_materialize(None, tensor, shard_op, idx, device="cpu", dtype=None), + 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") @@ -785,6 +875,178 @@ def test_ernie4_5_vl_moe_conversion_reversed(self): self.assertTrue(compare_state_dicts(reversed_state_dict, state_dict)) +class TestDtensorShardOperation(unittest.TestCase): + """Unit tests for DtensorShardOperation.shard_tensor — one test per code path. + + Branch coverage map: + + shard_tensor() + ├── A: no sharding placements → full copy [test_no_shard_returns_full_tensor] + ├── B: expert path (tensor_idx set, ndim mismatch) + │ ├── B1: has_expert_sharding=False → fall through to C [test_expert_shaped_tp_only_no_expert_sharding] + │ ├── B2: not owns_local_expert → None [test_expert_filtering] + │ ├── B3: owned, no inner placements → full copy [test_expert_filtering] + │ └── B4: owned, with inner placements → _shard_nd [test_expert_filtering_preserves_inner_sharding] + └── C: _shard_nd() + ├── C1: _can_shard_on_read=False → _materialize_and_split [test_nd_strided_plus_shard_same_dim_fallback] + ├── C2: has_strided=False → contiguous slice + │ ├── 1D mesh [test_1d_shard_fast_path] + │ ├── 2D mesh [test_nd_contiguous_single_slice] + │ ├── negative dim [test_negative_dim_normalizes_correctly] + │ └── uneven division [test_contiguous_shard_uneven_division] + └── C3: has_strided=True → _compute_dim_ranges + _slice_and_read + ├── _StridedShard → _strided_ranges [test_nd_strided_shard_disjoint_ranges] + └── _source_tensor_needs_packing → contiguous [test_prepacked_strided_shard_uses_contiguous_source_slice] + + _slice_and_read (tested directly) + ├── all single ranges → simple slice [test_slice_and_read_all_single_ranges] + └── two multi-range dims → ValueError [test_slice_and_read_raises_on_two_multi_range_dims] + """ + + def test_no_shard_returns_full_tensor(self): + """Replicate-only → full copy.""" + mesh = FakeMesh(shape=(2,), rank=0) + op = _make_dtensor_shard_op(mesh, [Replicate()], param_shape=(4, 4), local_shape=(4, 4)) + tensor = torch.arange(16).reshape(4, 4).float() + torch.testing.assert_close(op.shard_tensor(tensor), tensor) + + def test_1d_shard_fast_path(self): + # TODO(3outeille): double check fast path + tensor = torch.arange(16).reshape(4, 4).float() + for rank, expected in [(0, tensor[:2]), (1, tensor[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, msg=f"rank {rank}") + + def test_nd_contiguous_single_slice(self): + """nD Shard on different dims → single slice read per rank.""" + tensor = torch.arange(64).reshape(8, 8).float() + expected = {0: tensor[:4, :4], 1: tensor[:4, 4:], 2: tensor[4:, :4], 3: tensor[4:, 4:]} + 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_nd_strided_shard_disjoint_ranges(self): + """_StridedShard on its own dim → multiple slice reads + cat.""" + tensor = torch.arange(64).reshape(8, 8).float() + # Shard(0) splits rows; _StridedShard(1, split_factor=2) produces disjoint col ranges + expected = { + 0: torch.cat([tensor[:4, :2], tensor[:4, 4:6]], dim=1), + 1: torch.cat([tensor[:4, 2:4], tensor[:4, 6:8]], dim=1), + 2: torch.cat([tensor[4:, :2], tensor[4:, 4:6]], dim=1), + 3: torch.cat([tensor[4:, 2:4], tensor[4:, 6:8]], dim=1), + } + 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=(8, 8), + local_shape=(4, 4), + ) + torch.testing.assert_close(op.shard_tensor(tensor), expected[rank], msg=f"rank {rank}") + + def test_nd_strided_plus_shard_same_dim_fallback(self): + """_StridedShard + Shard on same dim → materialize-then-split fallback.""" + tensor = torch.arange(16).reshape(4, 4).float() + expected = {0: tensor[[0]], 1: tensor[[2]], 2: tensor[[1]], 3: tensor[[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_prepacked_strided_shard_uses_contiguous_source_slice(self): + """Pre-concat w1/w3 tensors should shard contiguously before gate/up packing.""" + tensor = torch.arange(8).reshape(4, 2).float() + for rank, expected in [(0, tensor[:2]), (1, tensor[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(tensor, tensor_idx=0), expected, msg=f"rank {rank}") + + def test_expert_shaped_tp_only_no_expert_sharding(self): + """Expert-shaped param with TP on dim 1 but no expert sharding on dim 0 → regular _shard_nd path.""" + tensor = torch.arange(8).reshape(4, 2).float() + # Shard(1) on 3D param maps to dim 0 of the 2D checkpoint tensor (ndim_diff=1) + for rank, expected in [(0, tensor[:2]), (1, tensor[2:])]: + mesh = FakeMesh(shape=(2,), rank=rank) + op = _make_dtensor_shard_op(mesh, [Shard(1)], param_shape=(4, 4, 2), local_shape=(4, 2, 2)) + torch.testing.assert_close(op.shard_tensor(tensor, tensor_idx=0), expected, msg=f"rank {rank}") + + def test_expert_filtering(self): + """Mixtral-style experts: skip non-owned, return owned.""" + mesh = FakeMesh(shape=(2,), rank=1) + op = _make_dtensor_shard_op(mesh, [Shard(0)], param_shape=(4, 2, 2), local_shape=(2, 2, 2)) + expert_tensor = torch.ones(2, 2) + # rank 1 owns experts 2,3 (offset=2) + self.assertIsNone(op.shard_tensor(expert_tensor, tensor_idx=0)) + torch.testing.assert_close(op.shard_tensor(expert_tensor, tensor_idx=2), expert_tensor) + + def test_expert_filtering_preserves_inner_sharding(self): + """MoE expert ownership checks should still apply TP sharding on inner dims.""" + tensor = torch.arange(8).reshape(4, 2).float() + expected = { + 0: tensor[:2], + 1: tensor[2:], + 2: None, + 3: None, + } + 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(tensor, tensor_idx=1) + if expected[rank] is None: + self.assertIsNone(shard) + else: + torch.testing.assert_close(shard, expected[rank], msg=f"rank {rank}") + + def test_negative_dim_normalizes_correctly(self): + """Shard(-1) on a 2D tensor should shard the last dimension.""" + tensor = torch.arange(16).reshape(4, 4).float() + for rank, expected in [(0, tensor[:, :2]), (1, tensor[:, 2:])]: + mesh = FakeMesh(shape=(2,), rank=rank) + op = _make_dtensor_shard_op(mesh, [Shard(-1)], param_shape=(4, 4), local_shape=(4, 2)) + torch.testing.assert_close(op.shard_tensor(tensor), expected, msg=f"rank {rank}") + + def test_contiguous_shard_uneven_division(self): + """Shard(0) on 5 rows across 2 ranks → rank 0 gets 3 rows, rank 1 gets 2.""" + tensor = torch.arange(20).reshape(5, 4).float() + expected = {0: tensor[:3], 1: tensor[3:]} + for rank in range(2): + mesh = FakeMesh(shape=(2,), rank=rank) + local_rows = 3 if rank == 0 else 2 + op = _make_dtensor_shard_op(mesh, [Shard(0)], param_shape=(5, 4), local_shape=(local_rows, 4)) + torch.testing.assert_close(op.shard_tensor(tensor), expected[rank], msg=f"rank {rank}") + + def test_slice_and_read_all_single_ranges(self): + """When every dim has exactly one range, _slice_and_read takes the simple slice path (no concat).""" + 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, 4)) + dim_ranges = {0: [(0, 4)], 1: [(2, 6)]} + result = op._slice_and_read(tensor, [8, 8], dim_ranges, None, None) + torch.testing.assert_close(result, tensor[0:4, 2:6]) + + def test_slice_and_read_raises_on_two_multi_range_dims(self): + """Multiple disjoint ranges on two different dims → ValueError.""" + 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, 4)) + dim_ranges = {0: [(0, 2), (4, 6)], 1: [(0, 2), (4, 6)]} + with self.assertRaises(ValueError): + op._slice_and_read(tensor, [8, 8], dim_ranges, None, None) + + class TestConversionMapping(unittest.TestCase): def test_register_checkpoint_conversion_mapping(self): register_checkpoint_conversion_mapping( From d4400d5b52b719d9d3205b1205357e8cf09af829 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 14 Apr 2026 16:53:04 +0000 Subject: [PATCH 015/116] do monkey patching for rotary --- src/transformers/distributed/patches.py | 58 +++++++++++++++++++ .../integrations/tensor_parallel.py | 5 ++ .../models/afmoe/modeling_afmoe.py | 5 -- .../models/apertus/modeling_apertus.py | 5 -- .../models/arcee/modeling_arcee.py | 5 -- src/transformers/models/aria/modeling_aria.py | 5 -- .../models/bitnet/modeling_bitnet.py | 5 -- src/transformers/models/blt/modeling_blt.py | 5 -- .../models/chameleon/modeling_chameleon.py | 5 -- src/transformers/models/csm/modeling_csm.py | 5 -- src/transformers/models/cwm/modeling_cwm.py | 5 -- src/transformers/models/dbrx/modeling_dbrx.py | 5 -- .../deepseek_v3/modeling_deepseek_v3.py | 5 -- src/transformers/models/dia/modeling_dia.py | 5 -- .../models/diffllama/modeling_diffllama.py | 5 -- src/transformers/models/doge/modeling_doge.py | 5 -- .../models/dots1/modeling_dots1.py | 5 -- src/transformers/models/emu3/modeling_emu3.py | 5 -- .../models/eurobert/modeling_eurobert.py | 5 -- .../models/exaone4/modeling_exaone4.py | 5 -- .../models/exaone_moe/modeling_exaone_moe.py | 5 -- .../models/falcon/modeling_falcon.py | 5 -- .../models/falcon_h1/modeling_falcon_h1.py | 5 -- .../models/gemma/modeling_gemma.py | 5 -- .../models/gemma2/modeling_gemma2.py | 5 -- .../models/gemma3/modeling_gemma3.py | 5 -- .../glm4_moe_lite/modeling_glm4_moe_lite.py | 5 -- .../modeling_gpt_neox_japanese.py | 5 -- .../models/granite/modeling_granite.py | 5 -- .../models/granitemoe/modeling_granitemoe.py | 5 -- .../modeling_granitemoehybrid.py | 5 -- .../modeling_granitemoeshared.py | 5 -- .../higgs_audio_v2/modeling_higgs_audio_v2.py | 5 -- .../modeling_hunyuan_v1_dense.py | 5 -- .../hunyuan_v1_moe/modeling_hunyuan_v1_moe.py | 5 -- .../models/jais2/modeling_jais2.py | 5 -- .../models/jamba/modeling_jamba.py | 5 -- .../models/jetmoe/modeling_jetmoe.py | 5 -- .../modeling_jina_embeddings_v3.py | 5 -- .../modeling_kyutai_speech_to_text.py | 5 -- src/transformers/models/lasr/modeling_lasr.py | 5 -- src/transformers/models/lfm2/modeling_lfm2.py | 5 -- .../models/lfm2_moe/modeling_lfm2_moe.py | 5 -- .../models/llama/modeling_llama.py | 5 -- src/transformers/models/mimi/modeling_mimi.py | 5 -- .../models/minimax/modeling_minimax.py | 5 -- .../models/ministral/modeling_ministral.py | 5 -- .../models/ministral3/modeling_ministral3.py | 5 -- .../models/mistral/modeling_mistral.py | 5 -- .../models/mistral4/modeling_mistral4.py | 5 -- .../models/mixtral/modeling_mixtral.py | 5 -- .../models/mllama/modeling_mllama.py | 5 -- .../models/moshi/modeling_moshi.py | 5 -- .../models/nanochat/modeling_nanochat.py | 5 -- .../models/nemotron_h/modeling_nemotron_h.py | 5 -- .../models/nomic_bert/modeling_nomic_bert.py | 5 -- .../models/olmoe/modeling_olmoe.py | 5 -- .../models/parakeet/modeling_parakeet.py | 5 -- .../models/persimmon/modeling_persimmon.py | 5 -- src/transformers/models/phi/modeling_phi.py | 5 -- .../models/phimoe/modeling_phimoe.py | 5 -- .../models/qwen2/modeling_qwen2.py | 5 -- .../models/qwen2_moe/modeling_qwen2_moe.py | 5 -- .../models/qwen3/modeling_qwen3.py | 5 -- .../models/qwen3/modular_qwen3.py | 5 -- .../models/qwen3_moe/modeling_qwen3_moe.py | 5 -- .../qwen3_omni_moe/modeling_qwen3_omni_moe.py | 5 -- .../models/qwen3_vl/modeling_qwen3_vl.py | 5 -- .../qwen3_vl_moe/modeling_qwen3_vl_moe.py | 5 -- .../modeling_recurrent_gemma.py | 5 -- .../models/seed_oss/modeling_seed_oss.py | 5 -- .../models/smollm3/modeling_smollm3.py | 5 -- .../models/solar_open/modeling_solar_open.py | 5 -- .../models/stablelm/modeling_stablelm.py | 5 -- .../models/starcoder2/modeling_starcoder2.py | 5 -- .../models/t5gemma/modeling_t5gemma.py | 5 -- .../models/t5gemma2/modeling_t5gemma2.py | 5 -- .../models/timesfm2_5/modeling_timesfm2_5.py | 5 -- .../models/vaultgemma/modeling_vaultgemma.py | 5 -- .../modeling_voxtral_realtime.py | 5 -- .../models/youtu/modeling_youtu.py | 5 -- .../models/zamba2/modeling_zamba2.py | 5 -- 82 files changed, 63 insertions(+), 400 deletions(-) create mode 100644 src/transformers/distributed/patches.py diff --git a/src/transformers/distributed/patches.py b/src/transformers/distributed/patches.py new file mode 100644 index 000000000000..984498ef680a --- /dev/null +++ b/src/transformers/distributed/patches.py @@ -0,0 +1,58 @@ +# 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. +""" +Monkey-patches for DTensor-aware operations. + +These patches are applied at model-loading time (during ``from_pretrained``) +so that modeling files stay free of DTensor-specific code. +""" + +from __future__ import annotations + +import sys +from functools import wraps + +from torch.distributed.tensor import DTensor, Replicate + + +def _make_dtensor_rotary_wrapper(original_fn): + """Return a wrapper that converts cos/sin to replicated DTensors when q is a DTensor.""" + + @wraps(original_fn) + def _dtensor_apply_rotary_pos_emb(q, k, cos, sin, *args, **kwargs): + if isinstance(q, DTensor) and not isinstance(cos, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) + return original_fn(q, k, cos, sin, *args, **kwargs) + + return _dtensor_apply_rotary_pos_emb + + +def patch_dtensor_ops(model): + """Monkey-patch DTensor-aware wrappers onto the model's modeling module. + + Finds the Python module where the model class is defined and wraps + ``apply_rotary_pos_emb`` (if present) so that cos/sin tensors are + automatically promoted to replicated DTensors when the query is a DTensor. + + Called from ``apply_tensor_parallel`` after ``parallelize_module``. + """ + model_module = sys.modules.get(type(model).__module__) + if model_module is None: + return + + original_fn = getattr(model_module, "apply_rotary_pos_emb", None) + if original_fn is not None: + model_module.apply_rotary_pos_emb = _make_dtensor_rotary_wrapper(original_fn) diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index 759280defe8f..39ab9c0ec78c 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -28,6 +28,7 @@ from torch.distributed.tensor.parallel.style import ParallelStyle from torch.distributed.tensor.placement_types import _StridedShard +from ..distributed.patches import patch_dtensor_ops from ..utils import logging from ..utils.import_utils import is_torch_available @@ -631,6 +632,10 @@ def apply_tensor_parallel(model, tp_mesh, tp_plan): parallelize_module(model, tp_mesh, parallelize_plan) + # Patch DTensor-aware operations (e.g. rotary embeddings) onto the + # model's modeling module so modeling files stay free of DTensor code. + patch_dtensor_ops(model) + # 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 diff --git a/src/transformers/models/afmoe/modeling_afmoe.py b/src/transformers/models/afmoe/modeling_afmoe.py index 819c30446755..421119b33deb 100644 --- a/src/transformers/models/afmoe/modeling_afmoe.py +++ b/src/transformers/models/afmoe/modeling_afmoe.py @@ -23,7 +23,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -278,10 +277,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/apertus/modeling_apertus.py b/src/transformers/models/apertus/modeling_apertus.py index e616764f8fc8..7d14dd3d14c8 100644 --- a/src/transformers/models/apertus/modeling_apertus.py +++ b/src/transformers/models/apertus/modeling_apertus.py @@ -23,7 +23,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2CLS, ACT2FN from ...cache_utils import Cache, DynamicCache @@ -171,10 +170,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/arcee/modeling_arcee.py b/src/transformers/models/arcee/modeling_arcee.py index a30a2e03642b..8d2d05bf2952 100644 --- a/src/transformers/models/arcee/modeling_arcee.py +++ b/src/transformers/models/arcee/modeling_arcee.py @@ -23,7 +23,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from transformers.utils import auto_docstring @@ -176,10 +175,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/aria/modeling_aria.py b/src/transformers/models/aria/modeling_aria.py index 715c721b0b39..e66b12438940 100644 --- a/src/transformers/models/aria/modeling_aria.py +++ b/src/transformers/models/aria/modeling_aria.py @@ -23,7 +23,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -405,10 +404,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/bitnet/modeling_bitnet.py b/src/transformers/models/bitnet/modeling_bitnet.py index 78ae2a49b77b..14c1581b250f 100644 --- a/src/transformers/models/bitnet/modeling_bitnet.py +++ b/src/transformers/models/bitnet/modeling_bitnet.py @@ -22,7 +22,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache @@ -107,10 +106,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/blt/modeling_blt.py b/src/transformers/models/blt/modeling_blt.py index 5fa5262c78c4..778f7ba80cf6 100644 --- a/src/transformers/models/blt/modeling_blt.py +++ b/src/transformers/models/blt/modeling_blt.py @@ -25,7 +25,6 @@ import torch.distributions import torch.nn as nn import torch.nn.functional as F -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -287,10 +286,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/chameleon/modeling_chameleon.py b/src/transformers/models/chameleon/modeling_chameleon.py index 69e6cc834336..af69779959e4 100644 --- a/src/transformers/models/chameleon/modeling_chameleon.py +++ b/src/transformers/models/chameleon/modeling_chameleon.py @@ -21,7 +21,6 @@ import torch import torch.nn.functional as F from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache @@ -181,10 +180,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/csm/modeling_csm.py b/src/transformers/models/csm/modeling_csm.py index 2f03aeaf66ea..eb78dca8faf5 100644 --- a/src/transformers/models/csm/modeling_csm.py +++ b/src/transformers/models/csm/modeling_csm.py @@ -24,7 +24,6 @@ import torch import torch.nn as nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -227,10 +226,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/cwm/modeling_cwm.py b/src/transformers/models/cwm/modeling_cwm.py index e6f98a0ce250..3e0eb0504be0 100644 --- a/src/transformers/models/cwm/modeling_cwm.py +++ b/src/transformers/models/cwm/modeling_cwm.py @@ -24,7 +24,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache @@ -136,10 +135,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/dbrx/modeling_dbrx.py b/src/transformers/models/dbrx/modeling_dbrx.py index 0efbc1db2f9f..58735fb55c0b 100644 --- a/src/transformers/models/dbrx/modeling_dbrx.py +++ b/src/transformers/models/dbrx/modeling_dbrx.py @@ -23,7 +23,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -135,10 +134,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index 96c958b7d552..ab998cc99c21 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -11,7 +11,6 @@ import torch import torch.nn.functional as F from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -276,10 +275,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/dia/modeling_dia.py b/src/transformers/models/dia/modeling_dia.py index 4eba95828b53..629dfd4cdb35 100644 --- a/src/transformers/models/dia/modeling_dia.py +++ b/src/transformers/models/dia/modeling_dia.py @@ -23,7 +23,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -230,10 +229,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/diffllama/modeling_diffllama.py b/src/transformers/models/diffllama/modeling_diffllama.py index 8e5c98875c38..d80ccd572dc3 100644 --- a/src/transformers/models/diffllama/modeling_diffllama.py +++ b/src/transformers/models/diffllama/modeling_diffllama.py @@ -26,7 +26,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -163,10 +162,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/doge/modeling_doge.py b/src/transformers/models/doge/modeling_doge.py index 31157deba6d5..4aad59b52a9a 100644 --- a/src/transformers/models/doge/modeling_doge.py +++ b/src/transformers/models/doge/modeling_doge.py @@ -27,7 +27,6 @@ import torch import torch.nn.functional as F from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -165,10 +164,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/dots1/modeling_dots1.py b/src/transformers/models/dots1/modeling_dots1.py index b8a364b760c4..c474161e3043 100644 --- a/src/transformers/models/dots1/modeling_dots1.py +++ b/src/transformers/models/dots1/modeling_dots1.py @@ -23,7 +23,6 @@ import torch import torch.nn.functional as F from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -175,10 +174,6 @@ def rotate_half(x): def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/emu3/modeling_emu3.py b/src/transformers/models/emu3/modeling_emu3.py index 584e32e505a5..2481decd7aeb 100644 --- a/src/transformers/models/emu3/modeling_emu3.py +++ b/src/transformers/models/emu3/modeling_emu3.py @@ -28,7 +28,6 @@ import torch import torch.nn as nn import torch.nn.functional as F -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -86,10 +85,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/eurobert/modeling_eurobert.py b/src/transformers/models/eurobert/modeling_eurobert.py index 9fef764976ca..b93dd0649f14 100644 --- a/src/transformers/models/eurobert/modeling_eurobert.py +++ b/src/transformers/models/eurobert/modeling_eurobert.py @@ -24,7 +24,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss from ...activations import ACT2FN @@ -91,10 +90,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/exaone4/modeling_exaone4.py b/src/transformers/models/exaone4/modeling_exaone4.py index 9100d0ffefec..fab10b9b6937 100644 --- a/src/transformers/models/exaone4/modeling_exaone4.py +++ b/src/transformers/models/exaone4/modeling_exaone4.py @@ -24,7 +24,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache @@ -161,10 +160,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/exaone_moe/modeling_exaone_moe.py b/src/transformers/models/exaone_moe/modeling_exaone_moe.py index cad1ee6be2d2..2836a3c2245d 100644 --- a/src/transformers/models/exaone_moe/modeling_exaone_moe.py +++ b/src/transformers/models/exaone_moe/modeling_exaone_moe.py @@ -25,7 +25,6 @@ import torch import torch.nn as nn import torch.nn.functional as F -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -93,10 +92,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/falcon/modeling_falcon.py b/src/transformers/models/falcon/modeling_falcon.py index f9582748823b..016b3209b6b1 100644 --- a/src/transformers/models/falcon/modeling_falcon.py +++ b/src/transformers/models/falcon/modeling_falcon.py @@ -19,7 +19,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, LayerNorm, MSELoss from torch.nn import functional as F @@ -95,10 +94,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/falcon_h1/modeling_falcon_h1.py b/src/transformers/models/falcon_h1/modeling_falcon_h1.py index 33a8de337953..37b5da9df4b3 100644 --- a/src/transformers/models/falcon_h1/modeling_falcon_h1.py +++ b/src/transformers/models/falcon_h1/modeling_falcon_h1.py @@ -29,7 +29,6 @@ import torch import torch.nn.functional as F from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -147,10 +146,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/gemma/modeling_gemma.py b/src/transformers/models/gemma/modeling_gemma.py index 54be043b2eef..c6c5a55b8790 100644 --- a/src/transformers/models/gemma/modeling_gemma.py +++ b/src/transformers/models/gemma/modeling_gemma.py @@ -25,7 +25,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -191,10 +190,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/gemma2/modeling_gemma2.py b/src/transformers/models/gemma2/modeling_gemma2.py index 7c99443931c6..20673571b2d2 100644 --- a/src/transformers/models/gemma2/modeling_gemma2.py +++ b/src/transformers/models/gemma2/modeling_gemma2.py @@ -23,7 +23,6 @@ import torch import torch.nn as nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -176,10 +175,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/gemma3/modeling_gemma3.py b/src/transformers/models/gemma3/modeling_gemma3.py index ef8d1c885bca..3ecd6344dc07 100644 --- a/src/transformers/models/gemma3/modeling_gemma3.py +++ b/src/transformers/models/gemma3/modeling_gemma3.py @@ -24,7 +24,6 @@ import torch import torch.nn as nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -259,10 +258,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed 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 b4786fb7f23e..d59fd2ab996e 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 @@ -26,7 +26,6 @@ import torch import torch.nn as nn import torch.nn.functional as F -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -141,10 +140,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/gpt_neox_japanese/modeling_gpt_neox_japanese.py b/src/transformers/models/gpt_neox_japanese/modeling_gpt_neox_japanese.py index a2f423a94578..e334ce023d67 100755 --- a/src/transformers/models/gpt_neox_japanese/modeling_gpt_neox_japanese.py +++ b/src/transformers/models/gpt_neox_japanese/modeling_gpt_neox_japanese.py @@ -19,7 +19,6 @@ import torch from torch import Tensor, nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -148,10 +147,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/granite/modeling_granite.py b/src/transformers/models/granite/modeling_granite.py index c0c7765dcb20..934345fe6723 100644 --- a/src/transformers/models/granite/modeling_granite.py +++ b/src/transformers/models/granite/modeling_granite.py @@ -24,7 +24,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache @@ -70,10 +69,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/granitemoe/modeling_granitemoe.py b/src/transformers/models/granitemoe/modeling_granitemoe.py index ce188272d5c0..5fb53d6afe49 100644 --- a/src/transformers/models/granitemoe/modeling_granitemoe.py +++ b/src/transformers/models/granitemoe/modeling_granitemoe.py @@ -24,7 +24,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from torch.nn import functional as F from ... import initialization as init @@ -295,10 +294,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py index a55da168936b..dadffaea0072 100644 --- a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py +++ b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py @@ -23,7 +23,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from torch.nn import functional as F from ... import initialization as init @@ -76,10 +75,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py b/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py index 8f152f18f99f..71f8c6eaff7d 100644 --- a/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py +++ b/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py @@ -23,7 +23,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from torch.nn import functional as F from ... import initialization as init @@ -283,10 +282,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/higgs_audio_v2/modeling_higgs_audio_v2.py b/src/transformers/models/higgs_audio_v2/modeling_higgs_audio_v2.py index 03534ee50ad5..a0f106167721 100644 --- a/src/transformers/models/higgs_audio_v2/modeling_higgs_audio_v2.py +++ b/src/transformers/models/higgs_audio_v2/modeling_higgs_audio_v2.py @@ -24,7 +24,6 @@ import torch import torch.nn as nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -111,10 +110,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed 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 bf924c636482..d1652d78cbbc 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 @@ -23,7 +23,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from transformers.cache_utils import Cache @@ -110,10 +109,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed 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 0c61827781c9..19779da0528c 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 @@ -24,7 +24,6 @@ import torch import torch.nn.functional as F from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -113,10 +112,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/jais2/modeling_jais2.py b/src/transformers/models/jais2/modeling_jais2.py index 4f303e45be9d..5e6a37c0172d 100644 --- a/src/transformers/models/jais2/modeling_jais2.py +++ b/src/transformers/models/jais2/modeling_jais2.py @@ -24,7 +24,6 @@ import torch import torch.nn as nn -from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache @@ -84,10 +83,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/jamba/modeling_jamba.py b/src/transformers/models/jamba/modeling_jamba.py index b98efc90c175..ae618fb4a2b3 100755 --- a/src/transformers/models/jamba/modeling_jamba.py +++ b/src/transformers/models/jamba/modeling_jamba.py @@ -26,7 +26,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -103,10 +102,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/jetmoe/modeling_jetmoe.py b/src/transformers/models/jetmoe/modeling_jetmoe.py index 10de88e75f26..d3ee0bb14875 100644 --- a/src/transformers/models/jetmoe/modeling_jetmoe.py +++ b/src/transformers/models/jetmoe/modeling_jetmoe.py @@ -23,7 +23,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from torch.nn import functional as F from ... import initialization as init @@ -388,10 +387,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/jina_embeddings_v3/modeling_jina_embeddings_v3.py b/src/transformers/models/jina_embeddings_v3/modeling_jina_embeddings_v3.py index 163fc9e157e9..a55ffe0151c3 100644 --- a/src/transformers/models/jina_embeddings_v3/modeling_jina_embeddings_v3.py +++ b/src/transformers/models/jina_embeddings_v3/modeling_jina_embeddings_v3.py @@ -24,7 +24,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss from ... import initialization as init @@ -191,10 +190,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed 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 b60706bacbdd..b16274332baf 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 @@ -25,7 +25,6 @@ import torch import torch.nn as nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -387,10 +386,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/lasr/modeling_lasr.py b/src/transformers/models/lasr/modeling_lasr.py index eb9a2742b62f..7ecea9099410 100644 --- a/src/transformers/models/lasr/modeling_lasr.py +++ b/src/transformers/models/lasr/modeling_lasr.py @@ -24,7 +24,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...integrations import use_kernel_func_from_hub, use_kernelized_func @@ -161,10 +160,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/lfm2/modeling_lfm2.py b/src/transformers/models/lfm2/modeling_lfm2.py index c58329ead347..ef753e3b2893 100644 --- a/src/transformers/models/lfm2/modeling_lfm2.py +++ b/src/transformers/models/lfm2/modeling_lfm2.py @@ -23,7 +23,6 @@ import torch import torch.nn.functional as F from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin @@ -181,10 +180,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py b/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py index 3a881abbee90..0369ae31b8ae 100644 --- a/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py +++ b/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py @@ -24,7 +24,6 @@ import torch import torch.nn.functional as F from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...cache_utils import Cache, DynamicCache @@ -257,10 +256,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/llama/modeling_llama.py b/src/transformers/models/llama/modeling_llama.py index 3a966dc30798..9d659c7c6f08 100644 --- a/src/transformers/models/llama/modeling_llama.py +++ b/src/transformers/models/llama/modeling_llama.py @@ -21,7 +21,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache @@ -164,10 +163,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/mimi/modeling_mimi.py b/src/transformers/models/mimi/modeling_mimi.py index 92e5e22b8754..30480d1d1c03 100644 --- a/src/transformers/models/mimi/modeling_mimi.py +++ b/src/transformers/models/mimi/modeling_mimi.py @@ -20,7 +20,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -607,10 +606,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/minimax/modeling_minimax.py b/src/transformers/models/minimax/modeling_minimax.py index 9f35283377db..69497f83cad8 100644 --- a/src/transformers/models/minimax/modeling_minimax.py +++ b/src/transformers/models/minimax/modeling_minimax.py @@ -25,7 +25,6 @@ import torch import torch.nn.functional as F from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -352,10 +351,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/ministral/modeling_ministral.py b/src/transformers/models/ministral/modeling_ministral.py index e45e9f3616e5..af4f7fbeae59 100644 --- a/src/transformers/models/ministral/modeling_ministral.py +++ b/src/transformers/models/ministral/modeling_ministral.py @@ -24,7 +24,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache @@ -92,10 +91,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/ministral3/modeling_ministral3.py b/src/transformers/models/ministral3/modeling_ministral3.py index ce88bde32cd2..6aacf4c8ce3a 100644 --- a/src/transformers/models/ministral3/modeling_ministral3.py +++ b/src/transformers/models/ministral3/modeling_ministral3.py @@ -9,7 +9,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache @@ -61,10 +60,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/mistral/modeling_mistral.py b/src/transformers/models/mistral/modeling_mistral.py index 3d992a8e8aa3..b79dea36c9e9 100644 --- a/src/transformers/models/mistral/modeling_mistral.py +++ b/src/transformers/models/mistral/modeling_mistral.py @@ -9,7 +9,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache @@ -77,10 +76,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/mistral4/modeling_mistral4.py b/src/transformers/models/mistral4/modeling_mistral4.py index 1ce9b9cd3e4e..006ddad187bf 100644 --- a/src/transformers/models/mistral4/modeling_mistral4.py +++ b/src/transformers/models/mistral4/modeling_mistral4.py @@ -23,7 +23,6 @@ import torch import torch.nn.functional as F from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -284,10 +283,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/mixtral/modeling_mixtral.py b/src/transformers/models/mixtral/modeling_mixtral.py index a5f805083a46..991851dbadd3 100644 --- a/src/transformers/models/mixtral/modeling_mixtral.py +++ b/src/transformers/models/mixtral/modeling_mixtral.py @@ -29,7 +29,6 @@ import torch import torch.nn.functional as F from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -250,10 +249,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/mllama/modeling_mllama.py b/src/transformers/models/mllama/modeling_mllama.py index d2373afd68d8..3b9d12b9a225 100644 --- a/src/transformers/models/mllama/modeling_mllama.py +++ b/src/transformers/models/mllama/modeling_mllama.py @@ -20,7 +20,6 @@ import torch import torch.nn.functional as F from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -496,10 +495,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/moshi/modeling_moshi.py b/src/transformers/models/moshi/modeling_moshi.py index 51aeaaa50ffd..a967445c18ec 100644 --- a/src/transformers/models/moshi/modeling_moshi.py +++ b/src/transformers/models/moshi/modeling_moshi.py @@ -20,7 +20,6 @@ import torch import torch.nn as nn -from torch.distributed.tensor import DTensor, Replicate from torch.nn import CrossEntropyLoss from ... import initialization as init @@ -361,10 +360,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/nanochat/modeling_nanochat.py b/src/transformers/models/nanochat/modeling_nanochat.py index 7d351cf56306..bff55846a0a5 100644 --- a/src/transformers/models/nanochat/modeling_nanochat.py +++ b/src/transformers/models/nanochat/modeling_nanochat.py @@ -24,7 +24,6 @@ import torch import torch.nn as nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -181,10 +180,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/nemotron_h/modeling_nemotron_h.py b/src/transformers/models/nemotron_h/modeling_nemotron_h.py index 1cb01c08e1f8..9e264e5cfdcc 100644 --- a/src/transformers/models/nemotron_h/modeling_nemotron_h.py +++ b/src/transformers/models/nemotron_h/modeling_nemotron_h.py @@ -26,7 +26,6 @@ import torch import torch.nn.functional as F from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -780,10 +779,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/nomic_bert/modeling_nomic_bert.py b/src/transformers/models/nomic_bert/modeling_nomic_bert.py index 26ffe02fcfb7..f2836a1ec0f6 100644 --- a/src/transformers/models/nomic_bert/modeling_nomic_bert.py +++ b/src/transformers/models/nomic_bert/modeling_nomic_bert.py @@ -24,7 +24,6 @@ import torch import torch.nn as nn -from torch.distributed.tensor import DTensor, Replicate from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss from ... import initialization as init @@ -190,10 +189,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/olmoe/modeling_olmoe.py b/src/transformers/models/olmoe/modeling_olmoe.py index bf20e0a58c1e..5d89ec741529 100644 --- a/src/transformers/models/olmoe/modeling_olmoe.py +++ b/src/transformers/models/olmoe/modeling_olmoe.py @@ -22,7 +22,6 @@ import torch import torch.nn.functional as F from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -176,10 +175,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/parakeet/modeling_parakeet.py b/src/transformers/models/parakeet/modeling_parakeet.py index ff759581dab6..501a573f8494 100644 --- a/src/transformers/models/parakeet/modeling_parakeet.py +++ b/src/transformers/models/parakeet/modeling_parakeet.py @@ -24,7 +24,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -214,10 +213,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/persimmon/modeling_persimmon.py b/src/transformers/models/persimmon/modeling_persimmon.py index ae75f666cd18..e0516ed7da9a 100644 --- a/src/transformers/models/persimmon/modeling_persimmon.py +++ b/src/transformers/models/persimmon/modeling_persimmon.py @@ -23,7 +23,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache @@ -152,10 +151,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/phi/modeling_phi.py b/src/transformers/models/phi/modeling_phi.py index 23d69f11775d..e3f97a01ee4c 100644 --- a/src/transformers/models/phi/modeling_phi.py +++ b/src/transformers/models/phi/modeling_phi.py @@ -9,7 +9,6 @@ import torch import torch.nn as nn -from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache @@ -126,10 +125,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/phimoe/modeling_phimoe.py b/src/transformers/models/phimoe/modeling_phimoe.py index fafdbff37c3e..23bc944c522a 100644 --- a/src/transformers/models/phimoe/modeling_phimoe.py +++ b/src/transformers/models/phimoe/modeling_phimoe.py @@ -24,7 +24,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -150,10 +149,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/qwen2/modeling_qwen2.py b/src/transformers/models/qwen2/modeling_qwen2.py index 021543db455f..9263e1d42937 100644 --- a/src/transformers/models/qwen2/modeling_qwen2.py +++ b/src/transformers/models/qwen2/modeling_qwen2.py @@ -9,7 +9,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache @@ -142,10 +141,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py b/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py index eb7585373a61..d4150d0a74d7 100644 --- a/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py +++ b/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py @@ -29,7 +29,6 @@ import torch import torch.nn.functional as F from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -188,10 +187,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/qwen3/modeling_qwen3.py b/src/transformers/models/qwen3/modeling_qwen3.py index 550debc37eec..6eb68e4855d5 100644 --- a/src/transformers/models/qwen3/modeling_qwen3.py +++ b/src/transformers/models/qwen3/modeling_qwen3.py @@ -23,7 +23,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache @@ -195,10 +194,6 @@ def rotate_half(x): def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/qwen3/modular_qwen3.py b/src/transformers/models/qwen3/modular_qwen3.py index 74a8447d1bb7..c18e0e030079 100644 --- a/src/transformers/models/qwen3/modular_qwen3.py +++ b/src/transformers/models/qwen3/modular_qwen3.py @@ -16,7 +16,6 @@ from collections.abc import Callable import torch -from torch.distributed.tensor import DTensor, Replicate from ...cache_utils import Cache from ...modeling_flash_attention_utils import FlashAttentionKwargs @@ -54,10 +53,6 @@ def rotate_half(x): def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py b/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py index cb827352bd57..b4168d12df5c 100644 --- a/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py +++ b/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py @@ -24,7 +24,6 @@ import torch import torch.nn.functional as F from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -95,10 +94,6 @@ def rotate_half(x): def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed 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 7f1c1655771a..700e20e74a55 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 @@ -27,7 +27,6 @@ import numpy as np import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from torch.nn import Parameter from torch.nn import functional as F @@ -1442,10 +1441,6 @@ def extra_repr(self): def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py b/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py index 62e97f14b514..43946f06bf1e 100644 --- a/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py +++ b/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py @@ -26,7 +26,6 @@ import torch import torch.nn as nn import torch.nn.functional as F -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -409,10 +408,6 @@ def extra_repr(self): def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py b/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py index 4abf05fd8316..679a4357e305 100644 --- a/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py +++ b/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py @@ -26,7 +26,6 @@ import torch import torch.nn as nn import torch.nn.functional as F -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -186,10 +185,6 @@ def eager_attention_forward( def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py b/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py index 3f45e4ce4cf3..6e9c072b8860 100644 --- a/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py +++ b/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py @@ -20,7 +20,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -161,10 +160,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/seed_oss/modeling_seed_oss.py b/src/transformers/models/seed_oss/modeling_seed_oss.py index 861789644d6d..1ebc8f10a272 100644 --- a/src/transformers/models/seed_oss/modeling_seed_oss.py +++ b/src/transformers/models/seed_oss/modeling_seed_oss.py @@ -23,7 +23,6 @@ import torch import torch.nn as nn -from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache @@ -113,10 +112,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/smollm3/modeling_smollm3.py b/src/transformers/models/smollm3/modeling_smollm3.py index 2cf346b08d8d..8d911e414b0f 100644 --- a/src/transformers/models/smollm3/modeling_smollm3.py +++ b/src/transformers/models/smollm3/modeling_smollm3.py @@ -23,7 +23,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache @@ -140,10 +139,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/solar_open/modeling_solar_open.py b/src/transformers/models/solar_open/modeling_solar_open.py index 68fcb6a64dde..dfa30292455f 100644 --- a/src/transformers/models/solar_open/modeling_solar_open.py +++ b/src/transformers/models/solar_open/modeling_solar_open.py @@ -23,7 +23,6 @@ import torch import torch.nn.functional as F from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -249,10 +248,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/stablelm/modeling_stablelm.py b/src/transformers/models/stablelm/modeling_stablelm.py index 45f0b03e574c..9b9e0430e985 100755 --- a/src/transformers/models/stablelm/modeling_stablelm.py +++ b/src/transformers/models/stablelm/modeling_stablelm.py @@ -23,7 +23,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache @@ -151,10 +150,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/starcoder2/modeling_starcoder2.py b/src/transformers/models/starcoder2/modeling_starcoder2.py index d944e3afa422..8b89a1d1745c 100644 --- a/src/transformers/models/starcoder2/modeling_starcoder2.py +++ b/src/transformers/models/starcoder2/modeling_starcoder2.py @@ -28,7 +28,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache @@ -96,10 +95,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/t5gemma/modeling_t5gemma.py b/src/transformers/models/t5gemma/modeling_t5gemma.py index 95eec7fb9f29..a6b9b5392194 100644 --- a/src/transformers/models/t5gemma/modeling_t5gemma.py +++ b/src/transformers/models/t5gemma/modeling_t5gemma.py @@ -23,7 +23,6 @@ import torch import torch.nn as nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -190,10 +189,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/t5gemma2/modeling_t5gemma2.py b/src/transformers/models/t5gemma2/modeling_t5gemma2.py index 6e587b1a0e30..2e0dddc17876 100644 --- a/src/transformers/models/t5gemma2/modeling_t5gemma2.py +++ b/src/transformers/models/t5gemma2/modeling_t5gemma2.py @@ -24,7 +24,6 @@ import torch import torch.nn as nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -201,10 +200,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/timesfm2_5/modeling_timesfm2_5.py b/src/transformers/models/timesfm2_5/modeling_timesfm2_5.py index 6ae3b994b47b..e7b4e799d20b 100644 --- a/src/transformers/models/timesfm2_5/modeling_timesfm2_5.py +++ b/src/transformers/models/timesfm2_5/modeling_timesfm2_5.py @@ -26,7 +26,6 @@ import torch import torch.nn as nn import torch.nn.functional as F -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -230,10 +229,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/vaultgemma/modeling_vaultgemma.py b/src/transformers/models/vaultgemma/modeling_vaultgemma.py index 46e152403f71..f0a2e48d20b8 100644 --- a/src/transformers/models/vaultgemma/modeling_vaultgemma.py +++ b/src/transformers/models/vaultgemma/modeling_vaultgemma.py @@ -24,7 +24,6 @@ import torch import torch.nn as nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -108,10 +107,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py b/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py index 624e2ca7971d..07325b0ea559 100644 --- a/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py +++ b/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py @@ -26,7 +26,6 @@ import torch import torch.nn as nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -280,10 +279,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/youtu/modeling_youtu.py b/src/transformers/models/youtu/modeling_youtu.py index 6ed3bc109921..d40bef358da6 100644 --- a/src/transformers/models/youtu/modeling_youtu.py +++ b/src/transformers/models/youtu/modeling_youtu.py @@ -31,7 +31,6 @@ import torch import torch.nn.functional as F from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -181,10 +180,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/zamba2/modeling_zamba2.py b/src/transformers/models/zamba2/modeling_zamba2.py index b9eb255c399f..6e4ea7dcf2d8 100644 --- a/src/transformers/models/zamba2/modeling_zamba2.py +++ b/src/transformers/models/zamba2/modeling_zamba2.py @@ -25,7 +25,6 @@ import torch from torch import nn -from torch.distributed.tensor import DTensor, Replicate from ... import initialization as init from ...activations import ACT2FN @@ -220,10 +219,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if isinstance(q, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed From 6793503baea12a9c7072d909bb278c822063b7d4 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 14 Apr 2026 17:05:35 +0000 Subject: [PATCH 016/116] Revert modeling file diffs to match fsdp-core-model-loading base Restores modeling files to their base branch versions so the PR diff only shows the distributed/patches.py monkey-patch approach instead of noisy function moves in modeling files. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../models/dots1/modeling_dots1.py | 54 ++++++++++---- .../models/nanochat/modeling_nanochat.py | 52 ++++++------- .../models/qwen3/modeling_qwen3.py | 49 ++++++++---- .../models/qwen3_5/modeling_qwen3_5.py | 74 +++++++++---------- .../qwen3_5_moe/modeling_qwen3_5_moe.py | 74 +++++++++---------- .../models/qwen3_moe/modeling_qwen3_moe.py | 54 ++++++++++---- .../models/qwen3_next/modeling_qwen3_next.py | 74 +++++++++---------- .../qwen3_omni_moe/modeling_qwen3_omni_moe.py | 26 ++++++- .../models/qwen3_vl/modeling_qwen3_vl.py | 21 +++++- .../qwen3_vl_moe/modeling_qwen3_vl_moe.py | 26 ++++++- 10 files changed, 319 insertions(+), 185 deletions(-) diff --git a/src/transformers/models/dots1/modeling_dots1.py b/src/transformers/models/dots1/modeling_dots1.py index c474161e3043..399194648663 100644 --- a/src/transformers/models/dots1/modeling_dots1.py +++ b/src/transformers/models/dots1/modeling_dots1.py @@ -28,7 +28,12 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func +from ...integrations import ( + use_experts_implementation, + use_kernel_forward_from_hub, + use_kernel_func_from_hub, + use_kernelized_func, +) from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -128,6 +133,39 @@ def forward(self, x, position_ids): return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +@use_kernel_func_from_hub("rotary_pos_emb") +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, @@ -165,20 +203,6 @@ def eager_attention_forward( return attn_output, attn_weights -def rotate_half(x): - x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] - return torch.cat((-x2, x1), dim=-1) - - -def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): - cos = cos.unsqueeze(unsqueeze_dim) - sin = sin.unsqueeze(unsqueeze_dim) - q_embed = (q * cos) + (rotate_half(q) * sin) - k_embed = (k * cos) + (rotate_half(k) * sin) - return q_embed, k_embed - - @use_kernelized_func(apply_rotary_pos_emb) class Dots1Attention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" diff --git a/src/transformers/models/nanochat/modeling_nanochat.py b/src/transformers/models/nanochat/modeling_nanochat.py index bff55846a0a5..9205b89cd360 100644 --- a/src/transformers/models/nanochat/modeling_nanochat.py +++ b/src/transformers/models/nanochat/modeling_nanochat.py @@ -122,6 +122,32 @@ def forward(self, x, position_ids): return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) +@use_kernel_func_from_hub("rotary_pos_emb") +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, @@ -159,32 +185,6 @@ def eager_attention_forward( return attn_output, attn_weights -@use_kernel_func_from_hub("rotary_pos_emb") -def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): - """Applies Rotary Position Embedding to the query and key tensors. - - Args: - q (`torch.Tensor`): The query tensor. - k (`torch.Tensor`): The key tensor. - cos (`torch.Tensor`): The cosine part of the rotary embedding. - sin (`torch.Tensor`): The sine part of the rotary embedding. - unsqueeze_dim (`int`, *optional*, defaults to 1): - The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and - sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note - that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and - k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes - cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have - the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. - Returns: - `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. - """ - cos = cos.unsqueeze(unsqueeze_dim) - sin = sin.unsqueeze(unsqueeze_dim) - q_embed = (q * cos) + (rotate_half(q) * sin) - k_embed = (k * cos) + (rotate_half(k) * sin) - return q_embed, k_embed - - def rotate_half(x): """Rotates half the hidden dims of the input with flipped signs for NanoChat.""" x1 = x[..., : x.shape[-1] // 2] diff --git a/src/transformers/models/qwen3/modeling_qwen3.py b/src/transformers/models/qwen3/modeling_qwen3.py index 6eb68e4855d5..91715a33cf9d 100644 --- a/src/transformers/models/qwen3/modeling_qwen3.py +++ b/src/transformers/models/qwen3/modeling_qwen3.py @@ -27,7 +27,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import use_kernel_forward_from_hub, use_kernelized_func +from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -148,6 +148,39 @@ def forward(self, x, position_ids): return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +@use_kernel_func_from_hub("rotary_pos_emb") +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, @@ -185,20 +218,6 @@ def eager_attention_forward( return attn_output, attn_weights -def rotate_half(x): - x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] - return torch.cat((-x2, x1), dim=-1) - - -def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): - cos = cos.unsqueeze(unsqueeze_dim) - sin = sin.unsqueeze(unsqueeze_dim) - q_embed = (q * cos) + (rotate_half(q) * sin) - k_embed = (k * cos) + (rotate_half(k) * sin) - return q_embed, k_embed - - @use_kernelized_func(apply_rotary_pos_emb) class Qwen3Attention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" diff --git a/src/transformers/models/qwen3_5/modeling_qwen3_5.py b/src/transformers/models/qwen3_5/modeling_qwen3_5.py index 81fc2a40ea60..4dd3dfbaaf60 100644 --- a/src/transformers/models/qwen3_5/modeling_qwen3_5.py +++ b/src/transformers/models/qwen3_5/modeling_qwen3_5.py @@ -533,43 +533,6 @@ def forward( return output -def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: - """ - This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, - num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) - """ - batch, num_key_value_heads, slen, head_dim = hidden_states.shape - if n_rep == 1: - return hidden_states - hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) - return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) - - -def eager_attention_forward( - module: nn.Module, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attention_mask: torch.Tensor | None, - scaling: float, - dropout: float = 0.0, - **kwargs: Unpack[TransformersKwargs], -): - key_states = repeat_kv(key, module.num_key_value_groups) - value_states = repeat_kv(value, module.num_key_value_groups) - - attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling - if attention_mask is not None: - attn_weights = attn_weights + attention_mask - - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) - attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) - attn_output = torch.matmul(attn_weights, value_states) - attn_output = attn_output.transpose(1, 2).contiguous() - - return attn_output, attn_weights - - def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] @@ -616,6 +579,43 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): return q_embed, k_embed +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + @use_kernelized_func(apply_rotary_pos_emb) class Qwen3_5Attention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" 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 0e80ba9b5d5d..125ded124cf7 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 @@ -534,43 +534,6 @@ def forward( return output -def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: - """ - This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, - num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) - """ - batch, num_key_value_heads, slen, head_dim = hidden_states.shape - if n_rep == 1: - return hidden_states - hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) - return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) - - -def eager_attention_forward( - module: nn.Module, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attention_mask: torch.Tensor | None, - scaling: float, - dropout: float = 0.0, - **kwargs: Unpack[TransformersKwargs], -): - key_states = repeat_kv(key, module.num_key_value_groups) - value_states = repeat_kv(value, module.num_key_value_groups) - - attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling - if attention_mask is not None: - attn_weights = attn_weights + attention_mask - - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) - attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) - attn_output = torch.matmul(attn_weights, value_states) - attn_output = attn_output.transpose(1, 2).contiguous() - - return attn_output, attn_weights - - def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] @@ -617,6 +580,43 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): return q_embed, k_embed +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + @use_kernelized_func(apply_rotary_pos_emb) class Qwen3_5MoeAttention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" diff --git a/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py b/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py index b4168d12df5c..ddf84fc575b7 100644 --- a/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py +++ b/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py @@ -29,7 +29,12 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func +from ...integrations import ( + use_experts_implementation, + use_kernel_forward_from_hub, + use_kernel_func_from_hub, + use_kernelized_func, +) from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -48,6 +53,39 @@ from .configuration_qwen3_moe import Qwen3MoeConfig +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +@use_kernel_func_from_hub("rotary_pos_emb") +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, @@ -85,20 +123,6 @@ def eager_attention_forward( return attn_output, attn_weights -def rotate_half(x): - x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] - return torch.cat((-x2, x1), dim=-1) - - -def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): - cos = cos.unsqueeze(unsqueeze_dim) - sin = sin.unsqueeze(unsqueeze_dim) - q_embed = (q * cos) + (rotate_half(q) * sin) - k_embed = (k * cos) + (rotate_half(k) * sin) - return q_embed, k_embed - - @use_kernelized_func(apply_rotary_pos_emb) class Qwen3MoeAttention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" diff --git a/src/transformers/models/qwen3_next/modeling_qwen3_next.py b/src/transformers/models/qwen3_next/modeling_qwen3_next.py index af9239618f95..cd152e3d3e59 100644 --- a/src/transformers/models/qwen3_next/modeling_qwen3_next.py +++ b/src/transformers/models/qwen3_next/modeling_qwen3_next.py @@ -169,43 +169,6 @@ def extra_repr(self): return f"{tuple(self.weight.shape)}, eps={self.eps}" -def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: - """ - This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, - num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) - """ - batch, num_key_value_heads, slen, head_dim = hidden_states.shape - if n_rep == 1: - return hidden_states - hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) - return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) - - -def eager_attention_forward( - module: nn.Module, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attention_mask: torch.Tensor | None, - scaling: float, - dropout: float = 0.0, - **kwargs: Unpack[TransformersKwargs], -): - key_states = repeat_kv(key, module.num_key_value_groups) - value_states = repeat_kv(value, module.num_key_value_groups) - - attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling - if attention_mask is not None: - attn_weights = attn_weights + attention_mask - - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) - attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) - attn_output = torch.matmul(attn_weights, value_states) - attn_output = attn_output.transpose(1, 2).contiguous() - - return attn_output, attn_weights - - def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] @@ -252,6 +215,43 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): return q_embed, k_embed +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + @use_kernelized_func(apply_rotary_pos_emb) class Qwen3NextAttention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" 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 700e20e74a55..7b6c8b5b1bd4 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 @@ -34,7 +34,12 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func +from ...integrations import ( + use_experts_implementation, + use_kernel_forward_from_hub, + use_kernel_func_from_hub, + use_kernelized_func, +) from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -809,6 +814,7 @@ def _get_feat_extract_output_lengths(self, input_lengths: torch.LongTensor): def rotate_half(x): + """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) @@ -1438,7 +1444,25 @@ def extra_repr(self): return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" +@use_kernel_func_from_hub("rotary_pos_emb") def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) q_embed = (q * cos) + (rotate_half(q) * sin) diff --git a/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py b/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py index 43946f06bf1e..73678ee8c736 100644 --- a/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py +++ b/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py @@ -31,7 +31,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import use_kernel_forward_from_hub, use_kernelized_func +from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -122,6 +122,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: def rotate_half(x): + """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) @@ -405,7 +406,25 @@ def extra_repr(self): return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" +@use_kernel_func_from_hub("rotary_pos_emb") def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) q_embed = (q * cos) + (rotate_half(q) * sin) diff --git a/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py b/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py index 679a4357e305..ce405683fc94 100644 --- a/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py +++ b/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py @@ -31,7 +31,12 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func +from ...integrations import ( + use_experts_implementation, + use_kernel_forward_from_hub, + use_kernel_func_from_hub, + use_kernelized_func, +) from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -140,6 +145,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: def rotate_half(x): + """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) @@ -182,7 +188,25 @@ def eager_attention_forward( return attn_output, attn_weights +@use_kernel_func_from_hub("rotary_pos_emb") def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) q_embed = (q * cos) + (rotate_half(q) * sin) From b94351234790eefe75650a2d93aeb28beb3508a4 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 14 Apr 2026 17:21:39 +0000 Subject: [PATCH 017/116] Migrate all model TP plans from strings to TPStyle - Convert string plan values ("colwise", "rowwise", etc.) to TPStyle objects across 66+ model configs and modular files - Consolidate MoE expert sub-entries into TPStyle("moe_experts", ...) with shard_plan - Remove "replicated_with_grad_allreduce" entries (not needed for DTensor TP) - Migrate _tp_plan class attributes in modeling files from "colwise_gather_output" to TPStyle("colwise", "allgather") - Add TypeError in apply_tensor_parallel for unsupported plan values - Remove old TensorParallelLayer tests (API removed in DTensor refactor) - Regenerate auto-generated files via modular converter --- .../integrations/tensor_parallel.py | 8 +- .../models/afmoe/modeling_afmoe.py | 3 +- .../models/afmoe/modular_afmoe.py | 3 +- .../models/apertus/configuration_apertus.py | 15 +- .../models/apertus/modeling_apertus.py | 3 +- .../models/apertus/modular_apertus.py | 15 +- .../models/arcee/configuration_arcee.py | 13 +- .../models/arcee/modeling_arcee.py | 3 +- .../models/arcee/modular_arcee.py | 13 +- .../models/aria/configuration_aria.py | 15 +- src/transformers/models/aria/modeling_aria.py | 3 +- src/transformers/models/aria/modular_aria.py | 15 +- .../models/bamba/modeling_bamba.py | 3 +- .../models/cohere/configuration_cohere.py | 15 +- .../models/cohere/modeling_cohere.py | 3 +- .../models/cohere2/configuration_cohere2.py | 15 +- .../models/cohere2/modeling_cohere2.py | 3 +- .../models/cohere2/modular_cohere2.py | 15 +- src/transformers/models/cwm/modeling_cwm.py | 3 +- src/transformers/models/dbrx/modeling_dbrx.py | 3 +- src/transformers/models/dbrx/modular_dbrx.py | 3 +- .../deepseek_v2/configuration_deepseek_v2.py | 29 +- .../deepseek_v2/modeling_deepseek_v2.py | 3 +- .../models/deepseek_v2/modular_deepseek_v2.py | 29 +- .../deepseek_v3/configuration_deepseek_v3.py | 4 +- .../deepseek_v3/modeling_deepseek_v3.py | 3 +- .../models/diffllama/modeling_diffllama.py | 3 +- .../models/doge/configuration_doge.py | 23 +- src/transformers/models/doge/modeling_doge.py | 3 +- src/transformers/models/doge/modular_doge.py | 23 +- .../models/dots1/configuration_dots1.py | 31 +- .../models/dots1/modeling_dots1.py | 57 ++-- .../models/dots1/modular_dots1.py | 31 +- src/transformers/models/emu3/modeling_emu3.py | 3 +- .../models/ernie4_5/configuration_ernie4_5.py | 15 +- .../models/ernie4_5/modeling_ernie4_5.py | 3 +- .../configuration_ernie4_5_moe.py | 29 +- .../ernie4_5_moe/modeling_ernie4_5_moe.py | 3 +- .../configuration_ernie4_5_vl_moe.py | 29 +- .../modular_ernie4_5_vl_moe.py | 29 +- .../models/eurobert/modeling_eurobert.py | 3 +- .../models/eurobert/modular_eurobert.py | 3 +- .../models/exaone4/configuration_exaone4.py | 17 +- .../models/exaone4/modeling_exaone4.py | 3 +- .../models/exaone4/modular_exaone4.py | 17 +- .../exaone_moe/configuration_exaone_moe.py | 17 +- .../models/exaone_moe/modeling_exaone_moe.py | 3 +- .../models/falcon_h1/modeling_falcon_h1.py | 3 +- .../flex_olmo/configuration_flex_olmo.py | 25 +- .../models/flex_olmo/modeling_flex_olmo.py | 3 +- .../models/flex_olmo/modular_flex_olmo.py | 25 +- .../models/gemma/configuration_gemma.py | 15 +- .../models/gemma/modeling_gemma.py | 3 +- .../models/gemma/modular_gemma.py | 15 +- .../models/gemma2/configuration_gemma2.py | 15 +- .../models/gemma2/modeling_gemma2.py | 3 +- .../models/gemma2/modular_gemma2.py | 15 +- .../models/gemma3/configuration_gemma3.py | 17 +- .../models/gemma3/modeling_gemma3.py | 3 +- .../models/gemma3/modular_gemma3.py | 17 +- .../models/gemma3n/configuration_gemma3n.py | 18 +- .../models/gemma3n/modeling_gemma3n.py | 3 +- .../models/gemma3n/modular_gemma3n.py | 18 +- .../models/gemma4/configuration_gemma4.py | 41 ++- .../models/gemma4/modeling_gemma4.py | 3 +- src/transformers/models/glm/modeling_glm.py | 3 +- .../models/glm4/configuration_glm4.py | 15 +- src/transformers/models/glm4/modeling_glm4.py | 3 +- .../models/glm4_moe/configuration_glm4_moe.py | 29 +- .../models/glm4_moe/modeling_glm4_moe.py | 3 +- .../models/glm4_moe/modular_glm4_moe.py | 29 +- .../configuration_glm4_moe_lite.py | 21 +- .../glm4_moe_lite/modeling_glm4_moe_lite.py | 3 +- .../glm4_moe_lite/modular_glm4_moe_lite.py | 21 +- .../models/glm4v/configuration_glm4v.py | 15 +- .../models/glm4v/modular_glm4v.py | 15 +- .../glm4v_moe/configuration_glm4v_moe.py | 15 +- .../models/glm4v_moe/modular_glm4v_moe.py | 15 +- .../glm_image/configuration_glm_image.py | 15 +- .../glm_moe_dsa/configuration_glm_moe_dsa.py | 27 +- .../glm_moe_dsa/modeling_glm_moe_dsa.py | 3 +- .../models/glm_moe_dsa/modular_glm_moe_dsa.py | 27 +- .../models/glm_ocr/configuration_glm_ocr.py | 15 +- .../models/gpt_neox/configuration_gpt_neox.py | 9 +- .../models/gpt_neox/modeling_gpt_neox.py | 3 +- .../models/gpt_neox/modular_gpt_neox.py | 3 +- .../models/gpt_oss/modeling_gpt_oss.py | 3 +- .../models/granite/configuration_granite.py | 15 +- .../models/granite/modeling_granite.py | 3 +- .../models/granitemoe/modeling_granitemoe.py | 3 +- .../modeling_granitemoehybrid.py | 3 +- .../modeling_granitemoeshared.py | 3 +- .../models/helium/configuration_helium.py | 15 +- .../models/helium/modeling_helium.py | 3 +- .../modeling_hunyuan_v1_dense.py | 3 +- .../hunyuan_v1_moe/modeling_hunyuan_v1_moe.py | 3 +- .../models/jais2/configuration_jais2.py | 13 +- .../models/jais2/modeling_jais2.py | 3 +- .../models/jais2/modular_jais2.py | 13 +- .../models/jamba/modeling_jamba.py | 3 +- .../modeling_kyutai_speech_to_text.py | 3 +- src/transformers/models/lfm2/modeling_lfm2.py | 3 +- .../models/lfm2_moe/modeling_lfm2_moe.py | 3 +- .../models/llama/modeling_llama.py | 3 +- .../models/llama4/configuration_llama4.py | 39 ++- .../models/llama4/modeling_llama4.py | 3 +- .../configuration_longcat_flash.py | 21 +- .../longcat_flash/modeling_longcat_flash.py | 3 +- .../models/minimax/configuration_minimax.py | 17 +- .../models/minimax/modeling_minimax.py | 3 +- .../models/minimax/modular_minimax.py | 17 +- .../minimax_m2/configuration_minimax_m2.py | 17 +- .../models/minimax_m2/modeling_minimax_m2.py | 3 +- .../models/minimax_m2/modular_minimax_m2.py | 17 +- .../models/ministral/modeling_ministral.py | 3 +- .../ministral3/configuration_ministral3.py | 15 +- .../models/ministral3/modeling_ministral3.py | 3 +- .../models/mistral/modeling_mistral.py | 3 +- .../models/mistral4/configuration_mistral4.py | 21 +- .../models/mistral4/modeling_mistral4.py | 3 +- .../models/mixtral/configuration_mixtral.py | 8 +- .../models/mixtral/modeling_mixtral.py | 3 +- .../models/nanochat/configuration_nanochat.py | 13 +- .../models/nanochat/modeling_nanochat.py | 55 ++-- .../models/nanochat/modular_nanochat.py | 3 +- .../models/olmo/configuration_olmo.py | 15 +- src/transformers/models/olmo/modeling_olmo.py | 3 +- .../models/olmo2/configuration_olmo2.py | 23 +- .../models/olmo2/modeling_olmo2.py | 3 +- .../models/olmo2/modular_olmo2.py | 23 +- .../models/olmo3/configuration_olmo3.py | 23 +- .../models/olmo3/modeling_olmo3.py | 3 +- .../models/olmo3/modular_olmo3.py | 23 +- .../olmo_hybrid/configuration_olmo_hybrid.py | 23 +- .../olmo_hybrid/modeling_olmo_hybrid.py | 3 +- .../models/olmo_hybrid/modular_olmo_hybrid.py | 23 +- .../models/olmoe/configuration_olmoe.py | 17 +- .../models/olmoe/modeling_olmoe.py | 3 +- .../configuration_paddleocr_vl.py | 15 +- src/transformers/models/phi/modeling_phi.py | 3 +- .../models/phi3/configuration_phi3.py | 13 +- src/transformers/models/phi3/modeling_phi3.py | 3 +- .../configuration_phi4_multimodal.py | 13 +- .../modeling_phi4_multimodal.py | 3 +- .../models/phimoe/modeling_phimoe.py | 3 +- src/transformers/models/pi0/modeling_pi0.py | 3 +- src/transformers/models/pi0/modular_pi0.py | 3 +- .../models/qwen2/modeling_qwen2.py | 3 +- .../configuration_qwen2_5_omni.py | 15 +- .../qwen2_5_omni/modular_qwen2_5_omni.py | 15 +- .../qwen2_5_vl/configuration_qwen2_5_vl.py | 15 +- .../qwen2_moe/configuration_qwen2_moe.py | 15 +- .../models/qwen2_moe/modeling_qwen2_moe.py | 3 +- .../models/qwen2_moe/modular_qwen2_moe.py | 3 +- .../models/qwen2_vl/configuration_qwen2_vl.py | 15 +- .../models/qwen3/modeling_qwen3.py | 52 ++-- .../models/qwen3_5/configuration_qwen3_5.py | 17 +- .../models/qwen3_5/modeling_qwen3_5.py | 77 ++--- .../models/qwen3_5/modular_qwen3_5.py | 17 +- .../qwen3_5_moe/configuration_qwen3_5_moe.py | 25 +- .../qwen3_5_moe/modeling_qwen3_5_moe.py | 79 ++--- .../models/qwen3_5_moe/modular_qwen3_5_moe.py | 27 +- .../qwen3_moe/configuration_qwen3_moe.py | 25 +- .../models/qwen3_moe/modeling_qwen3_moe.py | 57 ++-- .../qwen3_next/configuration_qwen3_next.py | 31 +- .../models/qwen3_next/modeling_qwen3_next.py | 77 ++--- .../configuration_qwen3_omni_moe.py | 40 ++- .../qwen3_omni_moe/modeling_qwen3_omni_moe.py | 31 +- .../qwen3_omni_moe/modular_qwen3_omni_moe.py | 19 +- .../configuration_qwen3_vl_moe.py | 15 +- .../qwen3_vl_moe/modular_qwen3_vl_moe.py | 15 +- .../models/seed_oss/configuration_seed_oss.py | 15 +- .../models/seed_oss/modeling_seed_oss.py | 3 +- .../models/smollm3/configuration_smollm3.py | 15 +- .../models/smollm3/modeling_smollm3.py | 3 +- .../models/smollm3/modular_smollm3.py | 15 +- .../solar_open/configuration_solar_open.py | 17 +- .../models/solar_open/modeling_solar_open.py | 3 +- .../models/solar_open/modular_solar_open.py | 17 +- .../starcoder2/configuration_starcoder2.py | 13 +- .../models/starcoder2/modeling_starcoder2.py | 3 +- .../models/t5gemma/configuration_t5gemma.py | 15 +- .../models/t5gemma/modeling_t5gemma.py | 3 +- .../models/t5gemma/modular_t5gemma.py | 3 +- .../models/t5gemma2/configuration_t5gemma2.py | 33 +-- .../models/t5gemma2/modeling_t5gemma2.py | 3 +- .../models/t5gemma2/modular_t5gemma2.py | 3 +- .../vaultgemma/configuration_vaultgemma.py | 15 +- .../models/vaultgemma/modeling_vaultgemma.py | 3 +- .../modeling_voxtral_realtime.py | 3 +- .../models/youtu/configuration_youtu.py | 7 +- .../models/youtu/modeling_youtu.py | 3 +- .../models/youtu/modular_youtu.py | 7 +- tests/tensor_parallel/test_tensor_parallel.py | 277 ------------------ 194 files changed, 1404 insertions(+), 1491 deletions(-) diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index 39ab9c0ec78c..a4fd1cc38ca4 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -627,8 +627,14 @@ def apply_tensor_parallel(model, tp_mesh, tp_plan): # so _partition_fn can create DTensors with the correct placements. if isinstance(dtensor_style, MoEExpertsParallel) and style_value.shard_plan: dtensor_style._moe_shard_plan = style_value.shard_plan - else: + elif isinstance(style_value, ParallelStyle): parallelize_plan[name] = style_value + else: + raise TypeError( + f"Unsupported plan value for '{name}': {style_value!r} (type {type(style_value).__name__}). " + f"TP plan values must be TPStyle instances or ParallelStyle instances, not strings. " + f"Migrate string plan values to TPStyle (e.g., 'colwise' -> TPStyle('colwise', 'none'))." + ) parallelize_module(model, tp_mesh, parallelize_plan) diff --git a/src/transformers/models/afmoe/modeling_afmoe.py b/src/transformers/models/afmoe/modeling_afmoe.py index 421119b33deb..72366f370d90 100644 --- a/src/transformers/models/afmoe/modeling_afmoe.py +++ b/src/transformers/models/afmoe/modeling_afmoe.py @@ -34,6 +34,7 @@ use_kernel_func_from_hub, use_kernelized_func, ) +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast @@ -615,7 +616,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/afmoe/modular_afmoe.py b/src/transformers/models/afmoe/modular_afmoe.py index f3ff9f15b103..d07a2f1d2017 100644 --- a/src/transformers/models/afmoe/modular_afmoe.py +++ b/src/transformers/models/afmoe/modular_afmoe.py @@ -21,6 +21,7 @@ from ... import initialization as init from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast @@ -395,7 +396,7 @@ 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": TPStyle("colwise", "allgather")} _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..6864d11589f4 100644 --- a/src/transformers/models/apertus/configuration_apertus.py +++ b/src/transformers/models/apertus/configuration_apertus.py @@ -21,6 +21,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -46,14 +47,12 @@ class ApertusConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] default_theta = 12000000.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.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/apertus/modeling_apertus.py b/src/transformers/models/apertus/modeling_apertus.py index 7d14dd3d14c8..34bac7c18cf7 100644 --- a/src/transformers/models/apertus/modeling_apertus.py +++ b/src/transformers/models/apertus/modeling_apertus.py @@ -28,6 +28,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GenericForTokenClassification, GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -422,7 +423,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/apertus/modular_apertus.py b/src/transformers/models/apertus/modular_apertus.py index 3c9eb6d8b6ea..a901c3fba0fa 100644 --- a/src/transformers/models/apertus/modular_apertus.py +++ b/src/transformers/models/apertus/modular_apertus.py @@ -21,6 +21,7 @@ from ...activations import ACT2CLS from ...cache_utils import Cache from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...modeling_utils import ALL_ATTENTION_FUNCTIONS from ...processing_utils import Unpack @@ -64,14 +65,12 @@ class ApertusConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] default_theta = 12000000.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.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/arcee/configuration_arcee.py b/src/transformers/models/arcee/configuration_arcee.py index ee711d608204..0f6e8aa9034b 100644 --- a/src/transformers/models/arcee/configuration_arcee.py +++ b/src/transformers/models/arcee/configuration_arcee.py @@ -23,6 +23,7 @@ from transformers.utils import auto_docstring from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters @@ -46,12 +47,12 @@ class ArceeConfig(PreTrainedConfig): model_type = "arcee" 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.o_proj": "rowwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/arcee/modeling_arcee.py b/src/transformers/models/arcee/modeling_arcee.py index 8d2d05bf2952..06916f082039 100644 --- a/src/transformers/models/arcee/modeling_arcee.py +++ b/src/transformers/models/arcee/modeling_arcee.py @@ -30,6 +30,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import ( GenericForQuestionAnswering, @@ -424,7 +425,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/arcee/modular_arcee.py b/src/transformers/models/arcee/modular_arcee.py index 91cd8e13f1ed..316bb90db03a 100644 --- a/src/transformers/models/arcee/modular_arcee.py +++ b/src/transformers/models/arcee/modular_arcee.py @@ -17,6 +17,7 @@ from transformers.utils import auto_docstring, logging +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ..llama.configuration_llama import LlamaConfig from ..llama.modeling_llama import ( @@ -50,12 +51,12 @@ class ArceeConfig(LlamaConfig): model_type = "arcee" 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.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } vocab_size: int = 32000 diff --git a/src/transformers/models/aria/configuration_aria.py b/src/transformers/models/aria/configuration_aria.py index 7694656905dd..d240f4d5bba0 100644 --- a/src/transformers/models/aria/configuration_aria.py +++ b/src/transformers/models/aria/configuration_aria.py @@ -20,6 +20,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring from ...utils.type_validators import interval @@ -41,13 +42,13 @@ class AriaTextConfig(PreTrainedConfig): model_type = "aria_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.o_proj": "rowwise", - "layers.*.mlp.shared_experts.gate_proj": "colwise", - "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/aria/modeling_aria.py b/src/transformers/models/aria/modeling_aria.py index e66b12438940..7eebaec97e04 100644 --- a/src/transformers/models/aria/modeling_aria.py +++ b/src/transformers/models/aria/modeling_aria.py @@ -29,6 +29,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -758,7 +759,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: AriaTextConfig): diff --git a/src/transformers/models/aria/modular_aria.py b/src/transformers/models/aria/modular_aria.py index bfd5191e4135..d3484849b3a8 100644 --- a/src/transformers/models/aria/modular_aria.py +++ b/src/transformers/models/aria/modular_aria.py @@ -30,6 +30,7 @@ SizeDict, get_image_size, ) +from ...integrations.tensor_parallel import TPStyle from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_outputs import BaseModelOutputWithPooling from ...modeling_utils import PreTrainedModel @@ -110,13 +111,13 @@ class AriaTextConfig(LlamaConfig): model_type = "aria_text" base_config_key = "text_config" 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.shared_experts.gate_proj": "colwise", - "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), } intermediate_size: int = 4096 diff --git a/src/transformers/models/bamba/modeling_bamba.py b/src/transformers/models/bamba/modeling_bamba.py index 90129fc998b1..fd63eb7c58f1 100644 --- a/src/transformers/models/bamba/modeling_bamba.py +++ b/src/transformers/models/bamba/modeling_bamba.py @@ -35,6 +35,7 @@ from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernelized_func from ...integrations.hub_kernels import lazy_load_kernel +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -1071,7 +1072,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/cohere/configuration_cohere.py b/src/transformers/models/cohere/configuration_cohere.py index d52a3e008427..678b03eb8894 100644 --- a/src/transformers/models/cohere/configuration_cohere.py +++ b/src/transformers/models/cohere/configuration_cohere.py @@ -21,6 +21,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -50,13 +51,13 @@ class CohereConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] default_theta = 500000.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.o_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/cohere/modeling_cohere.py b/src/transformers/models/cohere/modeling_cohere.py index b8bf50af9bf4..80d50905bde0 100644 --- a/src/transformers/models/cohere/modeling_cohere.py +++ b/src/transformers/models/cohere/modeling_cohere.py @@ -36,6 +36,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -454,7 +455,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/cohere2/configuration_cohere2.py b/src/transformers/models/cohere2/configuration_cohere2.py index 48c2df360354..715749c51e0f 100644 --- a/src/transformers/models/cohere2/configuration_cohere2.py +++ b/src/transformers/models/cohere2/configuration_cohere2.py @@ -21,6 +21,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -49,13 +50,13 @@ class Cohere2Config(PreTrainedConfig): model_type = "cohere2" 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.o_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/cohere2/modeling_cohere2.py b/src/transformers/models/cohere2/modeling_cohere2.py index f43b2a0ef412..743031635387 100644 --- a/src/transformers/models/cohere2/modeling_cohere2.py +++ b/src/transformers/models/cohere2/modeling_cohere2.py @@ -28,6 +28,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -433,7 +434,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/cohere2/modular_cohere2.py b/src/transformers/models/cohere2/modular_cohere2.py index d19055a1b787..dd7421f320dd 100644 --- a/src/transformers/models/cohere2/modular_cohere2.py +++ b/src/transformers/models/cohere2/modular_cohere2.py @@ -20,6 +20,7 @@ from ...cache_utils import Cache, DynamicCache from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_outputs import BaseModelOutputWithPast from ...modeling_rope_utils import ( @@ -70,13 +71,13 @@ class Cohere2Config(PreTrainedConfig): model_type = "cohere2" 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.o_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/cwm/modeling_cwm.py b/src/transformers/models/cwm/modeling_cwm.py index 3e0eb0504be0..6e60b4ac31da 100644 --- a/src/transformers/models/cwm/modeling_cwm.py +++ b/src/transformers/models/cwm/modeling_cwm.py @@ -29,6 +29,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -426,7 +427,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/dbrx/modeling_dbrx.py b/src/transformers/models/dbrx/modeling_dbrx.py index 58735fb55c0b..7951f79b334f 100644 --- a/src/transformers/models/dbrx/modeling_dbrx.py +++ b/src/transformers/models/dbrx/modeling_dbrx.py @@ -29,6 +29,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_func_from_hub +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast @@ -642,7 +643,7 @@ 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": TPStyle("colwise", "allgather")} _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..34c1b3b6ac5c 100644 --- a/src/transformers/models/dbrx/modular_dbrx.py +++ b/src/transformers/models/dbrx/modular_dbrx.py @@ -23,6 +23,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import ( GradientCheckpointingLayer, @@ -430,7 +431,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: DbrxConfig): diff --git a/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py b/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py index 1b8005f8efef..626ba3a495e9 100644 --- a/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py +++ b/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py @@ -21,6 +21,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -52,20 +53,22 @@ class DeepseekV2Config(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise", - "layers.*.self_attn.q_b_proj": "colwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.q_b_proj": TPStyle("colwise", "none"), "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.*.mlp.shared_experts.gate_proj": "colwise", - "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.kv_b_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), + "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/deepseek_v2/modeling_deepseek_v2.py b/src/transformers/models/deepseek_v2/modeling_deepseek_v2.py index 3ef8266218f7..672b14742bb7 100644 --- a/src/transformers/models/deepseek_v2/modeling_deepseek_v2.py +++ b/src/transformers/models/deepseek_v2/modeling_deepseek_v2.py @@ -30,6 +30,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_experts_implementation, use_kernel_forward_from_hub +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GenericForSequenceClassification, GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -541,7 +542,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, 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..00683225189f 100644 --- a/src/transformers/models/deepseek_v2/modular_deepseek_v2.py +++ b/src/transformers/models/deepseek_v2/modular_deepseek_v2.py @@ -21,6 +21,7 @@ from ... import initialization as init from ...cache_utils import Cache +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters, dynamic_rope_update from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...utils import auto_docstring, logging @@ -67,20 +68,22 @@ class DeepseekV2Config(LlamaConfig): """ base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise", - "layers.*.self_attn.q_b_proj": "colwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.q_b_proj": TPStyle("colwise", "none"), "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.*.mlp.shared_experts.gate_proj": "colwise", - "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.kv_b_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), + "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } 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 a9216f4db2ab..651a2fc2af5f 100644 --- a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py @@ -54,8 +54,8 @@ class DeepseekV3Config(PreTrainedConfig): "moe_experts", "allreduce", shard_plan={ - "gate_up_proj": "packed_colwise", - "down_proj": "rowwise", + "gate_up_proj": TPStyle("packed_colwise", "none"), + "down_proj": TPStyle("rowwise", "allreduce"), }, ), "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index ab998cc99c21..fdf708ee9cfa 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -17,6 +17,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernel_func_from_hub +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -634,7 +635,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/diffllama/modeling_diffllama.py b/src/transformers/models/diffllama/modeling_diffllama.py index d80ccd572dc3..6c5b12cb6850 100644 --- a/src/transformers/models/diffllama/modeling_diffllama.py +++ b/src/transformers/models/diffllama/modeling_diffllama.py @@ -32,6 +32,7 @@ from ...cache_utils import Cache, DynamicCache, StaticCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import _flash_attention_forward, flash_attn_supports_top_left_mask from ...modeling_layers import ( @@ -660,7 +661,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/doge/configuration_doge.py b/src/transformers/models/doge/configuration_doge.py index 8518d9021458..e4ccda76fafc 100644 --- a/src/transformers/models/doge/configuration_doge.py +++ b/src/transformers/models/doge/configuration_doge.py @@ -22,6 +22,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -52,17 +53,17 @@ class DogeConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `DogeModel` 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.dt_proj": "rowwise", - "layers.*.self_attn.o_proj": "rowwise", - "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.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.dt_proj": TPStyle("rowwise", "allreduce"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.router_gate": TPStyle("colwise", "allgather"), + "layers.*.mlp.down_embed": TPStyle("vocab", "allreduce"), + "layers.*.mlp.up_embed": TPStyle("vocab", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/doge/modeling_doge.py b/src/transformers/models/doge/modeling_doge.py index 4aad59b52a9a..ebb6b4bc992c 100644 --- a/src/transformers/models/doge/modeling_doge.py +++ b/src/transformers/models/doge/modeling_doge.py @@ -34,6 +34,7 @@ from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub from ...integrations.flex_attention import compile_friendly_flex_attention +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_layers import GenericForSequenceClassification, GradientCheckpointingLayer from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast @@ -716,7 +717,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/doge/modular_doge.py b/src/transformers/models/doge/modular_doge.py index 8b78126c0a00..cf92f343eb26 100644 --- a/src/transformers/models/doge/modular_doge.py +++ b/src/transformers/models/doge/modular_doge.py @@ -29,6 +29,7 @@ from ...cache_utils import Cache from ...configuration_utils import PreTrainedConfig from ...integrations.flex_attention import compile_friendly_flex_attention +from ...integrations.tensor_parallel import TPStyle from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast from ...modeling_rope_utils import RopeParameters @@ -81,17 +82,17 @@ class DogeConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `DogeModel` 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.dt_proj": "rowwise", - "layers.*.self_attn.o_proj": "rowwise", - "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.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.dt_proj": TPStyle("rowwise", "allreduce"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.router_gate": TPStyle("colwise", "allgather"), + "layers.*.mlp.down_embed": TPStyle("vocab", "allreduce"), + "layers.*.mlp.up_embed": TPStyle("vocab", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/dots1/configuration_dots1.py b/src/transformers/models/dots1/configuration_dots1.py index 4d568bf4a565..e452acf1802c 100644 --- a/src/transformers/models/dots1/configuration_dots1.py +++ b/src/transformers/models/dots1/configuration_dots1.py @@ -20,6 +20,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -48,21 +49,21 @@ class Dots1Config(PreTrainedConfig): 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.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.*.mlp.shared_experts.gate_proj": "colwise", - "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), + "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { diff --git a/src/transformers/models/dots1/modeling_dots1.py b/src/transformers/models/dots1/modeling_dots1.py index 399194648663..e77ec0940223 100644 --- a/src/transformers/models/dots1/modeling_dots1.py +++ b/src/transformers/models/dots1/modeling_dots1.py @@ -28,12 +28,8 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import ( - use_experts_implementation, - use_kernel_forward_from_hub, - use_kernel_func_from_hub, - use_kernelized_func, -) +from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -133,39 +129,6 @@ def forward(self, x, position_ids): return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) -def rotate_half(x): - """Rotates half the hidden dims of the input.""" - x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] - return torch.cat((-x2, x1), dim=-1) - - -@use_kernel_func_from_hub("rotary_pos_emb") -def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): - """Applies Rotary Position Embedding to the query and key tensors. - - Args: - q (`torch.Tensor`): The query tensor. - k (`torch.Tensor`): The key tensor. - cos (`torch.Tensor`): The cosine part of the rotary embedding. - sin (`torch.Tensor`): The sine part of the rotary embedding. - unsqueeze_dim (`int`, *optional*, defaults to 1): - The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and - sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note - that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and - k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes - cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have - the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. - Returns: - `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. - """ - cos = cos.unsqueeze(unsqueeze_dim) - sin = sin.unsqueeze(unsqueeze_dim) - q_embed = (q * cos) + (rotate_half(q) * sin) - k_embed = (k * cos) + (rotate_half(k) * sin) - return q_embed, k_embed - - def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, @@ -203,6 +166,20 @@ def eager_attention_forward( return attn_output, attn_weights +def rotate_half(x): + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + @use_kernelized_func(apply_rotary_pos_emb) class Dots1Attention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" @@ -569,7 +546,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/dots1/modular_dots1.py b/src/transformers/models/dots1/modular_dots1.py index d390037d0820..86a9093a49b9 100644 --- a/src/transformers/models/dots1/modular_dots1.py +++ b/src/transformers/models/dots1/modular_dots1.py @@ -15,6 +15,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_outputs import CausalLMOutputWithPast from ...modeling_rope_utils import RopeParameters from ...processing_utils import Unpack @@ -63,21 +64,21 @@ class Dots1Config(PreTrainedConfig): 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.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.*.mlp.shared_experts.gate_proj": "colwise", - "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), + "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { diff --git a/src/transformers/models/emu3/modeling_emu3.py b/src/transformers/models/emu3/modeling_emu3.py index 2481decd7aeb..ec4890f5caa9 100644 --- a/src/transformers/models/emu3/modeling_emu3.py +++ b/src/transformers/models/emu3/modeling_emu3.py @@ -34,6 +34,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, CausalLMOutputWithPast @@ -1275,7 +1276,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} config: Emu3TextConfig diff --git a/src/transformers/models/ernie4_5/configuration_ernie4_5.py b/src/transformers/models/ernie4_5/configuration_ernie4_5.py index 896fb0b99402..f32929da5b70 100644 --- a/src/transformers/models/ernie4_5/configuration_ernie4_5.py +++ b/src/transformers/models/ernie4_5/configuration_ernie4_5.py @@ -16,6 +16,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -47,13 +48,13 @@ class Ernie4_5Config(PreTrainedConfig): default_theta = 500000.0 # Default tensor parallel plan for base model `Ernie4_5Model` 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", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/ernie4_5/modeling_ernie4_5.py b/src/transformers/models/ernie4_5/modeling_ernie4_5.py index ae533c7c6ef8..e8b33e2ade89 100644 --- a/src/transformers/models/ernie4_5/modeling_ernie4_5.py +++ b/src/transformers/models/ernie4_5/modeling_ernie4_5.py @@ -28,6 +28,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -422,7 +423,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, 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..22b38d1d5d33 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 @@ -16,6 +16,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -63,19 +64,21 @@ class Ernie4_5_MoeConfig(PreTrainedConfig): # Default tensor parallel plan for base model `Ernie4_5_MoE` 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.*.mlp.shared_experts.gate_proj": "colwise", - "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), + "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py b/src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py index 9c106a90010d..7e2c863211db 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 @@ -30,6 +30,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast @@ -654,7 +655,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, 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..610b9647ee75 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 @@ -20,6 +20,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring, logging @@ -49,10 +50,10 @@ class Ernie4_5_VLMoeVisionConfig(PreTrainedConfig): initializer_range: float = 0.02 base_model_tp_plan = { - "blocks.*.attn.qkv": "colwise", - "blocks.*.attn.proj": "rowwise", - "blocks.*.mlp.fc1": "colwise", - "blocks.*.mlp.fc2": "rowwise", + "blocks.*.attn.qkv": TPStyle("colwise", "none"), + "blocks.*.attn.proj": TPStyle("rowwise", "allreduce"), + "blocks.*.mlp.fc1": TPStyle("colwise", "none"), + "blocks.*.mlp.fc2": TPStyle("rowwise", "allreduce"), } intermediate_size: int = 4 * 1280 temporal_merge_size: int = 2 @@ -83,16 +84,16 @@ class Ernie4_5_VLMoeTextConfig(PreTrainedConfig): default_theta = 500000.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.o_proj": "rowwise", - "layers.*.mlp.shared_experts.gate_proj": "colwise", - "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), 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 42bbb44b70a5..43e5d780f8c9 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 @@ -37,6 +37,7 @@ PILImageResampling, SizeDict, ) +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -93,10 +94,10 @@ class Ernie4_5_VLMoeVisionConfig(Qwen2VLVisionConfig): model_type = "ernie4_5_vl_moe_vision" base_model_tp_plan = { - "blocks.*.attn.qkv": "colwise", - "blocks.*.attn.proj": "rowwise", - "blocks.*.mlp.fc1": "colwise", - "blocks.*.mlp.fc2": "rowwise", + "blocks.*.attn.qkv": TPStyle("colwise", "none"), + "blocks.*.attn.proj": TPStyle("rowwise", "allreduce"), + "blocks.*.mlp.fc1": TPStyle("colwise", "none"), + "blocks.*.mlp.fc2": TPStyle("rowwise", "allreduce"), } hidden_size: int = 1280 @@ -131,16 +132,16 @@ class Ernie4_5_VLMoeTextConfig(Ernie4_5_MoeConfig): base_config_key = "text_config" 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.shared_experts.gate_proj": "colwise", - "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } ignore_keys_at_rope_validation = {"mrope_section"} diff --git a/src/transformers/models/eurobert/modeling_eurobert.py b/src/transformers/models/eurobert/modeling_eurobert.py index b93dd0649f14..ca6c39acb238 100644 --- a/src/transformers/models/eurobert/modeling_eurobert.py +++ b/src/transformers/models/eurobert/modeling_eurobert.py @@ -29,6 +29,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_bidirectional_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, MaskedLMOutput, SequenceClassifierOutput, TokenClassifierOutput @@ -408,7 +409,7 @@ 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": TPStyle("colwise", "allgather")} _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 f0a2f1b0b479..588508230dfb 100644 --- a/src/transformers/models/eurobert/modular_eurobert.py +++ b/src/transformers/models/eurobert/modular_eurobert.py @@ -18,6 +18,7 @@ from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss from ...configuration_utils import strict +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_bidirectional_mask from ...modeling_outputs import BaseModelOutput, MaskedLMOutput, SequenceClassifierOutput, TokenClassifierOutput from ...modeling_rope_utils import RopeParameters @@ -141,7 +142,7 @@ 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": TPStyle("colwise", "allgather")} _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..89ee40135153 100644 --- a/src/transformers/models/exaone4/configuration_exaone4.py +++ b/src/transformers/models/exaone4/configuration_exaone4.py @@ -21,6 +21,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -60,15 +61,13 @@ class Exaone4Config(PreTrainedConfig): 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.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/exaone4/modeling_exaone4.py b/src/transformers/models/exaone4/modeling_exaone4.py index fab10b9b6937..0155d7fa3a9f 100644 --- a/src/transformers/models/exaone4/modeling_exaone4.py +++ b/src/transformers/models/exaone4/modeling_exaone4.py @@ -29,6 +29,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_layers import ( GenericForQuestionAnswering, @@ -440,7 +441,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/exaone4/modular_exaone4.py b/src/transformers/models/exaone4/modular_exaone4.py index c6d9202170a0..cc152edb42dc 100644 --- a/src/transformers/models/exaone4/modular_exaone4.py +++ b/src/transformers/models/exaone4/modular_exaone4.py @@ -22,6 +22,7 @@ from ...cache_utils import Cache, DynamicCache from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_outputs import ( BaseModelOutputWithPast, @@ -89,15 +90,13 @@ class Exaone4Config(PreTrainedConfig): 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.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/exaone_moe/configuration_exaone_moe.py b/src/transformers/models/exaone_moe/configuration_exaone_moe.py index 81ded9366cdb..1f948a9eb5fb 100644 --- a/src/transformers/models/exaone_moe/configuration_exaone_moe.py +++ b/src/transformers/models/exaone_moe/configuration_exaone_moe.py @@ -21,6 +21,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...utils import auto_docstring @@ -66,15 +67,13 @@ class ExaoneMoeConfig(PreTrainedConfig): 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.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/exaone_moe/modeling_exaone_moe.py b/src/transformers/models/exaone_moe/modeling_exaone_moe.py index 2836a3c2245d..abe193821de9 100644 --- a/src/transformers/models/exaone_moe/modeling_exaone_moe.py +++ b/src/transformers/models/exaone_moe/modeling_exaone_moe.py @@ -31,6 +31,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernel_func_from_hub +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -563,7 +564,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/falcon_h1/modeling_falcon_h1.py b/src/transformers/models/falcon_h1/modeling_falcon_h1.py index 37b5da9df4b3..8913db392fb1 100644 --- a/src/transformers/models/falcon_h1/modeling_falcon_h1.py +++ b/src/transformers/models/falcon_h1/modeling_falcon_h1.py @@ -36,6 +36,7 @@ from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func from ...integrations.hub_kernels import lazy_load_kernel +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -1166,7 +1167,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, 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..9f8991d84000 100644 --- a/src/transformers/models/flex_olmo/configuration_flex_olmo.py +++ b/src/transformers/models/flex_olmo/configuration_flex_olmo.py @@ -22,6 +22,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -50,13 +51,23 @@ 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": TPStyle( + "colwise", "allgather" + ), # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.k_proj": TPStyle( + "colwise", "allgather" + ), # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.v_proj": TPStyle( + "colwise", "allgather" + ), # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.o_proj": TPStyle( + "vocab", "allreduce" + ), # input is replicated due to the added norm on q and k + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/flex_olmo/modeling_flex_olmo.py b/src/transformers/models/flex_olmo/modeling_flex_olmo.py index 100e6fa35554..3f38fb9e328c 100644 --- a/src/transformers/models/flex_olmo/modeling_flex_olmo.py +++ b/src/transformers/models/flex_olmo/modeling_flex_olmo.py @@ -31,6 +31,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast @@ -597,7 +598,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, 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..8a496547c204 100644 --- a/src/transformers/models/flex_olmo/modular_flex_olmo.py +++ b/src/transformers/models/flex_olmo/modular_flex_olmo.py @@ -18,6 +18,7 @@ from ...cache_utils import Cache, DynamicCache from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_outputs import MoeModelOutputWithPast from ...modeling_rope_utils import RopeParameters @@ -60,13 +61,23 @@ 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": TPStyle( + "colwise", "allgather" + ), # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.k_proj": TPStyle( + "colwise", "allgather" + ), # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.v_proj": TPStyle( + "colwise", "allgather" + ), # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.o_proj": TPStyle( + "vocab", "allreduce" + ), # input is replicated due to the added norm on q and k + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/gemma/configuration_gemma.py b/src/transformers/models/gemma/configuration_gemma.py index 0ef2e8b31b8d..bb676b5d8486 100644 --- a/src/transformers/models/gemma/configuration_gemma.py +++ b/src/transformers/models/gemma/configuration_gemma.py @@ -23,6 +23,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -47,13 +48,13 @@ class GemmaConfig(PreTrainedConfig): model_type = "gemma" 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.o_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/gemma/modeling_gemma.py b/src/transformers/models/gemma/modeling_gemma.py index c6c5a55b8790..c26458cfbc54 100644 --- a/src/transformers/models/gemma/modeling_gemma.py +++ b/src/transformers/models/gemma/modeling_gemma.py @@ -31,6 +31,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_func_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import ( GenericForSequenceClassification, @@ -450,7 +451,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/gemma/modular_gemma.py b/src/transformers/models/gemma/modular_gemma.py index 25f436473fbe..9a777a6be8ad 100644 --- a/src/transformers/models/gemma/modular_gemma.py +++ b/src/transformers/models/gemma/modular_gemma.py @@ -21,6 +21,7 @@ from ... import initialization as init from ...cache_utils import Cache, DynamicCache from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_outputs import BaseModelOutputWithPast from ...modeling_rope_utils import RopeParameters @@ -66,13 +67,13 @@ class GemmaConfig(PreTrainedConfig): model_type = "gemma" 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.o_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/gemma2/configuration_gemma2.py b/src/transformers/models/gemma2/configuration_gemma2.py index 11d7b012099f..852eb0e54e63 100644 --- a/src/transformers/models/gemma2/configuration_gemma2.py +++ b/src/transformers/models/gemma2/configuration_gemma2.py @@ -21,6 +21,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -51,13 +52,13 @@ class Gemma2Config(PreTrainedConfig): model_type = "gemma2" 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.o_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/gemma2/modeling_gemma2.py b/src/transformers/models/gemma2/modeling_gemma2.py index 20673571b2d2..83812655fb12 100644 --- a/src/transformers/models/gemma2/modeling_gemma2.py +++ b/src/transformers/models/gemma2/modeling_gemma2.py @@ -29,6 +29,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_func_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -476,7 +477,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/gemma2/modular_gemma2.py b/src/transformers/models/gemma2/modular_gemma2.py index 2edd9ef5f101..191ce9e14401 100644 --- a/src/transformers/models/gemma2/modular_gemma2.py +++ b/src/transformers/models/gemma2/modular_gemma2.py @@ -21,6 +21,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -78,13 +79,13 @@ class Gemma2Config(PreTrainedConfig): model_type = "gemma2" 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.o_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/gemma3/configuration_gemma3.py b/src/transformers/models/gemma3/configuration_gemma3.py index f25c9cef21fd..f7699c7cd994 100644 --- a/src/transformers/models/gemma3/configuration_gemma3.py +++ b/src/transformers/models/gemma3/configuration_gemma3.py @@ -23,6 +23,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...utils import auto_docstring, logging from ..siglip import SiglipVisionConfig @@ -58,15 +59,13 @@ class Gemma3TextConfig(PreTrainedConfig): model_type = "gemma3_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", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/gemma3/modeling_gemma3.py b/src/transformers/models/gemma3/modeling_gemma3.py index 3ecd6344dc07..1f8207f8b534 100644 --- a/src/transformers/models/gemma3/modeling_gemma3.py +++ b/src/transformers/models/gemma3/modeling_gemma3.py @@ -31,6 +31,7 @@ from ...configuration_utils import PreTrainedConfig from ...generation import GenerationMixin from ...integrations import use_kernel_func_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_masks_for_generate, create_sliding_window_causal_mask from ...modeling_layers import GenericForSequenceClassification, GradientCheckpointingLayer from ...modeling_outputs import ( @@ -591,7 +592,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} config: Gemma3TextConfig diff --git a/src/transformers/models/gemma3/modular_gemma3.py b/src/transformers/models/gemma3/modular_gemma3.py index 1e96f5acceb9..1ce5f8fa443c 100644 --- a/src/transformers/models/gemma3/modular_gemma3.py +++ b/src/transformers/models/gemma3/modular_gemma3.py @@ -22,6 +22,7 @@ from ... import initialization as init from ...cache_utils import Cache, DynamicCache from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_masks_for_generate, create_sliding_window_causal_mask from ...modeling_layers import GenericForSequenceClassification, GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, SequenceClassifierOutputWithPast @@ -86,15 +87,13 @@ class Gemma3TextConfig(Gemma2Config, PreTrainedConfig): model_type = "gemma3_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.o_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } default_theta = {"global": 1_000_000.0, "local": 10_000.0} diff --git a/src/transformers/models/gemma3n/configuration_gemma3n.py b/src/transformers/models/gemma3n/configuration_gemma3n.py index d78002ed76c3..21627324d286 100644 --- a/src/transformers/models/gemma3n/configuration_gemma3n.py +++ b/src/transformers/models/gemma3n/configuration_gemma3n.py @@ -24,6 +24,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...utils import auto_docstring, is_timm_available, logging, requires_backends @@ -79,16 +80,13 @@ 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.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/gemma3n/modeling_gemma3n.py b/src/transformers/models/gemma3n/modeling_gemma3n.py index 3a41bb261c43..ac502610284f 100644 --- a/src/transformers/models/gemma3n/modeling_gemma3n.py +++ b/src/transformers/models/gemma3n/modeling_gemma3n.py @@ -32,6 +32,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, CausalLMOutputWithPast @@ -1769,7 +1770,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} config: Gemma3nTextConfig diff --git a/src/transformers/models/gemma3n/modular_gemma3n.py b/src/transformers/models/gemma3n/modular_gemma3n.py index c531291fc584..41364bd92f87 100644 --- a/src/transformers/models/gemma3n/modular_gemma3n.py +++ b/src/transformers/models/gemma3n/modular_gemma3n.py @@ -26,6 +26,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS @@ -106,16 +107,13 @@ 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.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } default_theta = {"global": 1_000_000.0, "local": 10_000.0} diff --git a/src/transformers/models/gemma4/configuration_gemma4.py b/src/transformers/models/gemma4/configuration_gemma4.py index 55c6d97d9ffc..9be69c3c0860 100644 --- a/src/transformers/models/gemma4/configuration_gemma4.py +++ b/src/transformers/models/gemma4/configuration_gemma4.py @@ -17,6 +17,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...utils import auto_docstring, logging from ...utils.type_validators import interval @@ -123,18 +124,18 @@ 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", - "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.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "layers.*.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -226,15 +227,13 @@ class Gemma4VisionConfig(PreTrainedConfig): model_type = "gemma4_vision" base_model_tp_plan = { - "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.*.mlp.gate_proj": "colwise", - "encoder.layers.*.mlp.up_proj": "colwise", - "encoder.layers.*.mlp.down_proj": "rowwise", + "encoder.layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "encoder.layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "encoder.layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "encoder.layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "encoder.layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "encoder.layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "encoder.layers.*.mlp.down_proj": TPStyle("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 406aa0ac72cd..17195d973c9c 100644 --- a/src/transformers/models/gemma4/modeling_gemma4.py +++ b/src/transformers/models/gemma4/modeling_gemma4.py @@ -34,6 +34,7 @@ from ...configuration_utils import PreTrainedConfig from ...generation import GenerationMixin from ...integrations import use_experts_implementation, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import ( create_bidirectional_mask, create_causal_mask, @@ -1699,7 +1700,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} config: Gemma4TextConfig base_model_prefix = "model" diff --git a/src/transformers/models/glm/modeling_glm.py b/src/transformers/models/glm/modeling_glm.py index 712202580943..ddea84722156 100644 --- a/src/transformers/models/glm/modeling_glm.py +++ b/src/transformers/models/glm/modeling_glm.py @@ -28,6 +28,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import ( GenericForSequenceClassification, @@ -439,7 +440,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/glm4/configuration_glm4.py b/src/transformers/models/glm4/configuration_glm4.py index f33129607fec..a91eab20b12c 100644 --- a/src/transformers/models/glm4/configuration_glm4.py +++ b/src/transformers/models/glm4/configuration_glm4.py @@ -17,6 +17,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -40,12 +41,14 @@ class Glm4Config(PreTrainedConfig): model_type = "glm4" 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.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.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_up_proj": TPStyle( + "colwise", "allgather" + ), # we need to replicate here due to the `chunk` operation + "layers.*.mlp.down_proj": TPStyle("vocab", "allreduce"), # input is replicated due to the `chunk` operation } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/glm4/modeling_glm4.py b/src/transformers/models/glm4/modeling_glm4.py index e99930ae57f6..e896daac86bc 100644 --- a/src/transformers/models/glm4/modeling_glm4.py +++ b/src/transformers/models/glm4/modeling_glm4.py @@ -29,6 +29,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -444,7 +445,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, 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..b09c9d61eef9 100644 --- a/src/transformers/models/glm4_moe/configuration_glm4_moe.py +++ b/src/transformers/models/glm4_moe/configuration_glm4_moe.py @@ -20,6 +20,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -54,19 +55,21 @@ class Glm4MoeConfig(PreTrainedConfig): # Default tensor parallel plan for base model `Glm4Moe` 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", # 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.*.mlp.shared_experts.gate_proj": "colwise", - "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), + "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/glm4_moe/modeling_glm4_moe.py b/src/transformers/models/glm4_moe/modeling_glm4_moe.py index 1bc20c8322d9..7275715df0c5 100644 --- a/src/transformers/models/glm4_moe/modeling_glm4_moe.py +++ b/src/transformers/models/glm4_moe/modeling_glm4_moe.py @@ -30,6 +30,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -577,7 +578,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, 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..0a0f6d9da610 100644 --- a/src/transformers/models/glm4_moe/modular_glm4_moe.py +++ b/src/transformers/models/glm4_moe/modular_glm4_moe.py @@ -18,6 +18,7 @@ from torch import nn from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring, logging from ..cohere.modeling_cohere import CohereAttention @@ -67,19 +68,21 @@ class Glm4MoeConfig(PreTrainedConfig): # Default tensor parallel plan for base model `Glm4Moe` 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", # 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.*.mlp.shared_experts.gate_proj": "colwise", - "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), + "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), 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..194eef2a695b 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 @@ -22,6 +22,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -52,16 +53,18 @@ class Glm4MoeLiteConfig(PreTrainedConfig): model_type = "glm4_moe_lite" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_b_proj": "colwise", + "layers.*.self_attn.q_b_proj": TPStyle("colwise", "none"), "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.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.kv_b_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), 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 d59fd2ab996e..b7949d728429 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 @@ -32,6 +32,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernel_func_from_hub +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -651,7 +652,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, 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..e242621817a3 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 @@ -17,6 +17,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring from ..deepseek_v3.modeling_deepseek_v3 import DeepseekV3Attention @@ -60,16 +61,18 @@ class Glm4MoeLiteConfig(PreTrainedConfig): model_type = "glm4_moe_lite" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_b_proj": "colwise", + "layers.*.self_attn.q_b_proj": TPStyle("colwise", "none"), "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.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.kv_b_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/glm4v/configuration_glm4v.py b/src/transformers/models/glm4v/configuration_glm4v.py index 4f151aa38156..e887cd44f6d2 100644 --- a/src/transformers/models/glm4v/configuration_glm4v.py +++ b/src/transformers/models/glm4v/configuration_glm4v.py @@ -20,6 +20,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -90,12 +91,14 @@ class Glm4vTextConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `Glm4v` 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_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.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_up_proj": TPStyle( + "colwise", "allgather" + ), # we need to replicate here due to the `chunk` operation + "layers.*.mlp.down_proj": TPStyle("vocab", "allreduce"), # input is replicated due to the `chunk` operation } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/glm4v/modular_glm4v.py b/src/transformers/models/glm4v/modular_glm4v.py index 1ffd06532a8b..ff53d2e287c6 100644 --- a/src/transformers/models/glm4v/modular_glm4v.py +++ b/src/transformers/models/glm4v/modular_glm4v.py @@ -27,6 +27,7 @@ from ...configuration_utils import PreTrainedConfig from ...feature_extraction_utils import BatchFeature from ...image_utils import ImageInput +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -134,12 +135,14 @@ class Glm4vTextConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `Glm4v` 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_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.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_up_proj": TPStyle( + "colwise", "allgather" + ), # we need to replicate here due to the `chunk` operation + "layers.*.mlp.down_proj": TPStyle("vocab", "allreduce"), # input is replicated due to the `chunk` operation } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py b/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py index 0e4d6a9cb191..c95729daa67e 100644 --- a/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py +++ b/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py @@ -20,6 +20,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -53,13 +54,13 @@ class Glm4vMoeTextConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `Glm4vMoe` 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", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/glm4v_moe/modular_glm4v_moe.py b/src/transformers/models/glm4v_moe/modular_glm4v_moe.py index 0929f3797e22..05c48b004623 100644 --- a/src/transformers/models/glm4v_moe/modular_glm4v_moe.py +++ b/src/transformers/models/glm4v_moe/modular_glm4v_moe.py @@ -19,6 +19,7 @@ from ... import initialization as init from ...cache_utils import Cache, DynamicCache +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_outputs import MoeModelOutputWithPast @@ -85,13 +86,13 @@ class Glm4vMoeTextConfig(Glm4MoeConfig): keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `Glm4vMoe` 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", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/glm_image/configuration_glm_image.py b/src/transformers/models/glm_image/configuration_glm_image.py index fcd302e35560..de3b70617151 100644 --- a/src/transformers/models/glm_image/configuration_glm_image.py +++ b/src/transformers/models/glm_image/configuration_glm_image.py @@ -21,6 +21,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -103,12 +104,14 @@ class GlmImageTextConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `GlmImage` 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_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.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_up_proj": TPStyle( + "colwise", "allgather" + ), # we need to replicate here due to the `chunk` operation + "layers.*.mlp.down_proj": TPStyle("vocab", "allreduce"), # input is replicated due to the `chunk` operation } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), 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 1f9f2a766cf4..68285a8f6d71 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 @@ -21,6 +21,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -57,19 +58,21 @@ class GlmMoeDsaConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_b_proj": "colwise", + "layers.*.self_attn.q_b_proj": TPStyle("colwise", "none"), "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.*.mlp.shared_experts.gate_proj": "colwise", - "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.kv_b_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), + "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), 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 950deba0800e..e88d0f62ed58 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 @@ -30,6 +30,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_experts_implementation, use_kernel_forward_from_hub +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -805,7 +806,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, 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 fcae77cb2562..4789c82aac64 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 @@ -21,6 +21,7 @@ from ...cache_utils import Cache from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_utils import ALL_ATTENTION_FUNCTIONS from ...models.llama.modeling_llama import rotate_half @@ -103,19 +104,21 @@ class GlmMoeDsaConfig(Glm4MoeLiteConfig): ```""" base_model_tp_plan = { - "layers.*.self_attn.q_b_proj": "colwise", + "layers.*.self_attn.q_b_proj": TPStyle("colwise", "none"), "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.*.mlp.shared_experts.gate_proj": "colwise", - "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.kv_b_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), + "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } hidden_size: int = 6144 diff --git a/src/transformers/models/glm_ocr/configuration_glm_ocr.py b/src/transformers/models/glm_ocr/configuration_glm_ocr.py index d08710c5e8b5..fdad0285f965 100644 --- a/src/transformers/models/glm_ocr/configuration_glm_ocr.py +++ b/src/transformers/models/glm_ocr/configuration_glm_ocr.py @@ -21,6 +21,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -91,12 +92,14 @@ class GlmOcrTextConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `GlmOcr` 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_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.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_up_proj": TPStyle( + "colwise", "allgather" + ), # we need to replicate here due to the `chunk` operation + "layers.*.mlp.down_proj": TPStyle("vocab", "allreduce"), # input is replicated due to the `chunk` operation } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/gpt_neox/configuration_gpt_neox.py b/src/transformers/models/gpt_neox/configuration_gpt_neox.py index 256ddca26b49..782bea43357f 100644 --- a/src/transformers/models/gpt_neox/configuration_gpt_neox.py +++ b/src/transformers/models/gpt_neox/configuration_gpt_neox.py @@ -16,6 +16,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -46,10 +47,10 @@ class GPTNeoXConfig(PreTrainedConfig): model_type = "gpt_neox" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.attention.query_key_value": "colwise", - "layers.*.attention.dense": "rowwise", - "layers.*.mlp.dense_h_to_4h": "colwise", - "layers.*.mlp.dense_4h_to_h": "rowwise", + "layers.*.attention.query_key_value": TPStyle("colwise", "none"), + "layers.*.attention.dense": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.dense_h_to_4h": TPStyle("colwise", "none"), + "layers.*.mlp.dense_4h_to_h": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_in": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/gpt_neox/modeling_gpt_neox.py b/src/transformers/models/gpt_neox/modeling_gpt_neox.py index 10e4b5922add..7e4426f258c1 100755 --- a/src/transformers/models/gpt_neox/modeling_gpt_neox.py +++ b/src/transformers/models/gpt_neox/modeling_gpt_neox.py @@ -13,6 +13,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -394,7 +395,7 @@ def set_input_embeddings(self, value): ) class GPTNeoXForCausalLM(GPTNeoXPreTrainedModel, GenerationMixin): _tied_weights_keys = {"embed_out.weight": "gpt_neox.embed_in.weight"} - _tp_plan = {"embed_out": "colwise_gather_output"} + _tp_plan = {"embed_out": TPStyle("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..5833150550e2 100644 --- a/src/transformers/models/gpt_neox/modular_gpt_neox.py +++ b/src/transformers/models/gpt_neox/modular_gpt_neox.py @@ -7,6 +7,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -341,7 +342,7 @@ def forward( ) class GPTNeoXForCausalLM(GPTNeoXPreTrainedModel, GenerationMixin): _tied_weights_keys = {"embed_out.weight": "gpt_neox.embed_in.weight"} - _tp_plan = {"embed_out": "colwise_gather_output"} + _tp_plan = {"embed_out": TPStyle("colwise", "allgather")} _pp_plan = {"embed_out": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/gpt_oss/modeling_gpt_oss.py b/src/transformers/models/gpt_oss/modeling_gpt_oss.py index 18f31ea90379..02527fb44c44 100644 --- a/src/transformers/models/gpt_oss/modeling_gpt_oss.py +++ b/src/transformers/models/gpt_oss/modeling_gpt_oss.py @@ -28,6 +28,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_layers import ( GenericForSequenceClassification, @@ -588,7 +589,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/granite/configuration_granite.py b/src/transformers/models/granite/configuration_granite.py index e026cbbe5ff3..6696c4d33685 100644 --- a/src/transformers/models/granite/configuration_granite.py +++ b/src/transformers/models/granite/configuration_granite.py @@ -21,6 +21,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -47,13 +48,13 @@ class GraniteConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `GraniteModel` 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", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/granite/modeling_granite.py b/src/transformers/models/granite/modeling_granite.py index 934345fe6723..b6062282fe8a 100644 --- a/src/transformers/models/granite/modeling_granite.py +++ b/src/transformers/models/granite/modeling_granite.py @@ -29,6 +29,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -445,7 +446,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/granitemoe/modeling_granitemoe.py b/src/transformers/models/granitemoe/modeling_granitemoe.py index 5fb53d6afe49..ee63a27d4e11 100644 --- a/src/transformers/models/granitemoe/modeling_granitemoe.py +++ b/src/transformers/models/granitemoe/modeling_granitemoe.py @@ -31,6 +31,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast @@ -626,7 +627,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: GraniteMoeConfig): diff --git a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py index dadffaea0072..7441fc47997c 100644 --- a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py +++ b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py @@ -31,6 +31,7 @@ from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func from ...integrations.hub_kernels import lazy_load_kernel +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, MoeCausalLMOutputWithPast, MoeModelOutputWithPast @@ -1306,7 +1307,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: GraniteMoeHybridConfig): diff --git a/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py b/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py index 71f8c6eaff7d..a698eae304ac 100644 --- a/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py +++ b/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py @@ -30,6 +30,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast @@ -695,7 +696,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: GraniteMoeSharedConfig): diff --git a/src/transformers/models/helium/configuration_helium.py b/src/transformers/models/helium/configuration_helium.py index caa966b58a9d..8e2f44ffca57 100644 --- a/src/transformers/models/helium/configuration_helium.py +++ b/src/transformers/models/helium/configuration_helium.py @@ -17,6 +17,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -41,13 +42,13 @@ class HeliumConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] default_theta = 100000.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.o_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/helium/modeling_helium.py b/src/transformers/models/helium/modeling_helium.py index 8283fcb19e28..e3a8de1b8b28 100644 --- a/src/transformers/models/helium/modeling_helium.py +++ b/src/transformers/models/helium/modeling_helium.py @@ -29,6 +29,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import ( GenericForSequenceClassification, @@ -423,7 +424,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): 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..481a644d891b 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 @@ -31,6 +31,7 @@ from ...cache_utils import DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GenericForSequenceClassification, GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -461,7 +462,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, 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..75f4033f69f5 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 @@ -35,6 +35,7 @@ use_kernel_func_from_hub, use_kernelized_func, ) +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GenericForSequenceClassification, GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -550,7 +551,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/jais2/configuration_jais2.py b/src/transformers/models/jais2/configuration_jais2.py index 7139661d7575..4886c5c45acf 100644 --- a/src/transformers/models/jais2/configuration_jais2.py +++ b/src/transformers/models/jais2/configuration_jais2.py @@ -22,6 +22,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring from ...utils.type_validators import interval @@ -47,12 +48,12 @@ class Jais2Config(PreTrainedConfig): model_type = "jais2" 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.o_proj": "rowwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/jais2/modeling_jais2.py b/src/transformers/models/jais2/modeling_jais2.py index 5e6a37c0172d..ca1de79c6fbb 100644 --- a/src/transformers/models/jais2/modeling_jais2.py +++ b/src/transformers/models/jais2/modeling_jais2.py @@ -29,6 +29,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_func_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -397,7 +398,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/jais2/modular_jais2.py b/src/transformers/models/jais2/modular_jais2.py index 3fbdf2c8cd46..90279d121eaa 100644 --- a/src/transformers/models/jais2/modular_jais2.py +++ b/src/transformers/models/jais2/modular_jais2.py @@ -16,6 +16,7 @@ import torch.nn as nn from huggingface_hub.dataclasses import strict +from ...integrations.tensor_parallel import TPStyle from ...utils import auto_docstring, can_return_tuple from ..llama.configuration_llama import LlamaConfig from ..llama.modeling_llama import ( @@ -31,12 +32,12 @@ @strict class Jais2Config(LlamaConfig): 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.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } vocab_size: int = 150272 diff --git a/src/transformers/models/jamba/modeling_jamba.py b/src/transformers/models/jamba/modeling_jamba.py index ae618fb4a2b3..85c00fe0f600 100755 --- a/src/transformers/models/jamba/modeling_jamba.py +++ b/src/transformers/models/jamba/modeling_jamba.py @@ -38,6 +38,7 @@ use_kernel_func_from_hub, use_kernelized_func, ) +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GenericForSequenceClassification, GradientCheckpointingLayer from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast @@ -843,7 +844,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: JambaConfig): 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..8ddd75155b56 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 @@ -30,6 +30,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache, StaticCache from ...generation import GenerationConfig, GenerationMixin +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import flash_attn_supports_top_left_mask, is_flash_attn_available from ...modeling_layers import GradientCheckpointingLayer @@ -874,7 +875,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _keep_in_fp32_modules_strict = ["codec_model"] output_modalities = ("audio", "text") diff --git a/src/transformers/models/lfm2/modeling_lfm2.py b/src/transformers/models/lfm2/modeling_lfm2.py index ef753e3b2893..8071817aa46c 100644 --- a/src/transformers/models/lfm2/modeling_lfm2.py +++ b/src/transformers/models/lfm2/modeling_lfm2.py @@ -27,6 +27,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -540,7 +541,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py b/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py index 0369ae31b8ae..4a1935876cca 100644 --- a/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py +++ b/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py @@ -34,6 +34,7 @@ use_kernel_func_from_hub, use_kernelized_func, ) +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, MoeModelOutputWithPast @@ -630,7 +631,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/llama/modeling_llama.py b/src/transformers/models/llama/modeling_llama.py index 9d659c7c6f08..0c1aca1edff8 100644 --- a/src/transformers/models/llama/modeling_llama.py +++ b/src/transformers/models/llama/modeling_llama.py @@ -26,6 +26,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import ( GenericForQuestionAnswering, @@ -428,7 +429,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/llama4/configuration_llama4.py b/src/transformers/models/llama4/configuration_llama4.py index 79cfd063f4d4..1d8c350d4b20 100644 --- a/src/transformers/models/llama4/configuration_llama4.py +++ b/src/transformers/models/llama4/configuration_llama4.py @@ -17,6 +17,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring, logging @@ -42,13 +43,13 @@ class Llama4VisionConfig(PreTrainedConfig): """ base_model_tp_plan = { - "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", - "vision_adapter.mlp.fc1": "colwise", - "vision_adapter.mlp.fc2": "rowwise", - "patch_embedding.linear": "colwise_gather_output", + "model.layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "model.layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "model.layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "model.layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "vision_adapter.mlp.fc1": TPStyle("colwise", "none"), + "vision_adapter.mlp.fc2": TPStyle("rowwise", "allreduce"), + "patch_embedding.linear": TPStyle("colwise", "allgather"), } model_type = "llama4_vision_model" base_config_key = "vision_config" @@ -110,18 +111,16 @@ class Llama4TextConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] default_theta = 500000.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.o_proj": "rowwise", - "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.gate_proj": "colwise", - "layers.*.feed_forward.up_proj": "colwise", - "layers.*.feed_forward.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.feed_forward.shared_expert.gate_proj": TPStyle("colwise", "none"), + "layers.*.feed_forward.shared_expert.up_proj": TPStyle("colwise", "none"), + "layers.*.feed_forward.shared_expert.down_proj": TPStyle("rowwise", "allreduce"), + "layers.*.feed_forward.gate_proj": TPStyle("colwise", "none"), + "layers.*.feed_forward.up_proj": TPStyle("colwise", "none"), + "layers.*.feed_forward.down_proj": TPStyle("rowwise", "allreduce"), } base_model_ep_plan = { "layers.*.self_attn.q_proj": "colwise", @@ -179,7 +178,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 08d50bd63f72..660dddd91bff 100644 --- a/src/transformers/models/llama4/modeling_llama4.py +++ b/src/transformers/models/llama4/modeling_llama4.py @@ -28,6 +28,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_chunked_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -590,7 +591,7 @@ 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": TPStyle("colwise", "allgather")} 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..5547c24deeb5 100644 --- a/src/transformers/models/longcat_flash/configuration_longcat_flash.py +++ b/src/transformers/models/longcat_flash/configuration_longcat_flash.py @@ -17,6 +17,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -54,17 +55,19 @@ class LongcatFlashConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] default_theta = 10000000.0 base_model_tp_plan = { - "layers.*.self_attn.*.q_b_proj": "colwise", + "layers.*.self_attn.*.q_b_proj": TPStyle("colwise", "none"), "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.*.kv_b_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.*.o_proj": TPStyle("rowwise", "allreduce"), "layers.*.mlp.experts.identity_expert": "moe_identity_expert", - "layers.*.mlp.experts": "moe_tp_experts", - "layers.*.mlps.*.gate_proj": "colwise", - "layers.*.mlps.*.up_proj": "colwise", - "layers.*.mlps.*.down_proj": "rowwise", + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), + "layers.*.mlps.*.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlps.*.up_proj": TPStyle("colwise", "none"), + "layers.*.mlps.*.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { diff --git a/src/transformers/models/longcat_flash/modeling_longcat_flash.py b/src/transformers/models/longcat_flash/modeling_longcat_flash.py index d5ac6e237742..a1675db42374 100644 --- a/src/transformers/models/longcat_flash/modeling_longcat_flash.py +++ b/src/transformers/models/longcat_flash/modeling_longcat_flash.py @@ -31,6 +31,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -650,7 +651,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _keys_to_ignore_on_load_unexpected = [r"model\.mtp.*"] diff --git a/src/transformers/models/minimax/configuration_minimax.py b/src/transformers/models/minimax/configuration_minimax.py index 9a3a0023725e..699a768bc311 100644 --- a/src/transformers/models/minimax/configuration_minimax.py +++ b/src/transformers/models/minimax/configuration_minimax.py @@ -21,6 +21,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -62,13 +63,15 @@ class MiniMaxConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] default_theta = 1000000.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.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_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/minimax/modeling_minimax.py b/src/transformers/models/minimax/modeling_minimax.py index 69497f83cad8..fde19bce52dd 100644 --- a/src/transformers/models/minimax/modeling_minimax.py +++ b/src/transformers/models/minimax/modeling_minimax.py @@ -36,6 +36,7 @@ use_kernel_func_from_hub, use_kernelized_func, ) +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -789,7 +790,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/minimax/modular_minimax.py b/src/transformers/models/minimax/modular_minimax.py index 0bd400458129..1dd83f5d148a 100644 --- a/src/transformers/models/minimax/modular_minimax.py +++ b/src/transformers/models/minimax/modular_minimax.py @@ -23,6 +23,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -88,13 +89,15 @@ class MiniMaxConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] default_theta = 1000000.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.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_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/minimax_m2/configuration_minimax_m2.py b/src/transformers/models/minimax_m2/configuration_minimax_m2.py index 75acb8b755d7..0ab84d885690 100644 --- a/src/transformers/models/minimax_m2/configuration_minimax_m2.py +++ b/src/transformers/models/minimax_m2/configuration_minimax_m2.py @@ -22,6 +22,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -48,13 +49,15 @@ 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": TPStyle("colwise", "allgather"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "allgather"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "allgather"), + "layers.*.self_attn.o_proj": TPStyle("vocab", "allreduce"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/minimax_m2/modeling_minimax_m2.py b/src/transformers/models/minimax_m2/modeling_minimax_m2.py index d19274262810..8597f43f40f7 100644 --- a/src/transformers/models/minimax_m2/modeling_minimax_m2.py +++ b/src/transformers/models/minimax_m2/modeling_minimax_m2.py @@ -31,6 +31,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast @@ -588,7 +589,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/minimax_m2/modular_minimax_m2.py b/src/transformers/models/minimax_m2/modular_minimax_m2.py index a9938a555c62..b0994cf4cd89 100644 --- a/src/transformers/models/minimax_m2/modular_minimax_m2.py +++ b/src/transformers/models/minimax_m2/modular_minimax_m2.py @@ -21,6 +21,7 @@ from ... import initialization as init from ...cache_utils import Cache, DynamicCache from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_outputs import MoeModelOutputWithPast from ...modeling_rope_utils import RopeParameters @@ -67,13 +68,15 @@ 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": TPStyle("colwise", "allgather"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "allgather"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "allgather"), + "layers.*.self_attn.o_proj": TPStyle("vocab", "allreduce"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/ministral/modeling_ministral.py b/src/transformers/models/ministral/modeling_ministral.py index af4f7fbeae59..45d3897db33c 100644 --- a/src/transformers/models/ministral/modeling_ministral.py +++ b/src/transformers/models/ministral/modeling_ministral.py @@ -29,6 +29,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -429,7 +430,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/ministral3/configuration_ministral3.py b/src/transformers/models/ministral3/configuration_ministral3.py index 76d55710d5ae..82cefa62a805 100644 --- a/src/transformers/models/ministral3/configuration_ministral3.py +++ b/src/transformers/models/ministral3/configuration_ministral3.py @@ -16,6 +16,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring, logging @@ -55,13 +56,13 @@ class Ministral3Config(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `MistralModel` 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", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/ministral3/modeling_ministral3.py b/src/transformers/models/ministral3/modeling_ministral3.py index 6aacf4c8ce3a..e158c819586e 100644 --- a/src/transformers/models/ministral3/modeling_ministral3.py +++ b/src/transformers/models/ministral3/modeling_ministral3.py @@ -14,6 +14,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -412,7 +413,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/mistral/modeling_mistral.py b/src/transformers/models/mistral/modeling_mistral.py index b79dea36c9e9..523fe2f34c98 100644 --- a/src/transformers/models/mistral/modeling_mistral.py +++ b/src/transformers/models/mistral/modeling_mistral.py @@ -14,6 +14,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -401,7 +402,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/mistral4/configuration_mistral4.py b/src/transformers/models/mistral4/configuration_mistral4.py index 0e16e0a14f45..74fed8abab2a 100644 --- a/src/transformers/models/mistral4/configuration_mistral4.py +++ b/src/transformers/models/mistral4/configuration_mistral4.py @@ -16,6 +16,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -47,15 +48,17 @@ 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.shared_experts.gate_proj": "colwise", - "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), + "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/mistral4/modeling_mistral4.py b/src/transformers/models/mistral4/modeling_mistral4.py index 006ddad187bf..15f4e9f43e98 100644 --- a/src/transformers/models/mistral4/modeling_mistral4.py +++ b/src/transformers/models/mistral4/modeling_mistral4.py @@ -29,6 +29,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernel_func_from_hub +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -640,7 +641,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/mixtral/configuration_mixtral.py b/src/transformers/models/mixtral/configuration_mixtral.py index 7b6ab7aaa974..9e4e22ef770b 100644 --- a/src/transformers/models/mixtral/configuration_mixtral.py +++ b/src/transformers/models/mixtral/configuration_mixtral.py @@ -53,8 +53,8 @@ class MixtralConfig(PreTrainedConfig): "moe_experts", "allreduce", shard_plan={ - "gate_up_proj": "packed_colwise", - "down_proj": "rowwise", + "gate_up_proj": TPStyle("packed_colwise", "none"), + "down_proj": TPStyle("rowwise", "allreduce"), }, ), } @@ -74,8 +74,8 @@ class MixtralConfig(PreTrainedConfig): "moe_experts", "allreduce", shard_plan={ - "gate_up_proj": "packed_colwise", - "down_proj": "rowwise", + "gate_up_proj": TPStyle("packed_colwise", "none"), + "down_proj": TPStyle("rowwise", "allreduce"), }, ), "norm": TPStyle("activation", "none"), diff --git a/src/transformers/models/mixtral/modeling_mixtral.py b/src/transformers/models/mixtral/modeling_mixtral.py index 991851dbadd3..1965f7ee8da4 100644 --- a/src/transformers/models/mixtral/modeling_mixtral.py +++ b/src/transformers/models/mixtral/modeling_mixtral.py @@ -40,6 +40,7 @@ use_kernel_func_from_hub, use_kernelized_func, ) +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -580,7 +581,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/nanochat/configuration_nanochat.py b/src/transformers/models/nanochat/configuration_nanochat.py index bd915e06dab4..0752837e4753 100644 --- a/src/transformers/models/nanochat/configuration_nanochat.py +++ b/src/transformers/models/nanochat/configuration_nanochat.py @@ -16,6 +16,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PretrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -43,12 +44,12 @@ class NanoChatConfig(PretrainedConfig): 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.o_proj": "rowwise", - "layers.*.mlp.fc1": "colwise", - "layers.*.mlp.fc2": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.fc1": TPStyle("colwise", "none"), + "layers.*.mlp.fc2": TPStyle("rowwise", "allreduce"), } vocab_size: int = 50304 diff --git a/src/transformers/models/nanochat/modeling_nanochat.py b/src/transformers/models/nanochat/modeling_nanochat.py index 9205b89cd360..b14faff93e67 100644 --- a/src/transformers/models/nanochat/modeling_nanochat.py +++ b/src/transformers/models/nanochat/modeling_nanochat.py @@ -30,6 +30,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_func_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -122,32 +123,6 @@ def forward(self, x, position_ids): return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) -@use_kernel_func_from_hub("rotary_pos_emb") -def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): - """Applies Rotary Position Embedding to the query and key tensors. - - Args: - q (`torch.Tensor`): The query tensor. - k (`torch.Tensor`): The key tensor. - cos (`torch.Tensor`): The cosine part of the rotary embedding. - sin (`torch.Tensor`): The sine part of the rotary embedding. - unsqueeze_dim (`int`, *optional*, defaults to 1): - The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and - sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note - that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and - k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes - cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have - the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. - Returns: - `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. - """ - cos = cos.unsqueeze(unsqueeze_dim) - sin = sin.unsqueeze(unsqueeze_dim) - q_embed = (q * cos) + (rotate_half(q) * sin) - k_embed = (k * cos) + (rotate_half(k) * sin) - return q_embed, k_embed - - def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, @@ -185,6 +160,32 @@ def eager_attention_forward( return attn_output, attn_weights +@use_kernel_func_from_hub("rotary_pos_emb") +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + def rotate_half(x): """Rotates half the hidden dims of the input with flipped signs for NanoChat.""" x1 = x[..., : x.shape[-1] // 2] @@ -432,7 +433,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/nanochat/modular_nanochat.py b/src/transformers/models/nanochat/modular_nanochat.py index 713cc29b81eb..486ec255089e 100644 --- a/src/transformers/models/nanochat/modular_nanochat.py +++ b/src/transformers/models/nanochat/modular_nanochat.py @@ -20,6 +20,7 @@ from ... import initialization as init from ...cache_utils import Cache, DynamicCache +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel @@ -198,7 +199,7 @@ def forward( @auto_docstring class NanoChatForCausalLM(Gemma2ForCausalLM): - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} def forward(self, **super_kwargs) -> CausalLMOutputWithPast: r""" diff --git a/src/transformers/models/olmo/configuration_olmo.py b/src/transformers/models/olmo/configuration_olmo.py index 186cc3a704fb..1e626e42d98e 100644 --- a/src/transformers/models/olmo/configuration_olmo.py +++ b/src/transformers/models/olmo/configuration_olmo.py @@ -21,6 +21,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -50,13 +51,13 @@ class OlmoConfig(PreTrainedConfig): model_type = "olmo" 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.o_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/olmo/modeling_olmo.py b/src/transformers/models/olmo/modeling_olmo.py index a0949e957057..47a7e2d78edf 100644 --- a/src/transformers/models/olmo/modeling_olmo.py +++ b/src/transformers/models/olmo/modeling_olmo.py @@ -34,6 +34,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -425,7 +426,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/olmo2/configuration_olmo2.py b/src/transformers/models/olmo2/configuration_olmo2.py index f879c0b8367f..96dbdd2908fb 100644 --- a/src/transformers/models/olmo2/configuration_olmo2.py +++ b/src/transformers/models/olmo2/configuration_olmo2.py @@ -26,6 +26,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -53,13 +54,21 @@ 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.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle( + "colwise", "allgather" + ), # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.k_proj": TPStyle( + "colwise", "allgather" + ), # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.v_proj": TPStyle( + "colwise", "allgather" + ), # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.o_proj": TPStyle( + "vocab", "allreduce" + ), # input is replicated due to the added norm on q and k + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/olmo2/modeling_olmo2.py b/src/transformers/models/olmo2/modeling_olmo2.py index 01fe9fb158b8..502d800ce288 100644 --- a/src/transformers/models/olmo2/modeling_olmo2.py +++ b/src/transformers/models/olmo2/modeling_olmo2.py @@ -35,6 +35,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -429,7 +430,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/olmo2/modular_olmo2.py b/src/transformers/models/olmo2/modular_olmo2.py index c2c0d0ecaa30..94ca63e2a7c0 100644 --- a/src/transformers/models/olmo2/modular_olmo2.py +++ b/src/transformers/models/olmo2/modular_olmo2.py @@ -26,6 +26,7 @@ from transformers.utils.generic import TransformersKwargs from ...cache_utils import Cache +from ...integrations.tensor_parallel import TPStyle from ...modeling_utils import ALL_ATTENTION_FUNCTIONS from ...processing_utils import Unpack from ...utils import auto_docstring, logging @@ -66,13 +67,21 @@ 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.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle( + "colwise", "allgather" + ), # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.k_proj": TPStyle( + "colwise", "allgather" + ), # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.v_proj": TPStyle( + "colwise", "allgather" + ), # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.o_proj": TPStyle( + "vocab", "allreduce" + ), # input is replicated due to the added norm on q and k + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } 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..fb7f445cf079 100644 --- a/src/transformers/models/olmo3/configuration_olmo3.py +++ b/src/transformers/models/olmo3/configuration_olmo3.py @@ -21,6 +21,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -48,13 +49,21 @@ 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.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle( + "colwise", "allgather" + ), # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.k_proj": TPStyle( + "colwise", "allgather" + ), # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.v_proj": TPStyle( + "colwise", "allgather" + ), # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.o_proj": TPStyle( + "vocab", "allreduce" + ), # input is replicated due to the added norm on q and k + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/olmo3/modeling_olmo3.py b/src/transformers/models/olmo3/modeling_olmo3.py index 5baa8e5f24ed..04dbdb27b1c2 100644 --- a/src/transformers/models/olmo3/modeling_olmo3.py +++ b/src/transformers/models/olmo3/modeling_olmo3.py @@ -28,6 +28,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -433,7 +434,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/olmo3/modular_olmo3.py b/src/transformers/models/olmo3/modular_olmo3.py index 81325db7cc1f..fba53df7502c 100644 --- a/src/transformers/models/olmo3/modular_olmo3.py +++ b/src/transformers/models/olmo3/modular_olmo3.py @@ -19,6 +19,7 @@ from huggingface_hub.dataclasses import strict from ...cache_utils import Cache, DynamicCache +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_outputs import BaseModelOutputWithPast from ...modeling_utils import ALL_ATTENTION_FUNCTIONS @@ -62,13 +63,21 @@ 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.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle( + "colwise", "allgather" + ), # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.k_proj": TPStyle( + "colwise", "allgather" + ), # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.v_proj": TPStyle( + "colwise", "allgather" + ), # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.o_proj": TPStyle( + "vocab", "allreduce" + ), # input is replicated due to the added norm on q and k + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } 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..65e9c2043045 100644 --- a/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py +++ b/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py @@ -21,6 +21,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring from ...utils.type_validators import interval @@ -73,13 +74,21 @@ 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.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle( + "colwise", "allgather" + ), # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.k_proj": TPStyle( + "colwise", "allgather" + ), # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.v_proj": TPStyle( + "colwise", "allgather" + ), # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.o_proj": TPStyle( + "vocab", "allreduce" + ), # input is replicated due to the added norm on q and k + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py index 8ecb44916a57..092cfb94e6c5 100644 --- a/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py +++ b/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py @@ -32,6 +32,7 @@ from ...cache_utils import Cache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -1032,7 +1033,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py index 089f29309007..2518a5bc247c 100644 --- a/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py +++ b/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py @@ -27,6 +27,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_outputs import BaseModelOutputWithPast from ...modeling_rope_utils import dynamic_rope_update @@ -120,13 +121,21 @@ 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.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle( + "colwise", "allgather" + ), # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.k_proj": TPStyle( + "colwise", "allgather" + ), # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.v_proj": TPStyle( + "colwise", "allgather" + ), # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.o_proj": TPStyle( + "vocab", "allreduce" + ), # input is replicated due to the added norm on q and k + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } vocab_size: int = 100352 diff --git a/src/transformers/models/olmoe/configuration_olmoe.py b/src/transformers/models/olmoe/configuration_olmoe.py index 16bedbe698f8..17def985f001 100644 --- a/src/transformers/models/olmoe/configuration_olmoe.py +++ b/src/transformers/models/olmoe/configuration_olmoe.py @@ -14,6 +14,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -46,13 +47,15 @@ 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": TPStyle("colwise", "allgather"), # due to the norm, we have to gather + "layers.*.self_attn.k_proj": TPStyle("colwise", "allgather"), # due to the norm, we have to gather + "layers.*.self_attn.v_proj": TPStyle("colwise", "allgather"), # due to the norm, we have to gather + "layers.*.self_attn.o_proj": TPStyle("vocab", "allreduce"), # due to the norm, we have to gather + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), } vocab_size: int = 50304 diff --git a/src/transformers/models/olmoe/modeling_olmoe.py b/src/transformers/models/olmoe/modeling_olmoe.py index 5d89ec741529..6d648ff9b879 100644 --- a/src/transformers/models/olmoe/modeling_olmoe.py +++ b/src/transformers/models/olmoe/modeling_olmoe.py @@ -33,6 +33,7 @@ use_kernel_func_from_hub, use_kernelized_func, ) +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast @@ -604,7 +605,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py b/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py index 343a22ade814..4da2e361d0aa 100644 --- a/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py +++ b/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py @@ -28,6 +28,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -95,13 +96,13 @@ class PaddleOCRTextConfig(PreTrainedConfig): default_theta = 500000.0 # Default tensor parallel plan for base model `PaddleOCRTextModel` 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", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/phi/modeling_phi.py b/src/transformers/models/phi/modeling_phi.py index e3f97a01ee4c..700fc1ddfd38 100644 --- a/src/transformers/models/phi/modeling_phi.py +++ b/src/transformers/models/phi/modeling_phi.py @@ -14,6 +14,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_func_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import ( GenericForSequenceClassification, @@ -406,7 +407,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/phi3/configuration_phi3.py b/src/transformers/models/phi3/configuration_phi3.py index f85502f8205d..8e5f50697762 100644 --- a/src/transformers/models/phi3/configuration_phi3.py +++ b/src/transformers/models/phi3/configuration_phi3.py @@ -17,6 +17,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -47,10 +48,14 @@ 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": TPStyle( + "colwise", "allgather" + ), # we need to replicate here due to the slicing of qkv + "layers.*.self_attn.o_proj": TPStyle("vocab", "allreduce"), # input is replicated due to the slicing of qkv + "layers.*.mlp.gate_up_proj": TPStyle( + "colwise", "allgather" + ), # we need to replicate here due to the `chunk` operation + "layers.*.mlp.down_proj": TPStyle("vocab", "allreduce"), # input is replicated due to the `chunk` operation } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/phi3/modeling_phi3.py b/src/transformers/models/phi3/modeling_phi3.py index b07735f8a2e6..e537449ffd56 100644 --- a/src/transformers/models/phi3/modeling_phi3.py +++ b/src/transformers/models/phi3/modeling_phi3.py @@ -29,6 +29,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -432,7 +433,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, 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..62e0925d0e3e 100644 --- a/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py +++ b/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py @@ -23,6 +23,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -181,10 +182,14 @@ 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": TPStyle( + "colwise", "allgather" + ), # we need to replicate here due to the slicing of qkv + "layers.*.self_attn.o_proj": TPStyle("vocab", "allreduce"), # input is replicated due to the slicing of qkv + "layers.*.mlp.gate_up_proj": TPStyle( + "colwise", "allgather" + ), # we need to replicate here due to the `chunk` operation + "layers.*.mlp.down_proj": TPStyle("vocab", "allreduce"), # input is replicated due to the `chunk` operation } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py b/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py index 9e6c0339098d..22d2886f6e17 100644 --- a/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py +++ b/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py @@ -32,6 +32,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_bidirectional_mask, create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -1596,7 +1597,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/phimoe/modeling_phimoe.py b/src/transformers/models/phimoe/modeling_phimoe.py index 23bc944c522a..2846fd2aefb0 100644 --- a/src/transformers/models/phimoe/modeling_phimoe.py +++ b/src/transformers/models/phimoe/modeling_phimoe.py @@ -30,6 +30,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_experts_implementation, use_kernel_func_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_layers import GenericForSequenceClassification, GradientCheckpointingLayer from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast @@ -772,7 +773,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/pi0/modeling_pi0.py b/src/transformers/models/pi0/modeling_pi0.py index 8fd8abe48d7b..a816d633bd04 100644 --- a/src/transformers/models/pi0/modeling_pi0.py +++ b/src/transformers/models/pi0/modeling_pi0.py @@ -27,6 +27,7 @@ from ... import initialization as init from ...cache_utils import Cache +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_bidirectional_mask from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from ...modeling_utils import PreTrainedModel @@ -226,7 +227,7 @@ def forward( class PI0ForConditionalGeneration(PI0PreTrainedModel): """PI0 model with action projection heads and flow matching.""" - _tp_plan = {"action_out_proj": "colwise_gather_output"} + _tp_plan = {"action_out_proj": TPStyle("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 f79ac3c2775a..dfa25e9fc121 100644 --- a/src/transformers/models/pi0/modular_pi0.py +++ b/src/transformers/models/pi0/modular_pi0.py @@ -27,6 +27,7 @@ from ...configuration_utils import PreTrainedConfig from ...feature_extraction_utils import BatchFeature from ...image_utils import ImageInput, make_nested_list_of_images +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_bidirectional_mask from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from ...modeling_utils import PreTrainedModel @@ -476,7 +477,7 @@ def forward( class PI0ForConditionalGeneration(PI0PreTrainedModel): """PI0 model with action projection heads and flow matching.""" - _tp_plan = {"action_out_proj": "colwise_gather_output"} + _tp_plan = {"action_out_proj": TPStyle("colwise", "allgather")} def __init__(self, config: PI0Config): super().__init__(config) diff --git a/src/transformers/models/qwen2/modeling_qwen2.py b/src/transformers/models/qwen2/modeling_qwen2.py index 9263e1d42937..d4e8d1f59d8c 100644 --- a/src/transformers/models/qwen2/modeling_qwen2.py +++ b/src/transformers/models/qwen2/modeling_qwen2.py @@ -14,6 +14,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -416,7 +417,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, 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 c950e15664ec..2f8b406315ae 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 @@ -21,6 +21,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring, logging @@ -149,13 +150,13 @@ class Qwen2_5OmniTextConfig(PreTrainedConfig): # Default tensor parallel plan for base model `Qwen25OmniText` 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", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), 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 6dd4e5727fc6..75ba364d0656 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 @@ -30,6 +30,7 @@ from ...cache_utils import Cache from ...configuration_utils import PreTrainedConfig from ...generation import GenerationMixin +from ...integrations.tensor_parallel import TPStyle from ...modeling_outputs import BaseModelOutputWithPooling, ModelOutput from ...modeling_rope_utils import RopeParameters from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel @@ -161,13 +162,13 @@ class Qwen2_5OmniTextConfig(PreTrainedConfig): # Default tensor parallel plan for base model `Qwen25OmniText` 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", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), 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 911a5543ba48..4b21dda3712a 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 @@ -27,6 +27,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -88,13 +89,13 @@ class Qwen2_5_VLTextConfig(PreTrainedConfig): default_theta = 1000000.0 # Default tensor parallel plan for base model `Qwen2_5_VL` 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", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py b/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py index 7f961976b58c..c0f5a7452fe8 100644 --- a/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py +++ b/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py @@ -16,6 +16,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -54,13 +55,13 @@ class Qwen2MoeConfig(PreTrainedConfig): # Default tensor parallel plan for base model `Qwen2Moe` 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", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py b/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py index d4150d0a74d7..542dddd48fe4 100644 --- a/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py +++ b/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py @@ -40,6 +40,7 @@ use_kernel_func_from_hub, use_kernelized_func, ) +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_layers import ( GenericForQuestionAnswering, @@ -617,7 +618,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, 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..eb624179143b 100644 --- a/src/transformers/models/qwen2_moe/modular_qwen2_moe.py +++ b/src/transformers/models/qwen2_moe/modular_qwen2_moe.py @@ -25,6 +25,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_layers import ( GenericForQuestionAnswering, @@ -230,7 +231,7 @@ 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": TPStyle("colwise", "allgather")} _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 536bca3be654..574f43b541df 100644 --- a/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py +++ b/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py @@ -18,6 +18,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -65,13 +66,13 @@ class Qwen2VLTextConfig(PreTrainedConfig): default_theta = 1000000.0 # Default tensor parallel plan for base model `Qwen2VL` 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", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/qwen3/modeling_qwen3.py b/src/transformers/models/qwen3/modeling_qwen3.py index 91715a33cf9d..beeab1982123 100644 --- a/src/transformers/models/qwen3/modeling_qwen3.py +++ b/src/transformers/models/qwen3/modeling_qwen3.py @@ -27,7 +27,8 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func +from ...integrations import use_kernel_forward_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -148,39 +149,6 @@ def forward(self, x, position_ids): return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) -def rotate_half(x): - """Rotates half the hidden dims of the input.""" - x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] - return torch.cat((-x2, x1), dim=-1) - - -@use_kernel_func_from_hub("rotary_pos_emb") -def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): - """Applies Rotary Position Embedding to the query and key tensors. - - Args: - q (`torch.Tensor`): The query tensor. - k (`torch.Tensor`): The key tensor. - cos (`torch.Tensor`): The cosine part of the rotary embedding. - sin (`torch.Tensor`): The sine part of the rotary embedding. - unsqueeze_dim (`int`, *optional*, defaults to 1): - The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and - sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note - that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and - k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes - cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have - the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. - Returns: - `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. - """ - cos = cos.unsqueeze(unsqueeze_dim) - sin = sin.unsqueeze(unsqueeze_dim) - q_embed = (q * cos) + (rotate_half(q) * sin) - k_embed = (k * cos) + (rotate_half(k) * sin) - return q_embed, k_embed - - def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, @@ -218,6 +186,20 @@ def eager_attention_forward( return attn_output, attn_weights +def rotate_half(x): + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + @use_kernelized_func(apply_rotary_pos_emb) class Qwen3Attention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" @@ -441,7 +423,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/qwen3_5/configuration_qwen3_5.py b/src/transformers/models/qwen3_5/configuration_qwen3_5.py index b200b920b18e..2f2f6ecca1b9 100644 --- a/src/transformers/models/qwen3_5/configuration_qwen3_5.py +++ b/src/transformers/models/qwen3_5/configuration_qwen3_5.py @@ -20,6 +20,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -57,15 +58,13 @@ class Qwen3_5TextConfig(PreTrainedConfig): 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.o_proj": "rowwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/qwen3_5/modeling_qwen3_5.py b/src/transformers/models/qwen3_5/modeling_qwen3_5.py index 4dd3dfbaaf60..fd4f46095111 100644 --- a/src/transformers/models/qwen3_5/modeling_qwen3_5.py +++ b/src/transformers/models/qwen3_5/modeling_qwen3_5.py @@ -32,6 +32,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GenericForSequenceClassification, GradientCheckpointingLayer @@ -533,6 +534,43 @@ def forward( return output +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] @@ -579,43 +617,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): return q_embed, k_embed -def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: - """ - This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, - num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) - """ - batch, num_key_value_heads, slen, head_dim = hidden_states.shape - if n_rep == 1: - return hidden_states - hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) - return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) - - -def eager_attention_forward( - module: nn.Module, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attention_mask: torch.Tensor | None, - scaling: float, - dropout: float = 0.0, - **kwargs: Unpack[TransformersKwargs], -): - key_states = repeat_kv(key, module.num_key_value_groups) - value_states = repeat_kv(value, module.num_key_value_groups) - - attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling - if attention_mask is not None: - attn_weights = attn_weights + attention_mask - - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) - attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) - attn_output = torch.matmul(attn_weights, value_states) - attn_output = attn_output.transpose(1, 2).contiguous() - - return attn_output, attn_weights - - @use_kernelized_func(apply_rotary_pos_emb) class Qwen3_5Attention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" @@ -1687,7 +1688,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} 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 8fddbc6115c1..f01e3392c237 100644 --- a/src/transformers/models/qwen3_5/modular_qwen3_5.py +++ b/src/transformers/models/qwen3_5/modular_qwen3_5.py @@ -22,6 +22,7 @@ from ... import initialization as init from ...cache_utils import Cache, DynamicCache +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GenericForSequenceClassification, GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling @@ -88,15 +89,13 @@ class Qwen3_5TextConfig(Qwen3NextConfig): base_config_key = "text_config" 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.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("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 a33b33af7eff..1f7455c5be9e 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 @@ -20,6 +20,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -57,18 +58,18 @@ class Qwen3_5MoeTextConfig(PreTrainedConfig): 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.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.*.mlp.shared_expert.gate_proj": "colwise", - "layers.*.mlp.shared_expert.up_proj": "colwise", - "layers.*.mlp.shared_expert.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), + "layers.*.mlp.shared_expert.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_expert.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_expert.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), 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 125ded124cf7..3dd042042ffe 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 @@ -32,6 +32,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_experts_implementation, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -534,6 +535,43 @@ def forward( return output +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] @@ -580,43 +618,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): return q_embed, k_embed -def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: - """ - This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, - num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) - """ - batch, num_key_value_heads, slen, head_dim = hidden_states.shape - if n_rep == 1: - return hidden_states - hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) - return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) - - -def eager_attention_forward( - module: nn.Module, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attention_mask: torch.Tensor | None, - scaling: float, - dropout: float = 0.0, - **kwargs: Unpack[TransformersKwargs], -): - key_states = repeat_kv(key, module.num_key_value_groups) - value_states = repeat_kv(value, module.num_key_value_groups) - - attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling - if attention_mask is not None: - attn_weights = attn_weights + attention_mask - - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) - attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) - attn_output = torch.matmul(attn_weights, value_states) - attn_output = attn_output.transpose(1, 2).contiguous() - - return attn_output, attn_weights - - @use_kernelized_func(apply_rotary_pos_emb) class Qwen3_5MoeAttention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" @@ -1894,7 +1895,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} config: Qwen3_5MoeTextConfig _keys_to_ignore_on_load_unexpected = [r"^mtp.*", r"^model.visual.*"] @@ -2001,7 +2002,7 @@ class Qwen3_5MoeForConditionalGeneration(Qwen3_5MoePreTrainedModel, GenerationMi # Reference: fix gemma3 grad acc #37208 accepts_loss_kwargs = False config: Qwen3_5MoeConfig - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} 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 f3b4b80aa3a6..7f0baa834584 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 @@ -17,6 +17,7 @@ from huggingface_hub.dataclasses import strict from ... import initialization as init +from ...integrations.tensor_parallel import TPStyle from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPooling from ...modeling_utils import PreTrainedModel @@ -86,18 +87,18 @@ class Qwen3_5MoeTextConfig(Qwen3NextConfig): base_config_key = "text_config" 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.*.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.*.mlp.shared_expert.gate_proj": "colwise", - "layers.*.mlp.shared_expert.up_proj": "colwise", - "layers.*.mlp.shared_expert.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), + "layers.*.mlp.shared_expert.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_expert.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_expert.down_proj": TPStyle("rowwise", "allreduce"), } ignore_keys_at_rope_validation = {"mrope_section", "mrope_interleaved"} @@ -248,7 +249,7 @@ def __init__(self, config): class Qwen3_5MoeForConditionalGeneration(Qwen3VLMoeForConditionalGeneration): - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} 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 da79e911215b..4a69ea88e43c 100644 --- a/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py +++ b/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py @@ -16,6 +16,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -54,18 +55,18 @@ class Qwen3MoeConfig(PreTrainedConfig): # Default tensor parallel plan for base model `Qwen3Moe` 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.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py b/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py index ddf84fc575b7..ceec4a562b7d 100644 --- a/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py +++ b/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py @@ -29,12 +29,8 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import ( - use_experts_implementation, - use_kernel_forward_from_hub, - use_kernel_func_from_hub, - use_kernelized_func, -) +from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -53,39 +49,6 @@ from .configuration_qwen3_moe import Qwen3MoeConfig -def rotate_half(x): - """Rotates half the hidden dims of the input.""" - x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] - return torch.cat((-x2, x1), dim=-1) - - -@use_kernel_func_from_hub("rotary_pos_emb") -def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): - """Applies Rotary Position Embedding to the query and key tensors. - - Args: - q (`torch.Tensor`): The query tensor. - k (`torch.Tensor`): The key tensor. - cos (`torch.Tensor`): The cosine part of the rotary embedding. - sin (`torch.Tensor`): The sine part of the rotary embedding. - unsqueeze_dim (`int`, *optional*, defaults to 1): - The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and - sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note - that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and - k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes - cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have - the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. - Returns: - `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. - """ - cos = cos.unsqueeze(unsqueeze_dim) - sin = sin.unsqueeze(unsqueeze_dim) - q_embed = (q * cos) + (rotate_half(q) * sin) - k_embed = (k * cos) + (rotate_half(k) * sin) - return q_embed, k_embed - - def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, @@ -123,6 +86,20 @@ def eager_attention_forward( return attn_output, attn_weights +def rotate_half(x): + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + @use_kernelized_func(apply_rotary_pos_emb) class Qwen3MoeAttention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" @@ -609,7 +586,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, 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..c8bcaf895d93 100644 --- a/src/transformers/models/qwen3_next/configuration_qwen3_next.py +++ b/src/transformers/models/qwen3_next/configuration_qwen3_next.py @@ -16,6 +16,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -59,21 +60,21 @@ class Qwen3NextConfig(PreTrainedConfig): 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", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "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.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.shared_expert.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_expert.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_expert.down_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/qwen3_next/modeling_qwen3_next.py b/src/transformers/models/qwen3_next/modeling_qwen3_next.py index cd152e3d3e59..a991c45fedac 100644 --- a/src/transformers/models/qwen3_next/modeling_qwen3_next.py +++ b/src/transformers/models/qwen3_next/modeling_qwen3_next.py @@ -30,6 +30,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_experts_implementation, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -169,6 +170,43 @@ def extra_repr(self): return f"{tuple(self.weight.shape)}, eps={self.eps}" +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] @@ -215,43 +253,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): return q_embed, k_embed -def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: - """ - This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, - num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) - """ - batch, num_key_value_heads, slen, head_dim = hidden_states.shape - if n_rep == 1: - return hidden_states - hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) - return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) - - -def eager_attention_forward( - module: nn.Module, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attention_mask: torch.Tensor | None, - scaling: float, - dropout: float = 0.0, - **kwargs: Unpack[TransformersKwargs], -): - key_states = repeat_kv(key, module.num_key_value_groups) - value_states = repeat_kv(value, module.num_key_value_groups) - - attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling - if attention_mask is not None: - attn_weights = attn_weights + attention_mask - - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) - attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) - attn_output = torch.matmul(attn_weights, value_states) - attn_output = attn_output.transpose(1, 2).contiguous() - - return attn_output, attn_weights - - @use_kernelized_func(apply_rotary_pos_emb) class Qwen3NextAttention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" @@ -1072,7 +1073,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, 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 04534187d73b..a56602fa2047 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 @@ -131,15 +131,13 @@ class Qwen3OmniMoeTextConfig(PreTrainedConfig): # Default tensor parallel plan for base model `Qwen3OmniMoeText` 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.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -375,18 +373,18 @@ class Qwen3OmniMoeTalkerTextConfig(PreTrainedConfig): # Default tensor parallel plan for base model `Qwen3OmniMoeTalkerText` 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.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), 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 7b6c8b5b1bd4..b3e30d387246 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 @@ -34,12 +34,8 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import ( - use_experts_implementation, - use_kernel_forward_from_hub, - use_kernel_func_from_hub, - use_kernelized_func, -) +from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -814,7 +810,6 @@ def _get_feat_extract_output_lengths(self, input_lengths: torch.LongTensor): def rotate_half(x): - """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) @@ -1444,25 +1439,7 @@ def extra_repr(self): return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" -@use_kernel_func_from_hub("rotary_pos_emb") def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): - """Applies Rotary Position Embedding to the query and key tensors. - - Args: - q (`torch.Tensor`): The query tensor. - k (`torch.Tensor`): The key tensor. - cos (`torch.Tensor`): The cosine part of the rotary embedding. - sin (`torch.Tensor`): The sine part of the rotary embedding. - unsqueeze_dim (`int`, *optional*, defaults to 1): - The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and - sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note - that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and - k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes - cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have - the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. - Returns: - `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. - """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) q_embed = (q * cos) + (rotate_half(q) * sin) @@ -2644,7 +2621,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} config_class = Qwen3OmniMoeTalkerCodePredictorConfig base_model_prefix = "talker.code_predictor" @@ -3028,7 +3005,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"codec_head": (["hidden_states"], ["logits"])} config_class = Qwen3OmniMoeTalkerConfig base_model_prefix = "talker" 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 d336784b3b49..8b8cbdc8218e 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 @@ -32,6 +32,7 @@ from ...feature_extraction_utils import BatchFeature from ...generation import GenerationMixin from ...image_utils import ImageInput +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( @@ -182,15 +183,13 @@ class Qwen3OmniMoeTextConfig(PreTrainedConfig): # Default tensor parallel plan for base model `Qwen3OmniMoeText` 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.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -1578,7 +1577,7 @@ def get_input_embeddings(self): class Qwen3OmniMoeTalkerForConditionalGeneration(Qwen3MoeForCausalLM): _tied_weights_keys = {"codec_head": "model.codec_embedding.weight"} - _tp_plan = {"codec_head": "colwise_gather_output"} + _tp_plan = {"codec_head": TPStyle("colwise", "allgather")} _pp_plan = {"codec_head": (["hidden_states"], ["logits"])} config_class = Qwen3OmniMoeTalkerConfig base_model_prefix = "talker" diff --git a/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py b/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py index dee0a4e04420..ec276a41d8e9 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 @@ -20,6 +20,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -56,13 +57,13 @@ class Qwen3VLMoeTextConfig(PreTrainedConfig): } # Default tensor parallel plan for base model `Qwen3VLMoe` 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", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), 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 1d5159d37f6a..1d6b4b3fcfd5 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 @@ -20,6 +20,7 @@ from ... import initialization as init from ...cache_utils import Cache, DynamicCache +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_outputs import MoeModelOutputWithPast @@ -83,13 +84,13 @@ class Qwen3VLMoeTextConfig(Qwen3MoeConfig): default_theta = 500000.0 # Default tensor parallel plan for base model `Qwen3VLMoe` 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", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/seed_oss/configuration_seed_oss.py b/src/transformers/models/seed_oss/configuration_seed_oss.py index b1221fcf53ce..7be4662965c1 100644 --- a/src/transformers/models/seed_oss/configuration_seed_oss.py +++ b/src/transformers/models/seed_oss/configuration_seed_oss.py @@ -16,6 +16,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -45,13 +46,13 @@ class SeedOssConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `SeedOssModel` 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", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/seed_oss/modeling_seed_oss.py b/src/transformers/models/seed_oss/modeling_seed_oss.py index 1ebc8f10a272..5c144edaa94c 100644 --- a/src/transformers/models/seed_oss/modeling_seed_oss.py +++ b/src/transformers/models/seed_oss/modeling_seed_oss.py @@ -28,6 +28,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import ( GenericForQuestionAnswering, @@ -429,7 +430,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/smollm3/configuration_smollm3.py b/src/transformers/models/smollm3/configuration_smollm3.py index f48c979a1dd3..a4898567d4c0 100644 --- a/src/transformers/models/smollm3/configuration_smollm3.py +++ b/src/transformers/models/smollm3/configuration_smollm3.py @@ -21,6 +21,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -55,13 +56,13 @@ class SmolLM3Config(PreTrainedConfig): default_theta = 2000000.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.o_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/smollm3/modeling_smollm3.py b/src/transformers/models/smollm3/modeling_smollm3.py index 8d911e414b0f..82ad64435bd8 100644 --- a/src/transformers/models/smollm3/modeling_smollm3.py +++ b/src/transformers/models/smollm3/modeling_smollm3.py @@ -28,6 +28,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -445,7 +446,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/smollm3/modular_smollm3.py b/src/transformers/models/smollm3/modular_smollm3.py index f75017ad2645..d62061154d50 100644 --- a/src/transformers/models/smollm3/modular_smollm3.py +++ b/src/transformers/models/smollm3/modular_smollm3.py @@ -19,6 +19,7 @@ from ...cache_utils import Cache from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_rope_utils import RopeParameters from ...modeling_utils import ALL_ATTENTION_FUNCTIONS @@ -71,13 +72,13 @@ class SmolLM3Config(PreTrainedConfig): default_theta = 2000000.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.o_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/solar_open/configuration_solar_open.py b/src/transformers/models/solar_open/configuration_solar_open.py index ac0016aa7791..9e35341a59bc 100644 --- a/src/transformers/models/solar_open/configuration_solar_open.py +++ b/src/transformers/models/solar_open/configuration_solar_open.py @@ -21,6 +21,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -38,13 +39,15 @@ class SolarOpenConfig(PreTrainedConfig): # Default tensor parallel plan for base model `SolarOpenModel` 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.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/solar_open/modeling_solar_open.py b/src/transformers/models/solar_open/modeling_solar_open.py index dfa30292455f..cb3f35fb3700 100644 --- a/src/transformers/models/solar_open/modeling_solar_open.py +++ b/src/transformers/models/solar_open/modeling_solar_open.py @@ -34,6 +34,7 @@ use_kernel_func_from_hub, use_kernelized_func, ) +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -552,7 +553,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, 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..8363cfec333e 100644 --- a/src/transformers/models/solar_open/modular_solar_open.py +++ b/src/transformers/models/solar_open/modular_solar_open.py @@ -16,6 +16,7 @@ from huggingface_hub.dataclasses import strict from torch import nn +from ...integrations.tensor_parallel import TPStyle from ...utils import auto_docstring, logging from ..glm4_moe.configuration_glm4_moe import Glm4MoeConfig from ..glm4_moe.modeling_glm4_moe import ( @@ -44,13 +45,15 @@ class SolarOpenConfig(Glm4MoeConfig): # Default tensor parallel plan for base model `SolarOpenModel` 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.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), } vocab_size: int = 196608 diff --git a/src/transformers/models/starcoder2/configuration_starcoder2.py b/src/transformers/models/starcoder2/configuration_starcoder2.py index 59efa94fc5f4..f55d8d17ab0c 100644 --- a/src/transformers/models/starcoder2/configuration_starcoder2.py +++ b/src/transformers/models/starcoder2/configuration_starcoder2.py @@ -16,6 +16,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -45,12 +46,12 @@ class Starcoder2Config(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `Starcoder2` 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.c_fc": "colwise", - "layers.*.mlp.c_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.c_fc": TPStyle("colwise", "none"), + "layers.*.mlp.c_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/starcoder2/modeling_starcoder2.py b/src/transformers/models/starcoder2/modeling_starcoder2.py index 8b89a1d1745c..26fcc4e8b435 100644 --- a/src/transformers/models/starcoder2/modeling_starcoder2.py +++ b/src/transformers/models/starcoder2/modeling_starcoder2.py @@ -33,6 +33,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_func_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -409,7 +410,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/t5gemma/configuration_t5gemma.py b/src/transformers/models/t5gemma/configuration_t5gemma.py index 9de40c832259..9a5f93ba8a53 100644 --- a/src/transformers/models/t5gemma/configuration_t5gemma.py +++ b/src/transformers/models/t5gemma/configuration_t5gemma.py @@ -23,6 +23,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -51,13 +52,13 @@ class T5GemmaModuleConfig(PreTrainedConfig): model_type = "t5_gemma_module" 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.o_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/t5gemma/modeling_t5gemma.py b/src/transformers/models/t5gemma/modeling_t5gemma.py index a6b9b5392194..65f4cb52d846 100644 --- a/src/transformers/models/t5gemma/modeling_t5gemma.py +++ b/src/transformers/models/t5gemma/modeling_t5gemma.py @@ -29,6 +29,7 @@ from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache from ...generation import GenerationMixin from ...integrations import use_kernel_func_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import ( create_bidirectional_mask, create_bidirectional_sliding_window_mask, @@ -945,7 +946,7 @@ def forward( class T5GemmaForConditionalGeneration(T5GemmaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.out_proj.weight": "model.decoder.embed_tokens.weight"} - _tp_plan = {"lm_head.out_proj": "colwise_gather_output"} + _tp_plan = {"lm_head.out_proj": TPStyle("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 c7d4a4051959..363e3ca76430 100644 --- a/src/transformers/models/t5gemma/modular_t5gemma.py +++ b/src/transformers/models/t5gemma/modular_t5gemma.py @@ -23,6 +23,7 @@ from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache from ...configuration_utils import PreTrainedConfig from ...generation import GenerationMixin +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import ( create_bidirectional_mask, create_bidirectional_sliding_window_mask, @@ -784,7 +785,7 @@ def forward( class T5GemmaForConditionalGeneration(T5GemmaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.out_proj.weight": "model.decoder.embed_tokens.weight"} - _tp_plan = {"lm_head.out_proj": "colwise_gather_output"} + _tp_plan = {"lm_head.out_proj": TPStyle("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..1f04bd70d9e6 100644 --- a/src/transformers/models/t5gemma2/configuration_t5gemma2.py +++ b/src/transformers/models/t5gemma2/configuration_t5gemma2.py @@ -23,6 +23,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...utils import auto_docstring, logging from ..siglip import SiglipVisionConfig @@ -45,15 +46,13 @@ class T5Gemma2TextConfig(PreTrainedConfig): model_type = "t5gemma2_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", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -220,15 +219,13 @@ class T5Gemma2DecoderConfig(PreTrainedConfig): model_type = "t5gemma2_decoder" 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", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/t5gemma2/modeling_t5gemma2.py b/src/transformers/models/t5gemma2/modeling_t5gemma2.py index 2e0dddc17876..3de40aeaec64 100644 --- a/src/transformers/models/t5gemma2/modeling_t5gemma2.py +++ b/src/transformers/models/t5gemma2/modeling_t5gemma2.py @@ -30,6 +30,7 @@ from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache, StaticCache from ...generation import GenerationConfig, GenerationMixin, GenerationMode from ...integrations import use_kernel_func_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_bidirectional_mask, create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -1185,7 +1186,7 @@ class T5Gemma2ForConditionalGeneration(T5Gemma2PreTrainedModel, GenerationMixin) _tied_weights_keys = { "lm_head.out_proj.weight": "model.encoder.text_model.embed_tokens.weight", } - _tp_plan = {"lm_head.out_proj": "colwise_gather_output"} + _tp_plan = {"lm_head.out_proj": TPStyle("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 2f0f3720a7cd..db8b82ec6b88 100644 --- a/src/transformers/models/t5gemma2/modular_t5gemma2.py +++ b/src/transformers/models/t5gemma2/modular_t5gemma2.py @@ -24,6 +24,7 @@ from ...cache_utils import DynamicCache, EncoderDecoderCache, StaticCache from ...configuration_utils import PreTrainedConfig from ...generation import GenerationConfig, GenerationMixin, GenerationMode +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_bidirectional_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_outputs import ( @@ -973,7 +974,7 @@ class T5Gemma2ForConditionalGeneration(T5Gemma2PreTrainedModel, GenerationMixin) _tied_weights_keys = { "lm_head.out_proj.weight": "model.encoder.text_model.embed_tokens.weight", } - _tp_plan = {"lm_head.out_proj": "colwise_gather_output"} + _tp_plan = {"lm_head.out_proj": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head.out_proj": (["hidden_states"], ["logits"])} def __init__(self, config: T5Gemma2Config): diff --git a/src/transformers/models/vaultgemma/configuration_vaultgemma.py b/src/transformers/models/vaultgemma/configuration_vaultgemma.py index a60b7e8edc0c..1fd184568ded 100644 --- a/src/transformers/models/vaultgemma/configuration_vaultgemma.py +++ b/src/transformers/models/vaultgemma/configuration_vaultgemma.py @@ -22,6 +22,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -50,13 +51,13 @@ class VaultGemmaConfig(PreTrainedConfig): model_type = "vaultgemma" 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.o_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/vaultgemma/modeling_vaultgemma.py b/src/transformers/models/vaultgemma/modeling_vaultgemma.py index f0a2e48d20b8..60122bfcb923 100644 --- a/src/transformers/models/vaultgemma/modeling_vaultgemma.py +++ b/src/transformers/models/vaultgemma/modeling_vaultgemma.py @@ -30,6 +30,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_func_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -467,7 +468,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py b/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py index 07325b0ea559..1dafcb67da56 100644 --- a/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py +++ b/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py @@ -32,6 +32,7 @@ from ...cache_utils import Cache, DynamicCache, StaticCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -823,7 +824,7 @@ def forward( @auto_docstring class VoxtralRealtimeTextForCausalLM(VoxtralRealtimeTextPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/youtu/configuration_youtu.py b/src/transformers/models/youtu/configuration_youtu.py index 61019d70391b..eaef7650cd34 100644 --- a/src/transformers/models/youtu/configuration_youtu.py +++ b/src/transformers/models/youtu/configuration_youtu.py @@ -27,6 +27,7 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig +from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -51,9 +52,9 @@ class YoutuConfig(PreTrainedConfig): model_type = "youtu" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/youtu/modeling_youtu.py b/src/transformers/models/youtu/modeling_youtu.py index d40bef358da6..07e7c639f979 100644 --- a/src/transformers/models/youtu/modeling_youtu.py +++ b/src/transformers/models/youtu/modeling_youtu.py @@ -37,6 +37,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub +from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -533,7 +534,7 @@ 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": TPStyle("colwise", "allgather")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/youtu/modular_youtu.py b/src/transformers/models/youtu/modular_youtu.py index b2de3a2df0a5..a4ee903c61e1 100644 --- a/src/transformers/models/youtu/modular_youtu.py +++ b/src/transformers/models/youtu/modular_youtu.py @@ -23,6 +23,7 @@ from torch import nn from ... import initialization as init +from ...integrations.tensor_parallel import TPStyle from ...modeling_utils import PreTrainedModel from ...utils import auto_docstring, logging from ..deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config @@ -60,9 +61,9 @@ class YoutuConfig(DeepseekV3Config): model_type = "youtu" base_model_tp_plan = { - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } attribute_map = {} diff --git a/tests/tensor_parallel/test_tensor_parallel.py b/tests/tensor_parallel/test_tensor_parallel.py index b2e19d91bd72..e81aa342d635 100644 --- a/tests/tensor_parallel/test_tensor_parallel.py +++ b/tests/tensor_parallel/test_tensor_parallel.py @@ -11,21 +11,9 @@ # 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, -) from transformers.testing_utils import TestCasePlus, is_tensor_parallel_test @@ -58,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.""" @@ -133,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) From 5ce6faa33519ca51e9224c9cb22615c0b4a2be0f Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 14 Apr 2026 17:26:57 +0000 Subject: [PATCH 018/116] Restore mxfp4.py to match base branch --- src/transformers/integrations/mxfp4.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/transformers/integrations/mxfp4.py b/src/transformers/integrations/mxfp4.py index d1e5506cbf20..482ef891771f 100644 --- a/src/transformers/integrations/mxfp4.py +++ b/src/transformers/integrations/mxfp4.py @@ -20,7 +20,7 @@ from torch import nn from contextlib import contextmanager -from ..core_model_loading import ConversionOps +from ..core_model_loading import ConversionOps, _IdentityOp from ..quantizers.quantizers_utils import get_module_from_name, should_convert_module @@ -145,6 +145,10 @@ def convert( dequantized = dequantize_convertops(param_data[f"{proj}_blocks"], param_data[f"{proj}_scales"]) return {full_layer_name: dequantized} + @property + def reverse_op(self) -> "ConversionOps": + return _IdentityOp() + class Mxfp4Deserialize(ConversionOps): def __init__(self, hf_quantizer): From b694f3641f898486aefaf1e9eaf53723c925a5b8 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 14 Apr 2026 17:31:27 +0000 Subject: [PATCH 019/116] Drop mla_kv_a_proj and moe_identity_expert from TP plans These string plan values have no TPStyle equivalent in the DTensor system. Remove them to avoid TypeError at apply_tensor_parallel time. Affected models: deepseek_v2, glm4_moe_lite, glm_moe_dsa, longcat_flash. --- .../models/deepseek_v2/configuration_deepseek_v2.py | 1 - src/transformers/models/deepseek_v2/modular_deepseek_v2.py | 1 - .../models/glm4_moe_lite/configuration_glm4_moe_lite.py | 1 - src/transformers/models/glm4_moe_lite/modular_glm4_moe_lite.py | 1 - .../models/glm_moe_dsa/configuration_glm_moe_dsa.py | 1 - src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py | 1 - .../models/longcat_flash/configuration_longcat_flash.py | 2 -- 7 files changed, 8 deletions(-) diff --git a/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py b/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py index 626ba3a495e9..52f03ec0cc58 100644 --- a/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py +++ b/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py @@ -55,7 +55,6 @@ class DeepseekV2Config(PreTrainedConfig): base_model_tp_plan = { "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), "layers.*.self_attn.q_b_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj", "layers.*.self_attn.kv_b_proj": TPStyle("colwise", "none"), "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), "layers.*.mlp.experts": TPStyle( diff --git a/src/transformers/models/deepseek_v2/modular_deepseek_v2.py b/src/transformers/models/deepseek_v2/modular_deepseek_v2.py index 00683225189f..30f439499c45 100644 --- a/src/transformers/models/deepseek_v2/modular_deepseek_v2.py +++ b/src/transformers/models/deepseek_v2/modular_deepseek_v2.py @@ -70,7 +70,6 @@ class DeepseekV2Config(LlamaConfig): base_model_tp_plan = { "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), "layers.*.self_attn.q_b_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj", "layers.*.self_attn.kv_b_proj": TPStyle("colwise", "none"), "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), "layers.*.mlp.experts": TPStyle( 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 194eef2a695b..012d1d06526b 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 @@ -54,7 +54,6 @@ class Glm4MoeLiteConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { "layers.*.self_attn.q_b_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj", "layers.*.self_attn.kv_b_proj": TPStyle("colwise", "none"), "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), "layers.*.mlp.experts": TPStyle( 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 e242621817a3..06741613822e 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 @@ -62,7 +62,6 @@ class Glm4MoeLiteConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { "layers.*.self_attn.q_b_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj", "layers.*.self_attn.kv_b_proj": TPStyle("colwise", "none"), "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), "layers.*.mlp.experts": TPStyle( 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 68285a8f6d71..88de8f2f2e3e 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 @@ -59,7 +59,6 @@ class GlmMoeDsaConfig(PreTrainedConfig): base_model_tp_plan = { "layers.*.self_attn.q_b_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj", "layers.*.self_attn.kv_b_proj": TPStyle("colwise", "none"), "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), "layers.*.mlp.experts": TPStyle( 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 4789c82aac64..7ca32fdd2143 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 @@ -105,7 +105,6 @@ class GlmMoeDsaConfig(Glm4MoeLiteConfig): base_model_tp_plan = { "layers.*.self_attn.q_b_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj", "layers.*.self_attn.kv_b_proj": TPStyle("colwise", "none"), "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), "layers.*.mlp.experts": TPStyle( diff --git a/src/transformers/models/longcat_flash/configuration_longcat_flash.py b/src/transformers/models/longcat_flash/configuration_longcat_flash.py index 5547c24deeb5..8002cb36b00f 100644 --- a/src/transformers/models/longcat_flash/configuration_longcat_flash.py +++ b/src/transformers/models/longcat_flash/configuration_longcat_flash.py @@ -56,10 +56,8 @@ class LongcatFlashConfig(PreTrainedConfig): default_theta = 10000000.0 base_model_tp_plan = { "layers.*.self_attn.*.q_b_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.*.kv_a_proj_with_mqa": "mla_kv_a_proj", "layers.*.self_attn.*.kv_b_proj": TPStyle("colwise", "none"), "layers.*.self_attn.*.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.experts.identity_expert": "moe_identity_expert", "layers.*.mlp.experts": TPStyle( "moe_experts", "allreduce", From 1b82460a469d2d31f8ee52b4bfb5c684bb649aa4 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 15 Apr 2026 14:13:07 +0000 Subject: [PATCH 020/116] more comments --- src/transformers/core_model_loading.py | 20 ++++++++++---- tests/utils/test_core_model_loading.py | 38 +++++++++++++------------- 2 files changed, 33 insertions(+), 25 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 33a07251c0f1..8b1be399400e 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -871,7 +871,7 @@ def shard_tensor( self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None ) -> torch.Tensor | None: """Return the local shard of ``param`` for this rank, dispatching to the appropriate strategy.""" - # Find which placements actually shard data. + # Find sharding placements (keep Shard and _StridedShard only) # _StridedShard.is_shard() returns False in PyTorch, so we also check for # the ``dim`` attribute that both Shard and _StridedShard have. sharding_placements = [ @@ -881,32 +881,38 @@ def shard_tensor( ] param_shape = list(param.shape) if isinstance(param, torch.Tensor) else param.get_shape() + # [A] No sharding placements -> Return full copy if not sharding_placements: return param[...].to(device=device, dtype=dtype) + if tensor_idx is not None and len(self.param.shape) == len(param_shape) + 1: - # Expert parallelism: dim 0 (expert dimension) is sharded across ranks. + # [B] Expert path: shard on expert dimension (dim 0). # When dim 0 is the only sharding placement, return the full expert or # skip it. When TP also shards an inner dim, keep applying the remaining # placements to the owned expert tensor. has_expert_sharding = any(self._normalize_param_dim(p.dim) == 0 for _, p in sharding_placements) if has_expert_sharding: + # [B2] This rank doesn't own the expert tensor -> skip it if not self._owns_local_expert(tensor_idx): return None inner_placements = [(i, p) for i, p in sharding_placements if self._normalize_param_dim(p.dim) != 0] + # [B3] Not composed with TP placements -> return full copy of the expert tensor if not inner_placements: return param[...].to(device=device, dtype=dtype) + # [B4] Composed with TP placements -> shard the expert's inner dims return self._shard_nd(param, inner_placements, param_shape, device, dtype) - + + # [B1] has_expert_sharding=False -> fall through to _shard_nd return self._shard_nd(param, sharding_placements, param_shape, device, dtype) def _shard_nd(self, param, sharding_placements, param_shape, device, dtype): """Handle multi-dimensional sharding, choosing the best strategy.""" + # [C1] Column Parallel when composed with FSDP. We choose the easier path but maybe we should do a better one? if not self._can_shard_on_read(sharding_placements): return self._materialize_and_split(param, sharding_placements, device, dtype) - # All placements are plain Shard on different dims. - # compute_local_shape_and_global_offset gives us one contiguous range per dim directly. + # [C2] All sharding placements are plain Shard on different dims -> single contiguous slice has_strided = any(not p.is_shard() for _, p in sharding_placements) if not has_strided: local_shape, global_offset = compute_local_shape_and_global_offset( @@ -919,6 +925,7 @@ def _shard_nd(self, param, sharding_placements, param_shape, device, dtype): slices[dim] = slice(offset, offset + local_shape[placement.dim]) return param[tuple(slices)].to(device=device, dtype=dtype) + # [C3] At least one _StridedShard (no same-dim conflict) -> _compute_dim_ranges + _slice_and_read dim_ranges = self._compute_dim_ranges(sharding_placements, param_shape) return self._slice_and_read(param, param_shape, dim_ranges, device, dtype) @@ -972,12 +979,13 @@ def _compute_dim_ranges(self, sharding_placements, param_shape) -> DtensorShardO if placement.is_shard(): new_ranges = self._contiguous_ranges(prev_ranges, rank, world_size) elif self._source_tensor_needs_packing(param_shape): - # _StridedShard only makes sense once the packed axis exists. While + # [C3a] _StridedShard only makes sense once the packed axis exists. While # loading pre-packed source tensors (e.g. w1/w3 before gate_up_proj # concatenation), take the contiguous chunk for this rank and let the # WeightConverter recreate the packed layout afterward. new_ranges = self._contiguous_ranges(prev_ranges, rank, world_size) else: + # [C3b] Normal strided -> disjoint ranges + cat new_ranges = self._strided_ranges(prev_ranges, rank, world_size, placement.split_factor) dim_ranges[dim] = new_ranges return dim_ranges diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py index 787cf7b903ad..d68ad2621437 100644 --- a/tests/utils/test_core_model_loading.py +++ b/tests/utils/test_core_model_loading.py @@ -878,29 +878,29 @@ def test_ernie4_5_vl_moe_conversion_reversed(self): class TestDtensorShardOperation(unittest.TestCase): """Unit tests for DtensorShardOperation.shard_tensor — one test per code path. - Branch coverage map: + Branch coverage map (labels [A]–[C3b] match comments in core_model_loading.py): shard_tensor() - ├── A: no sharding placements → full copy [test_no_shard_returns_full_tensor] - ├── B: expert path (tensor_idx set, ndim mismatch) - │ ├── B1: has_expert_sharding=False → fall through to C [test_expert_shaped_tp_only_no_expert_sharding] - │ ├── B2: not owns_local_expert → None [test_expert_filtering] - │ ├── B3: owned, no inner placements → full copy [test_expert_filtering] - │ └── B4: owned, with inner placements → _shard_nd [test_expert_filtering_preserves_inner_sharding] - └── C: _shard_nd() - ├── C1: _can_shard_on_read=False → _materialize_and_split [test_nd_strided_plus_shard_same_dim_fallback] - ├── C2: has_strided=False → contiguous slice - │ ├── 1D mesh [test_1d_shard_fast_path] - │ ├── 2D mesh [test_nd_contiguous_single_slice] - │ ├── negative dim [test_negative_dim_normalizes_correctly] - │ └── uneven division [test_contiguous_shard_uneven_division] - └── C3: has_strided=True → _compute_dim_ranges + _slice_and_read - ├── _StridedShard → _strided_ranges [test_nd_strided_shard_disjoint_ranges] - └── _source_tensor_needs_packing → contiguous [test_prepacked_strided_shard_uses_contiguous_source_slice] + ├── [A] no sharding placements → full copy [test_no_shard_returns_full_tensor] + ├── [B] expert path (tensor_idx set, ndim mismatch) + │ ├── [B1] has_expert_sharding=False → fall through to C [test_expert_shaped_tp_only_no_expert_sharding] + │ ├── [B2] not owns_local_expert → None [test_expert_filtering] + │ ├── [B3] owned, no inner placements → full copy [test_expert_filtering] + │ └── [B4] owned, with inner placements → _shard_nd [test_expert_filtering_preserves_inner_sharding] + └── [C] _shard_nd() + ├── [C1] _can_shard_on_read=False → _materialize_and_split [test_nd_strided_plus_shard_same_dim_fallback] + ├── [C2] has_strided=False → contiguous slice + │ ├── 1D mesh [test_1d_shard_fast_path] + │ ├── 2D mesh [test_nd_contiguous_single_slice] + │ ├── negative dim [test_negative_dim_normalizes_correctly] + │ └── uneven division [test_contiguous_shard_uneven_division] + └── [C3] has_strided=True → _compute_dim_ranges + _slice_and_read + ├── [C3a] _source_tensor_needs_packing → contiguous [test_prepacked_strided_shard_uses_contiguous_source_slice] + └── [C3b] _StridedShard → _strided_ranges [test_nd_strided_shard_disjoint_ranges] _slice_and_read (tested directly) - ├── all single ranges → simple slice [test_slice_and_read_all_single_ranges] - └── two multi-range dims → ValueError [test_slice_and_read_raises_on_two_multi_range_dims] + ├── all single ranges → simple slice [test_slice_and_read_all_single_ranges] + └── two multi-range dims → ValueError [test_slice_and_read_raises_on_two_multi_range_dims] """ def test_no_shard_returns_full_tensor(self): From 48f8d6f927ef69fd3076db6c88466984cdf6d6d5 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 15 Apr 2026 15:15:49 +0000 Subject: [PATCH 021/116] =?UTF-8?q?fix=20tp=20for=20most=20models.=20=20Py?= =?UTF-8?q?Torch=20doesn't=20implement=20all=20placement=20conversions=20(?= =?UTF-8?q?e.g.=20=5FStridedShard=E2=86=94Shard).=20We=20force=20replicate?= =?UTF-8?q?=20beforehand?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../integrations/tensor_parallel.py | 29 ++++++++++--------- tests/test_tensor_parallel_mixin.py | 4 ++- 2 files changed, 19 insertions(+), 14 deletions(-) diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index a4fd1cc38ca4..82973cf8e11f 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -124,9 +124,7 @@ def gather_full_state_dict(model) -> dict[str, torch.Tensor]: if isinstance(tensor, DTensor): # All ranks participate in the collective, only rank 0 keeps the result with torch.no_grad(): - full = tensor.redistribute( - placements=[Replicate()] * tensor.device_mesh.ndim, async_op=False - ).to_local() + full = _replicate_dtensor(tensor).to_local() if is_rank0: result[key] = _to_cpu_fresh(full) del full @@ -136,15 +134,19 @@ def gather_full_state_dict(model) -> dict[str, torch.Tensor]: return result -def _redistribute_dtensor(tensor: DTensor, target_placements: tuple) -> DTensor: - """Redistribute a DTensor via Replicate as an intermediate step. +def _replicate_dtensor(tensor: DTensor) -> DTensor: + """All-gather a DTensor to Replicate, handling _StridedShard placements. - PyTorch doesn't implement all placement conversions (e.g. _StridedShard↔Shard). - Going through Replicate first is always supported. + PyTorch's ``redistribute`` does not support ``_StridedShard`` as a source, + so we use each placement's ``_to_replicate_tensor`` directly. """ + mesh = tensor.device_mesh with torch.no_grad(): - replicated = tensor.redistribute(placements=[Replicate()] * tensor.device_mesh.ndim) - return replicated.redistribute(placements=target_placements) + local = tensor._local_tensor + for i, p in enumerate(tensor.placements): + if not p.is_replicate(): + local = p._to_replicate_tensor(local, mesh, i, tensor.shape) + return DTensor.from_local(local, mesh, [Replicate()] * mesh.ndim, run_check=False) def convert_strided_to_shard(state_dict: dict) -> dict[str, tuple]: @@ -158,7 +160,7 @@ def convert_strided_to_shard(state_dict: dict) -> dict[str, tuple]: elif isinstance(value, DTensor) and any(isinstance(p, _StridedShard) for p in value.placements): placement_map[key] = tuple(value.placements) shard_placements = tuple(Shard(p.dim) if isinstance(p, _StridedShard) else p for p in value.placements) - state_dict[key] = _redistribute_dtensor(value, shard_placements) + state_dict[key] = _replicate_dtensor(value).redistribute(placements=shard_placements) return placement_map @@ -173,7 +175,7 @@ def _resolve(d, dotted_key): for key, original_placements in placement_map.items(): container, leaf_key = _resolve(state_dict, key) if leaf_key in container and isinstance(container[leaf_key], DTensor): - container[leaf_key] = _redistribute_dtensor(container[leaf_key], original_placements) + container[leaf_key] = _replicate_dtensor(container[leaf_key]).redistribute(placements=original_placements) def verify_tp_plan(expected_keys: list[str], tp_plan: dict[str, str | TPStyle] | None): @@ -405,14 +407,15 @@ def _uses_partial_outputs(mod) -> bool: # partial hidden-state contribution that must be reduced. Under TP+FSDP, # FSDP can swap in full gathered expert weights for the current rank's # forward, in which case the local output is already complete. + intermediate = getattr(mod, "intermediate_dim", None) or getattr(mod, "intermediate_size", None) if hasattr(mod, "gate_up_proj"): gate_up_proj = mod.gate_up_proj.to_local() if isinstance(mod.gate_up_proj, DTensor) else mod.gate_up_proj - full_expert_out = 2 * mod.intermediate_dim + full_expert_out = 2 * intermediate sharded_dim = -1 if getattr(mod, "is_transposed", False) else -2 cached = gate_up_proj.shape[sharded_dim] != full_expert_out elif hasattr(mod, "up_proj"): up_proj = mod.up_proj.to_local() if isinstance(mod.up_proj, DTensor) else mod.up_proj - full_expert_out = mod.intermediate_dim + full_expert_out = intermediate sharded_dim = -1 if getattr(mod, "is_transposed", False) else -2 cached = up_proj.shape[sharded_dim] != full_expert_out else: diff --git a/tests/test_tensor_parallel_mixin.py b/tests/test_tensor_parallel_mixin.py index e07d2beb539b..f4a9e54a8691 100644 --- a/tests/test_tensor_parallel_mixin.py +++ b/tests/test_tensor_parallel_mixin.py @@ -48,7 +48,9 @@ def _to_local(tensor): # Partial gradients for us. if isinstance(tensor, DTensor) and any(not p.is_replicate() for p in tensor.placements): - tensor = tensor.redistribute(placements=(Replicate(),)) + from transformers.integrations.tensor_parallel import _replicate_dtensor + + tensor = _replicate_dtensor(tensor) return tensor.to_local() return tensor From 91b48242b3839569c71a69ef70ebf47437004182 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Thu, 16 Apr 2026 10:46:18 +0000 Subject: [PATCH 022/116] fix tp through _replicate_dtensor --- .../integrations/tensor_parallel.py | 49 ++++++++++++++++--- tests/test_tensor_parallel_mixin.py | 4 +- 2 files changed, 42 insertions(+), 11 deletions(-) diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index 82973cf8e11f..6a2a043c8b57 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -135,18 +135,51 @@ def gather_full_state_dict(model) -> dict[str, torch.Tensor]: def _replicate_dtensor(tensor: DTensor) -> DTensor: - """All-gather a DTensor to Replicate, handling _StridedShard placements. + """All-gather a DTensor to fully Replicate, handling ``_StridedShard``. - PyTorch's ``redistribute`` does not support ``_StridedShard`` as a source, - so we use each placement's ``_to_replicate_tensor`` directly. + 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 + + So we bypass ``redistribute`` and call each placement's low-level + ``_to_replicate_tensor`` (manual all-gather + interleaved reorder). + + We process mesh dims **right-to-left** (innermost first). Under TP+FSDP + the 2D mesh is ``(fsdp, tp)`` and both dims can shard the same tensor dim:: + + placements = (_StridedShard(dim=0), Shard(dim=0)) + local shape = [64, 1024] (global [256, 1024], fsdp=2, tp=2) + + Right-to-left means TP is gathered first (local grows to [128, 1024]), + then FSDP (grows to [256, 1024]). Each step must pass the correct + intermediate logical shape — the global shape divided by the mesh sizes + of dims not yet gathered (to the left). """ mesh = tensor.device_mesh + replicate_all = tuple(Replicate() for _ in range(mesh.ndim)) with torch.no_grad(): - local = tensor._local_tensor - for i, p in enumerate(tensor.placements): - if not p.is_replicate(): - local = p._to_replicate_tensor(local, mesh, i, tensor.shape) - return DTensor.from_local(local, mesh, [Replicate()] * mesh.ndim, run_check=False) + if any(isinstance(p, _StridedShard) for p in tensor.placements): + local = tensor._local_tensor + placements = tensor.placements + for i in reversed(range(mesh.ndim)): + p = placements[i] + if p.is_replicate(): + continue + # Compute the logical shape seen at this step: dims to the left + # (not yet gathered) still divide their tensor dimension. + logical_shape = list(tensor.shape) + for j in range(i): + pj = placements[j] + if not pj.is_replicate(): + logical_shape[pj.dim] //= mesh.size(j) + local = p._to_replicate_tensor(local, mesh, i, logical_shape) + return DTensor.from_local(local, mesh, replicate_all, run_check=False) + + return tensor.redistribute(placements=replicate_all) def convert_strided_to_shard(state_dict: dict) -> dict[str, tuple]: diff --git a/tests/test_tensor_parallel_mixin.py b/tests/test_tensor_parallel_mixin.py index f4a9e54a8691..e07d2beb539b 100644 --- a/tests/test_tensor_parallel_mixin.py +++ b/tests/test_tensor_parallel_mixin.py @@ -48,9 +48,7 @@ def _to_local(tensor): # Partial gradients for us. if isinstance(tensor, DTensor) and any(not p.is_replicate() for p in tensor.placements): - from transformers.integrations.tensor_parallel import _replicate_dtensor - - tensor = _replicate_dtensor(tensor) + tensor = tensor.redistribute(placements=(Replicate(),)) return tensor.to_local() return tensor From 44706eb1572decaa9282d1c06df6080e3747ea66 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Thu, 16 Apr 2026 11:06:15 +0000 Subject: [PATCH 023/116] revert small change --- tests/test_tensor_parallel_mixin.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_tensor_parallel_mixin.py b/tests/test_tensor_parallel_mixin.py index e07d2beb539b..159446b6e8f6 100644 --- a/tests/test_tensor_parallel_mixin.py +++ b/tests/test_tensor_parallel_mixin.py @@ -23,6 +23,7 @@ is_torch_available, ) from transformers.utils import is_torch_greater_or_equal, is_torchao_available +from transformers.integrations.tensor_parallel import _replicate_dtensor if is_torchao_available(): @@ -48,7 +49,7 @@ def _to_local(tensor): # Partial gradients for us. if isinstance(tensor, DTensor) and any(not p.is_replicate() for p in tensor.placements): - tensor = tensor.redistribute(placements=(Replicate(),)) + tensor = _replicate_dtensor(tensor) return tensor.to_local() return tensor From aa45f5ba2582d4dc1e4eef8346c78d44cc286edd Mon Sep 17 00:00:00 2001 From: 3outeille Date: Thu, 16 Apr 2026 12:30:56 +0000 Subject: [PATCH 024/116] push temporary fix for TP and strided shard for backward --- src/transformers/core_model_loading.py | 5 +- .../integrations/tensor_parallel.py | 110 +++++++++++------- tests/test_tensor_parallel_mixin.py | 5 +- 3 files changed, 75 insertions(+), 45 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 8b1be399400e..bd33512a0296 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -885,7 +885,6 @@ def shard_tensor( if not sharding_placements: return param[...].to(device=device, dtype=dtype) - if tensor_idx is not None and len(self.param.shape) == len(param_shape) + 1: # [B] Expert path: shard on expert dimension (dim 0). # When dim 0 is the only sharding placement, return the full expert or @@ -902,13 +901,13 @@ def shard_tensor( return param[...].to(device=device, dtype=dtype) # [B4] Composed with TP placements -> shard the expert's inner dims return self._shard_nd(param, inner_placements, param_shape, device, dtype) - + # [B1] has_expert_sharding=False -> fall through to _shard_nd return self._shard_nd(param, sharding_placements, param_shape, device, dtype) def _shard_nd(self, param, sharding_placements, param_shape, device, dtype): """Handle multi-dimensional sharding, choosing the best strategy.""" - # [C1] Column Parallel when composed with FSDP. We choose the easier path but maybe we should do a better one? + # [C1] Column Parallel when composed with FSDP. We choose the easier path but maybe we should do a better one? if not self._can_shard_on_read(sharding_placements): return self._materialize_and_split(param, sharding_placements, device, dtype) diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index 6a2a043c8b57..636d8ab15558 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -280,6 +280,70 @@ def output_hook(mod, inputs, output): return module +def _accumulate_local_param_grad(original_param: DTensor, local_grad: torch.Tensor) -> torch.Tensor: + """Copy a detached local grad into the original DTensor parameter. + + Packed ``_StridedShard`` parameters cannot rely on autograd through + ``DTensor.to_local()`` on older torch releases, so we materialize a local leaf + parameter for the forward and stitch its gradient back manually here. + """ + tensor_meta = original_param._spec.tensor_meta + detached_grad = local_grad.detach() + grad_dtensor = DTensor.from_local( + detached_grad, + original_param.device_mesh, + original_param.placements, + run_check=False, + shape=tensor_meta.shape, + stride=tensor_meta.stride, + ) + + with torch.no_grad(): + existing_grad = original_param.grad + if existing_grad is None: + original_param.grad = grad_dtensor + elif isinstance(existing_grad, DTensor): + existing_grad._local_tensor.add_(detached_grad) + else: + existing_grad.add_(detached_grad) + + return local_grad + + +def _swap_dtensor_params_for_local(module, shadow_attr: str) -> None: + """Temporarily replace DTensor params by detached local leaf params.""" + local_param_shadows = {} + for param_name, param in list(module.named_parameters(recurse=False)): + if not isinstance(param, DTensor): + continue + + local_param_shadows[param_name] = param + local_param = torch.nn.Parameter(param._local_tensor.detach(), requires_grad=param.requires_grad) + if param.requires_grad: + local_param.register_hook( + lambda grad, original_param=param: _accumulate_local_param_grad(original_param, grad) + ) + + module._parameters.pop(param_name) + setattr(module, param_name, local_param) + + if local_param_shadows: + shadow_stack = getattr(module, shadow_attr, None) + if shadow_stack is None: + shadow_stack = [] + setattr(module, shadow_attr, shadow_stack) + shadow_stack.append(local_param_shadows) + + +def _restore_dtensor_params(module, shadow_attr: str) -> None: + shadow_stack = getattr(module, shadow_attr, None) + if shadow_stack: + for param_name, param in shadow_stack.pop().items(): + if hasattr(module, param_name): + delattr(module, param_name) + module.register_parameter(param_name, param) + + class PackedColwiseParallel(ParallelStyle): """Column-wise parallel style for fused linear weights packed along the output dimension.""" @@ -325,27 +389,11 @@ def _prepare_input_fn(self, mod, inputs, device_mesh): input_tensor = input_tensor.redistribute(placements=self.input_layouts) input_tensor = input_tensor.to_local() - local_param_shadows = {} - for param_name, param in list(mod.named_parameters(recurse=False)): - if isinstance(param, DTensor): - local_param_shadows[param_name] = param - mod._parameters.pop(param_name) - setattr(mod, param_name, param.to_local()) - if local_param_shadows: - shadow_stack = getattr(mod, "_packed_local_param_shadows", None) - if shadow_stack is None: - shadow_stack = [] - mod._packed_local_param_shadows = shadow_stack - shadow_stack.append(local_param_shadows) + _swap_dtensor_params_for_local(mod, "_packed_local_param_shadows") return (input_tensor,) + inputs[1:] def _prepare_output_fn(self, mod, outputs, device_mesh): - shadow_stack = getattr(mod, "_packed_local_param_shadows", None) - if shadow_stack: - for param_name, param in shadow_stack.pop().items(): - if hasattr(mod, param_name): - delattr(mod, param_name) - mod.register_parameter(param_name, param) + _restore_dtensor_params(mod, "_packed_local_param_shadows") if outputs is None or self.use_local_output: return outputs @@ -472,31 +520,15 @@ def _prepare_input_fn(mod, inputs, device_mesh): if isinstance(top_k_weights, DTensor): top_k_weights = top_k_weights.to_local() top_k_weights = _AllReduceBackward.apply(top_k_weights, tp_group) - local_param_shadows = {} - for param_name, param in list(mod.named_parameters(recurse=False)): - if isinstance(param, DTensor): - # grouped_mm expects plain tensors, but we must restore the - # original DTensor params after the forward so save_pretrained - # still sees the canonical sharded weights. - local_param_shadows[param_name] = param - mod._parameters.pop(param_name) - setattr(mod, param_name, param.to_local()) - if local_param_shadows: - shadow_stack = getattr(mod, "_moe_local_param_shadows", None) - if shadow_stack is None: - shadow_stack = [] - mod._moe_local_param_shadows = shadow_stack - shadow_stack.append(local_param_shadows) + # grouped_mm expects plain tensors, but we must restore the original + # DTensor params after the forward so save_pretrained still sees the + # canonical sharded weights. + _swap_dtensor_params_for_local(mod, "_moe_local_param_shadows") return (hidden_states, top_k_index, top_k_weights) @staticmethod def _prepare_output_fn(output_layouts, mod, outputs, device_mesh): - shadow_stack = getattr(mod, "_moe_local_param_shadows", None) - if shadow_stack: - for param_name, param in shadow_stack.pop().items(): - if hasattr(mod, param_name): - delattr(mod, param_name) - mod.register_parameter(param_name, param) + _restore_dtensor_params(mod, "_moe_local_param_shadows") if outputs is None: return None # Plain TP expert weights produce partial outputs that need an all-reduce. diff --git a/tests/test_tensor_parallel_mixin.py b/tests/test_tensor_parallel_mixin.py index 159446b6e8f6..5d102d2ff0ec 100644 --- a/tests/test_tensor_parallel_mixin.py +++ b/tests/test_tensor_parallel_mixin.py @@ -17,13 +17,12 @@ from transformers import TorchAoConfig, set_seed from transformers.distributed import DistributedConfig -from transformers.integrations.tensor_parallel import _get_parameter_tp_plan +from transformers.integrations.tensor_parallel import _get_parameter_tp_plan, _replicate_dtensor from transformers.testing_utils import ( is_tensor_parallel_test, is_torch_available, ) from transformers.utils import is_torch_greater_or_equal, is_torchao_available -from transformers.integrations.tensor_parallel import _replicate_dtensor if is_torchao_available(): @@ -34,7 +33,7 @@ import torch import torch.distributed as dist import torch.multiprocessing as mp - from torch.distributed.tensor import DTensor, Replicate + from torch.distributed.tensor import DTensor from torch.multiprocessing.spawn import ProcessRaisedException From 0a566c5250d82afcea59993a5a50fa57e87dca9e Mon Sep 17 00:00:00 2001 From: 3outeille Date: Thu, 16 Apr 2026 13:27:48 +0000 Subject: [PATCH 025/116] refactor a bit --- .../integrations/tensor_parallel.py | 44 ++++++++++--------- 1 file changed, 23 insertions(+), 21 deletions(-) diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index 636d8ab15558..36a4bf2a1860 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -310,14 +310,17 @@ def _accumulate_local_param_grad(original_param: DTensor, local_grad: torch.Tens return local_grad -def _swap_dtensor_params_for_local(module, shadow_attr: str) -> None: - """Temporarily replace DTensor params by detached local leaf params.""" - local_param_shadows = {} +def _materialize_local_params(module, shadow_attr: str) -> None: + """Swap DTensor params for detached local leaf params during one forward.""" + if getattr(module, shadow_attr, None) is not None: + raise RuntimeError(f"{module.__class__.__name__} already has active local parameter shadows") + + param_shadows = {} for param_name, param in list(module.named_parameters(recurse=False)): if not isinstance(param, DTensor): continue - local_param_shadows[param_name] = param + param_shadows[param_name] = param local_param = torch.nn.Parameter(param._local_tensor.detach(), requires_grad=param.requires_grad) if param.requires_grad: local_param.register_hook( @@ -327,21 +330,20 @@ def _swap_dtensor_params_for_local(module, shadow_attr: str) -> None: module._parameters.pop(param_name) setattr(module, param_name, local_param) - if local_param_shadows: - shadow_stack = getattr(module, shadow_attr, None) - if shadow_stack is None: - shadow_stack = [] - setattr(module, shadow_attr, shadow_stack) - shadow_stack.append(local_param_shadows) + if param_shadows: + setattr(module, shadow_attr, param_shadows) + +def _restore_local_params(module, shadow_attr: str) -> None: + param_shadows = getattr(module, shadow_attr, None) + if param_shadows is None: + return -def _restore_dtensor_params(module, shadow_attr: str) -> None: - shadow_stack = getattr(module, shadow_attr, None) - if shadow_stack: - for param_name, param in shadow_stack.pop().items(): - if hasattr(module, param_name): - delattr(module, param_name) - module.register_parameter(param_name, param) + delattr(module, shadow_attr) + for param_name, param in param_shadows.items(): + if hasattr(module, param_name): + delattr(module, param_name) + module.register_parameter(param_name, param) class PackedColwiseParallel(ParallelStyle): @@ -389,11 +391,11 @@ def _prepare_input_fn(self, mod, inputs, device_mesh): input_tensor = input_tensor.redistribute(placements=self.input_layouts) input_tensor = input_tensor.to_local() - _swap_dtensor_params_for_local(mod, "_packed_local_param_shadows") + _materialize_local_params(mod, "_packed_local_params") return (input_tensor,) + inputs[1:] def _prepare_output_fn(self, mod, outputs, device_mesh): - _restore_dtensor_params(mod, "_packed_local_param_shadows") + _restore_local_params(mod, "_packed_local_params") if outputs is None or self.use_local_output: return outputs @@ -523,12 +525,12 @@ def _prepare_input_fn(mod, inputs, device_mesh): # grouped_mm expects plain tensors, but we must restore the original # DTensor params after the forward so save_pretrained still sees the # canonical sharded weights. - _swap_dtensor_params_for_local(mod, "_moe_local_param_shadows") + _materialize_local_params(mod, "_moe_local_params") return (hidden_states, top_k_index, top_k_weights) @staticmethod def _prepare_output_fn(output_layouts, mod, outputs, device_mesh): - _restore_dtensor_params(mod, "_moe_local_param_shadows") + _restore_local_params(mod, "_moe_local_params") if outputs is None: return None # Plain TP expert weights produce partial outputs that need an all-reduce. From 11a55d4614bb49b4c9f712f5325e2bd8d7dc0e10 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Thu, 16 Apr 2026 13:42:27 +0000 Subject: [PATCH 026/116] patches for rotary --- src/transformers/distributed/patches.py | 41 +++++++++++++++++++------ 1 file changed, 32 insertions(+), 9 deletions(-) diff --git a/src/transformers/distributed/patches.py b/src/transformers/distributed/patches.py index 984498ef680a..17d268ca837f 100644 --- a/src/transformers/distributed/patches.py +++ b/src/transformers/distributed/patches.py @@ -23,21 +23,44 @@ import sys from functools import wraps +import inspect + from torch.distributed.tensor import DTensor, Replicate def _make_dtensor_rotary_wrapper(original_fn): - """Return a wrapper that converts cos/sin to replicated DTensors when q is a DTensor.""" + """Return a wrapper that promotes cos/sin to replicated DTensors. + + Models use two ``apply_rotary_pos_emb`` signatures: + - ``(q, k, cos, sin, ...)`` — most models + - ``(x, cos, sin, ...)`` — gemma3n, gemma4, glm_moe_dsa + + We detect which one at patch time via parameter count and create + the matching wrapper. + """ + params = inspect.signature(original_fn).parameters + n_required = sum(1 for p in params.values() if p.default is inspect.Parameter.empty) + + if n_required >= 4: + + @wraps(original_fn) + def _wrapper(q, k, cos, sin, *args, **kwargs): + if isinstance(q, DTensor) and not isinstance(cos, DTensor): + replicate = (Replicate(),) * q.device_mesh.ndim + cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) + return original_fn(q, k, cos, sin, *args, **kwargs) + else: - @wraps(original_fn) - def _dtensor_apply_rotary_pos_emb(q, k, cos, sin, *args, **kwargs): - if isinstance(q, DTensor) and not isinstance(cos, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) - return original_fn(q, k, cos, sin, *args, **kwargs) + @wraps(original_fn) + def _wrapper(x, cos, sin, *args, **kwargs): + if isinstance(x, DTensor) and not isinstance(cos, DTensor): + replicate = (Replicate(),) * x.device_mesh.ndim + cos = DTensor.from_local(cos, x.device_mesh, replicate, run_check=False) + sin = DTensor.from_local(sin, x.device_mesh, replicate, run_check=False) + return original_fn(x, cos, sin, *args, **kwargs) - return _dtensor_apply_rotary_pos_emb + return _wrapper def patch_dtensor_ops(model): From 53490d98611bf0c0216c431c8770dbcef17bd067 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Thu, 16 Apr 2026 14:41:05 +0000 Subject: [PATCH 027/116] refactor MoEExpertsParallel --- src/transformers/distributed/patches.py | 3 +- .../integrations/tensor_parallel.py | 139 ++++++++---------- 2 files changed, 60 insertions(+), 82 deletions(-) diff --git a/src/transformers/distributed/patches.py b/src/transformers/distributed/patches.py index 17d268ca837f..3ffe15a37afe 100644 --- a/src/transformers/distributed/patches.py +++ b/src/transformers/distributed/patches.py @@ -20,11 +20,10 @@ from __future__ import annotations +import inspect import sys from functools import wraps -import inspect - from torch.distributed.tensor import DTensor, Replicate diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index 36a4bf2a1860..306b57863414 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -454,12 +454,20 @@ def backward(ctx, grad): class MoEExpertsParallel(ParallelStyle): - """Hybrid parallel style for MoE expert modules. - - Converts expert weights to DTensors based on the ``shard_plan`` (e.g. - ``{"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}``). - Communication uses DTensor ``from_local``/``to_local`` on activations only — - compatible with ``grouped_mm``. + """Tensor-parallel style for MoE expert modules. + + Shards expert weights as DTensors, then wraps the module's ``forward`` so + that grouped_mm (which needs plain tensors) works transparently. + + The wrapped forward does four things: + 1. Localize inputs — wrap hidden_states as Replicate DTensor then extract + local tensor (gives us an all-reduce on the backward gradient for free). + 2. Fix routing grads — routing weights are the same on all ranks, but their + backward gradient is partial; use allreduce-sum (not divide-by-world-size). + 3. Swap params — temporarily replace DTensor params with local tensors + for grouped_mm, restore them after so save_pretrained sees DTensors. + 4. Reduce output — each rank's output is partial (only its expert shard + contributed); all-reduce to get the complete hidden state. """ def __init__(self, output_layouts=None): @@ -480,84 +488,55 @@ def _partition_fn(name, module, device_mesh, shard_plan): dtensor = distribute_tensor(param.data, device_mesh, [placement]) module._parameters[param_name] = torch.nn.Parameter(dtensor, requires_grad=param.requires_grad) - @staticmethod - def _uses_partial_outputs(mod) -> bool: - cached = getattr(mod, "_moe_outputs_are_partial", None) - if cached is not None: - return cached - - # Under TP-only the expert MLP dimension is sharded, so each rank emits a - # partial hidden-state contribution that must be reduced. Under TP+FSDP, - # FSDP can swap in full gathered expert weights for the current rank's - # forward, in which case the local output is already complete. - intermediate = getattr(mod, "intermediate_dim", None) or getattr(mod, "intermediate_size", None) - if hasattr(mod, "gate_up_proj"): - gate_up_proj = mod.gate_up_proj.to_local() if isinstance(mod.gate_up_proj, DTensor) else mod.gate_up_proj - full_expert_out = 2 * intermediate - sharded_dim = -1 if getattr(mod, "is_transposed", False) else -2 - cached = gate_up_proj.shape[sharded_dim] != full_expert_out - elif hasattr(mod, "up_proj"): - up_proj = mod.up_proj.to_local() if isinstance(mod.up_proj, DTensor) else mod.up_proj - full_expert_out = intermediate - sharded_dim = -1 if getattr(mod, "is_transposed", False) else -2 - cached = up_proj.shape[sharded_dim] != full_expert_out - else: - cached = True - - mod._moe_outputs_are_partial = cached - return cached + def _apply(self, module, device_mesh): + self._partition_fn(module.__class__.__name__, module, device_mesh, self._moe_shard_plan) - @staticmethod - def _prepare_input_fn(mod, inputs, device_mesh): - hidden_states, top_k_index, top_k_weights = inputs[0], inputs[1], inputs[2] - # from_local([Replicate()]).to_local(): forward sees plain tensor, - # backward graph goes through DTensor all-reduce on gradient. - if not isinstance(hidden_states, DTensor): - hidden_states = DTensor.from_local(hidden_states, device_mesh, [Replicate()], run_check=False) - hidden_states = hidden_states.to_local() - # Route weights are replicated (same on all ranks), but their backward - # gradient is partial (each rank's contribution from its expert shard). - # Use allreduce-sum (not Replicate's allreduce-then-divide) to aggregate. + output_layouts = self.output_layouts + original_forward = module.forward tp_group = device_mesh.get_group() if device_mesh.ndim == 1 else device_mesh.get_group("tp") - if isinstance(top_k_weights, DTensor): - top_k_weights = top_k_weights.to_local() - top_k_weights = _AllReduceBackward.apply(top_k_weights, tp_group) - # grouped_mm expects plain tensors, but we must restore the original - # DTensor params after the forward so save_pretrained still sees the - # canonical sharded weights. - _materialize_local_params(mod, "_moe_local_params") - return (hidden_states, top_k_index, top_k_weights) - @staticmethod - def _prepare_output_fn(output_layouts, mod, outputs, device_mesh): - _restore_local_params(mod, "_moe_local_params") - if outputs is None: - return None - # Plain TP expert weights produce partial outputs that need an all-reduce. - # TP+FSDP can leave experts replicated across TP and sharded only across - # experts/FSDP, in which case the local output is already complete. - source_layout = Partial() if MoEExpertsParallel._uses_partial_outputs(mod) else Replicate() - if not isinstance(outputs, DTensor): - outputs = DTensor.from_local(outputs, device_mesh, [source_layout], run_check=False) - # MoE experts output 2D [num_tokens, hidden]. For SP reduce-scatter, - # Shard(1) means sequence dim in 3D, but in 2D the token dim is 0. - actual_layouts = output_layouts - if outputs.dim() == 2 and isinstance(output_layouts, Shard) and output_layouts.dim == 1: - actual_layouts = Shard(0) - if outputs.placements != (actual_layouts,): - outputs = outputs.redistribute(placements=(actual_layouts,)) - return outputs.to_local() + def tp_forward(hidden_states, top_k_index, top_k_weights): + # --- 1. Localize hidden_states (backward gets all-reduce for free) --- + if not isinstance(hidden_states, DTensor): + hidden_states = DTensor.from_local(hidden_states, device_mesh, [Replicate()], run_check=False) + hidden_states = hidden_states.to_local() + + # --- 2. Fix routing weight gradients (allreduce-sum, not ÷ world_size) --- + if isinstance(top_k_weights, DTensor): + top_k_weights = top_k_weights.to_local() + top_k_weights = _AllReduceBackward.apply(top_k_weights, tp_group) + + # --- 3. Swap DTensor params → local for grouped_mm --- + _materialize_local_params(module, "_moe_local_params") + + # --- 4. Run the original forward --- + output = original_forward(hidden_states, top_k_index, top_k_weights) + + # --- 5. Restore DTensor params (so save_pretrained sees them) --- + _restore_local_params(module, "_moe_local_params") + + # --- 6. Reduce partial output --- + 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, device_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 = 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() - def _apply(self, module, device_mesh): - # Don't use PyTorch's distribute_module — it would auto-convert all - # params to Replicate DTensors. We create DTensors with proper Shard - # placements in _partition_fn instead, and register hooks manually. - self._partition_fn(module.__class__.__name__, module, device_mesh, self._moe_shard_plan) - module.register_forward_pre_hook(lambda mod, inputs: self._prepare_input_fn(mod, inputs, device_mesh)) - module.register_forward_hook( - lambda mod, inputs, outputs: self._prepare_output_fn(self.output_layouts, mod, outputs, device_mesh), - always_call=True, - ) + module.forward = tp_forward return module From 0c09915546f1c694441cdaa6e4716171161533da Mon Sep 17 00:00:00 2001 From: 3outeille Date: Sat, 18 Apr 2026 10:25:36 +0000 Subject: [PATCH 028/116] fix tp for last models --- src/transformers/integrations/tensor_parallel.py | 2 +- .../models/flex_olmo/configuration_flex_olmo.py | 7 +------ .../models/flex_olmo/modular_flex_olmo.py | 7 +------ .../models/gemma3n/configuration_gemma3n.py | 8 ++++---- .../models/gemma3n/modular_gemma3n.py | 8 ++++---- .../models/gemma4/configuration_gemma4.py | 15 ++++++--------- .../longcat_flash/configuration_longcat_flash.py | 5 ----- 7 files changed, 17 insertions(+), 35 deletions(-) diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index 306b57863414..4ffe66f91f80 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -496,7 +496,7 @@ def _apply(self, module, device_mesh): tp_group = device_mesh.get_group() if device_mesh.ndim == 1 else device_mesh.get_group("tp") def tp_forward(hidden_states, top_k_index, top_k_weights): - # --- 1. Localize hidden_states (backward gets all-reduce for free) --- + # --- 1. Localize hidden_states (backward all-reduce via DTensor) --- if not isinstance(hidden_states, DTensor): hidden_states = DTensor.from_local(hidden_states, device_mesh, [Replicate()], run_check=False) hidden_states = hidden_states.to_local() diff --git a/src/transformers/models/flex_olmo/configuration_flex_olmo.py b/src/transformers/models/flex_olmo/configuration_flex_olmo.py index 9f8991d84000..f31dd4d14674 100644 --- a/src/transformers/models/flex_olmo/configuration_flex_olmo.py +++ b/src/transformers/models/flex_olmo/configuration_flex_olmo.py @@ -62,12 +62,7 @@ class FlexOlmoConfig(PreTrainedConfig): ), # we need to replicate here due to the added norm on q and k "layers.*.self_attn.o_proj": TPStyle( "vocab", "allreduce" - ), # input is replicated due to the added norm on q and k - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), + ) # input is replicated due to the added norm on q and k } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/flex_olmo/modular_flex_olmo.py b/src/transformers/models/flex_olmo/modular_flex_olmo.py index 8a496547c204..76b76e256734 100644 --- a/src/transformers/models/flex_olmo/modular_flex_olmo.py +++ b/src/transformers/models/flex_olmo/modular_flex_olmo.py @@ -72,12 +72,7 @@ class FlexOlmoConfig(PreTrainedConfig): ), # we need to replicate here due to the added norm on q and k "layers.*.self_attn.o_proj": TPStyle( "vocab", "allreduce" - ), # input is replicated due to the added norm on q and k - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), + ) # input is replicated due to the added norm on q and k } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/gemma3n/configuration_gemma3n.py b/src/transformers/models/gemma3n/configuration_gemma3n.py index 21627324d286..950f0ea78b47 100644 --- a/src/transformers/models/gemma3n/configuration_gemma3n.py +++ b/src/transformers/models/gemma3n/configuration_gemma3n.py @@ -80,10 +80,10 @@ class Gemma3nTextConfig(PreTrainedConfig): model_type = "gemma3n_text" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "allgather"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "allgather"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "allgather"), + "layers.*.self_attn.o_proj": TPStyle("vocab", "allreduce"), "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), diff --git a/src/transformers/models/gemma3n/modular_gemma3n.py b/src/transformers/models/gemma3n/modular_gemma3n.py index 41364bd92f87..53a2b066d1d7 100644 --- a/src/transformers/models/gemma3n/modular_gemma3n.py +++ b/src/transformers/models/gemma3n/modular_gemma3n.py @@ -107,10 +107,10 @@ class Gemma3nTextConfig(Gemma3TextConfig): model_type = "gemma3n_text" base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "allgather"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "allgather"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "allgather"), + "layers.*.self_attn.o_proj": TPStyle("vocab", "allreduce"), "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), diff --git a/src/transformers/models/gemma4/configuration_gemma4.py b/src/transformers/models/gemma4/configuration_gemma4.py index 9be69c3c0860..4bb3b142eb77 100644 --- a/src/transformers/models/gemma4/configuration_gemma4.py +++ b/src/transformers/models/gemma4/configuration_gemma4.py @@ -124,18 +124,15 @@ class Gemma4TextConfig(PreTrainedConfig): model_type = "gemma4_text" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), + # 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": TPStyle("colwise", "allgather"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "allgather"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "allgather"), + "layers.*.self_attn.o_proj": TPStyle("vocab", "allreduce"), "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), - "layers.*.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/longcat_flash/configuration_longcat_flash.py b/src/transformers/models/longcat_flash/configuration_longcat_flash.py index 8002cb36b00f..3ad063190541 100644 --- a/src/transformers/models/longcat_flash/configuration_longcat_flash.py +++ b/src/transformers/models/longcat_flash/configuration_longcat_flash.py @@ -58,11 +58,6 @@ class LongcatFlashConfig(PreTrainedConfig): "layers.*.self_attn.*.q_b_proj": TPStyle("colwise", "none"), "layers.*.self_attn.*.kv_b_proj": TPStyle("colwise", "none"), "layers.*.self_attn.*.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), "layers.*.mlps.*.gate_proj": TPStyle("colwise", "none"), "layers.*.mlps.*.up_proj": TPStyle("colwise", "none"), "layers.*.mlps.*.down_proj": TPStyle("rowwise", "allreduce"), From ebd03ecd724dbc97a61bc0482161091d0959fe4d Mon Sep 17 00:00:00 2001 From: 3outeille Date: Sat, 18 Apr 2026 10:50:38 +0000 Subject: [PATCH 029/116] refactor moe expert parallels --- .../integrations/tensor_parallel.py | 113 ++++++++---------- 1 file changed, 50 insertions(+), 63 deletions(-) diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index 4ffe66f91f80..67b3d748c08c 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -13,6 +13,7 @@ # limitations under the License. from __future__ import annotations +import contextlib import re from dataclasses import dataclass from typing import Literal @@ -310,40 +311,33 @@ def _accumulate_local_param_grad(original_param: DTensor, local_grad: torch.Tens return local_grad -def _materialize_local_params(module, shadow_attr: str) -> None: - """Swap DTensor params for detached local leaf params during one forward.""" - if getattr(module, shadow_attr, None) is not None: - raise RuntimeError(f"{module.__class__.__name__} already has active local parameter shadows") +@contextlib.contextmanager +def _local_dtensor_params(module): + """Temporarily swap DTensor params for local leaf params during one forward. - param_shadows = {} - for param_name, param in list(module.named_parameters(recurse=False)): + Needed because grouped_mm / fused ops on DTensors trigger broken autograd + paths for ``_StridedShard``. We run forward on a plain-tensor leaf param and + stitch its gradient back onto the original DTensor via a hook. Restores the + DTensor params on exit (even on exception). + """ + shadows = {} + for name, param in list(module.named_parameters(recurse=False)): if not isinstance(param, DTensor): continue - - param_shadows[param_name] = param - local_param = torch.nn.Parameter(param._local_tensor.detach(), requires_grad=param.requires_grad) + shadows[name] = param + local = torch.nn.Parameter(param._local_tensor.detach(), requires_grad=param.requires_grad) if param.requires_grad: - local_param.register_hook( - lambda grad, original_param=param: _accumulate_local_param_grad(original_param, grad) - ) - - module._parameters.pop(param_name) - setattr(module, param_name, local_param) - - if param_shadows: - setattr(module, shadow_attr, param_shadows) - + local.register_hook(lambda g, p=param: _accumulate_local_param_grad(p, g)) + module._parameters.pop(name) + setattr(module, name, local) -def _restore_local_params(module, shadow_attr: str) -> None: - param_shadows = getattr(module, shadow_attr, None) - if param_shadows is None: - return - - delattr(module, shadow_attr) - for param_name, param in param_shadows.items(): - if hasattr(module, param_name): - delattr(module, param_name) - module.register_parameter(param_name, param) + try: + yield + finally: + for name, param in shadows.items(): + if hasattr(module, name): + delattr(module, name) + module.register_parameter(name, param) class PackedColwiseParallel(ParallelStyle): @@ -383,36 +377,34 @@ def _partition_linear_fn(self, module, device_mesh): ), ) - def _prepare_input_fn(self, mod, inputs, device_mesh): - input_tensor = inputs[0] - if not isinstance(input_tensor, DTensor): - input_tensor = DTensor.from_local(input_tensor, device_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() - - _materialize_local_params(mod, "_packed_local_params") - return (input_tensor,) + inputs[1:] - - def _prepare_output_fn(self, mod, outputs, device_mesh): - _restore_local_params(mod, "_packed_local_params") - - if outputs is None or self.use_local_output: - return outputs - return DTensor.from_local( - outputs, device_mesh, (_StridedShard(dim=-1, split_factor=self.split_factor),), run_check=False - ) - def _apply(self, module, device_mesh): if not isinstance(module, torch.nn.Linear): raise NotImplementedError("PackedColwiseParallel currently only supports nn.Linear!") self._partition_linear_fn(module, device_mesh) - module.register_forward_pre_hook(lambda mod, inputs: self._prepare_input_fn(mod, inputs, device_mesh)) - module.register_forward_hook( - lambda mod, inputs, outputs: self._prepare_output_fn(mod, outputs, device_mesh), - always_call=True, - ) + + input_layouts = self.input_layouts + use_local_output = self.use_local_output + split_factor = self.split_factor + original_forward = module.forward + + def tp_forward(input_tensor, *args, **kwargs): + if not isinstance(input_tensor, DTensor): + input_tensor = DTensor.from_local(input_tensor, device_mesh, input_layouts, run_check=False) + elif input_tensor.placements != input_layouts: + input_tensor = input_tensor.redistribute(placements=input_layouts) + input_tensor = input_tensor.to_local() + + with _local_dtensor_params(module): + output = original_forward(input_tensor, *args, **kwargs) + + if output is None or use_local_output: + return output + return DTensor.from_local( + output, device_mesh, (_StridedShard(dim=-1, split_factor=split_factor),), run_check=False + ) + + module.forward = tp_forward return module def __repr__(self) -> str: @@ -506,16 +498,11 @@ def tp_forward(hidden_states, top_k_index, top_k_weights): top_k_weights = top_k_weights.to_local() top_k_weights = _AllReduceBackward.apply(top_k_weights, tp_group) - # --- 3. Swap DTensor params → local for grouped_mm --- - _materialize_local_params(module, "_moe_local_params") - - # --- 4. Run the original forward --- - output = original_forward(hidden_states, top_k_index, top_k_weights) - - # --- 5. Restore DTensor params (so save_pretrained sees them) --- - _restore_local_params(module, "_moe_local_params") + # --- 3. Run forward with local params (grouped_mm needs plain tensors) --- + with _local_dtensor_params(module): + output = original_forward(hidden_states, top_k_index, top_k_weights) - # --- 6. Reduce partial output --- + # --- 4. Reduce partial output --- if output is None: return None # Under TP-only each rank has a partial result; under TP+FSDP the From c08c071410b6c141dc9d8887fa365247c8330df2 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Sat, 18 Apr 2026 13:05:49 +0000 Subject: [PATCH 030/116] linting --- src/transformers/models/flex_olmo/configuration_flex_olmo.py | 2 +- src/transformers/models/flex_olmo/modular_flex_olmo.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/transformers/models/flex_olmo/configuration_flex_olmo.py b/src/transformers/models/flex_olmo/configuration_flex_olmo.py index f31dd4d14674..fb45082fd430 100644 --- a/src/transformers/models/flex_olmo/configuration_flex_olmo.py +++ b/src/transformers/models/flex_olmo/configuration_flex_olmo.py @@ -62,7 +62,7 @@ class FlexOlmoConfig(PreTrainedConfig): ), # we need to replicate here due to the added norm on q and k "layers.*.self_attn.o_proj": TPStyle( "vocab", "allreduce" - ) # input is replicated due to the added norm on q and k + ), # input is replicated due to the added norm on q and k } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/flex_olmo/modular_flex_olmo.py b/src/transformers/models/flex_olmo/modular_flex_olmo.py index 76b76e256734..68e21e53fa30 100644 --- a/src/transformers/models/flex_olmo/modular_flex_olmo.py +++ b/src/transformers/models/flex_olmo/modular_flex_olmo.py @@ -72,7 +72,7 @@ class FlexOlmoConfig(PreTrainedConfig): ), # we need to replicate here due to the added norm on q and k "layers.*.self_attn.o_proj": TPStyle( "vocab", "allreduce" - ) # input is replicated due to the added norm on q and k + ), # input is replicated due to the added norm on q and k } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), From 4804d0d21669b24a8797cdd1905ea5092964aa4a Mon Sep 17 00:00:00 2001 From: 3outeille Date: Sat, 18 Apr 2026 14:32:29 +0000 Subject: [PATCH 031/116] add sp plan for models --- .gitignore | 5 +++ run_compare.sh | 2 +- run_verify_all.sh | 2 +- .../integrations/tensor_parallel.py | 5 +-- .../models/apertus/configuration_apertus.py | 18 ++++++++++ .../models/apertus/modular_apertus.py | 18 ++++++++++ .../models/arcee/configuration_arcee.py | 15 ++++++++ .../models/arcee/modular_arcee.py | 15 ++++++++ .../models/aria/configuration_aria.py | 16 +++++++++ .../models/cohere/configuration_cohere.py | 18 ++++++++++ .../models/cohere2/configuration_cohere2.py | 16 +++++++++ .../models/cohere2/modular_cohere2.py | 16 +++++++++ .../models/cwm/configuration_cwm.py | 16 +++++++++ .../deepseek_v2/configuration_deepseek_v2.py | 1 + .../models/deepseek_v2/modular_deepseek_v2.py | 2 ++ .../models/ernie4_5/configuration_ernie4_5.py | 16 +++++++++ .../models/exaone4/configuration_exaone4.py | 18 ++++++++++ .../models/exaone4/modular_exaone4.py | 18 ++++++++++ .../exaone_moe/configuration_exaone_moe.py | 18 ++++++++++ .../models/gemma/configuration_gemma.py | 16 +++++++++ .../models/gemma/modular_gemma.py | 16 +++++++++ .../models/gemma2/configuration_gemma2.py | 16 +++++++++ .../models/gemma2/modular_gemma2.py | 16 +++++++++ .../models/gemma3/configuration_gemma3.py | 18 ++++++++++ .../models/gemma3/modular_gemma3.py | 18 ++++++++++ .../models/gemma3n/configuration_gemma3n.py | 18 ++++++++++ .../models/gemma3n/modular_gemma3n.py | 18 ++++++++++ .../models/glm/configuration_glm.py | 15 ++++++++ .../models/glm4/configuration_glm4.py | 15 ++++++++ .../models/gpt_neox/configuration_gpt_neox.py | 10 ++++++ .../models/granite/configuration_granite.py | 16 +++++++++ .../models/helium/configuration_helium.py | 16 +++++++++ .../configuration_higgs_audio_v2.py | 16 +++++++++ .../models/jais2/configuration_jais2.py | 15 ++++++++ .../models/jais2/modular_jais2.py | 15 ++++++++ .../models/llama/configuration_llama.py | 16 +++++++++ .../minimax_m2/configuration_minimax_m2.py | 18 ++++++++++ .../models/minimax_m2/modular_minimax_m2.py | 18 ++++++++++ .../ministral/configuration_ministral.py | 16 +++++++++ .../ministral3/configuration_ministral3.py | 16 +++++++++ .../models/mistral/configuration_mistral.py | 16 +++++++++ .../models/nanochat/configuration_nanochat.py | 18 ++++++++++ .../models/olmo/configuration_olmo.py | 16 +++++++++ .../models/olmo2/configuration_olmo2.py | 18 ++++++++++ .../models/olmo2/modular_olmo2.py | 18 ++++++++++ .../models/olmo3/configuration_olmo3.py | 18 ++++++++++ .../models/olmo3/modular_olmo3.py | 18 ++++++++++ .../olmo_hybrid/configuration_olmo_hybrid.py | 16 +++++++++ .../models/olmoe/configuration_olmoe.py | 18 ++++++++++ .../configuration_paddleocr_vl.py | 16 +++++++++ .../models/phi/configuration_phi.py | 15 ++++++++ .../models/phi3/configuration_phi3.py | 13 +++++++ .../configuration_phi4_multimodal.py | 13 +++++++ .../models/qwen2/configuration_qwen2.py | 16 +++++++++ .../qwen3_moe/configuration_qwen3_moe.py | 23 ++++++++++++ .../configuration_qwen3_omni_moe.py | 23 ++++++++++++ .../configuration_qwen3_vl_moe.py | 23 ++++++++++++ .../models/seed_oss/configuration_seed_oss.py | 16 +++++++++ .../models/smollm3/configuration_smollm3.py | 16 +++++++++ .../models/smollm3/modular_smollm3.py | 16 +++++++++ .../solar_open/configuration_solar_open.py | 18 ++++++++++ .../models/solar_open/modular_solar_open.py | 18 ++++++++++ .../starcoder2/configuration_starcoder2.py | 15 ++++++++ .../models/t5gemma/configuration_t5gemma.py | 16 +++++++++ .../models/t5gemma2/configuration_t5gemma2.py | 36 +++++++++++++++++++ .../vaultgemma/configuration_vaultgemma.py | 16 +++++++++ .../models/youtu/configuration_youtu.py | 11 ++++++ .../models/youtu/modular_youtu.py | 11 ++++++ train_fsdp_tp.py | 2 +- verify_loading.py | 18 +++++----- 70 files changed, 1066 insertions(+), 14 deletions(-) mode change 100644 => 100755 run_compare.sh mode change 100644 => 100755 run_verify_all.sh diff --git a/.gitignore b/.gitignore index d26d399331af..6039e82a4390 100644 --- a/.gitignore +++ b/.gitignore @@ -8,6 +8,11 @@ __pycache__/ # C extensions *.so +verify_ckpt* +checkpoints* +*result* +debug* + # tests and logs tests/fixtures/cached_*_text.txt logs/ diff --git a/run_compare.sh b/run_compare.sh old mode 100644 new mode 100755 index eb47e1841fa9..75e06e8dae96 --- a/run_compare.sh +++ b/run_compare.sh @@ -53,4 +53,4 @@ cat "${LOG_FSDP_ONLY}.phase1" "${LOG_FSDP_ONLY}.phase2" > "$LOG_FSDP_ONLY" echo "" echo "=== Full Loss & Grad Diff (steps 0-19) ===" -git diff --no-index --color --word-diff=color "$LOG_FSDP_TP" "$LOG_FSDP_ONLY" || true +git diff --no-index --color --word-diff=color "$LOG_FSDP_TP" "$LOG_FSDP_ONLY" || true \ No newline at end of file diff --git a/run_verify_all.sh b/run_verify_all.sh old mode 100644 new mode 100755 index 16aa3267fe9a..3a1c9c08d8f9 --- a/run_verify_all.sh +++ b/run_verify_all.sh @@ -157,4 +157,4 @@ if [ "$HAS_FAIL" -eq 1 ]; then echo -e " ${YELLOW}cat $LOGDIR/$mode.log${NC}" fi done -fi +fi \ No newline at end of file diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index 67b3d748c08c..59aa7ee760e2 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -636,8 +636,9 @@ def apply_tensor_parallel(model, tp_mesh, tp_plan): if tp_plan == "auto": enable_sp = getattr(getattr(model.config, "distributed_config", None), "enable_sequence_parallel", False) - if enable_sp and hasattr(model.config, "base_model_sp_plan"): - base_plan = model.config.base_model_sp_plan + sp_plan = getattr(model.config, "base_model_sp_plan", None) + if enable_sp and sp_plan is not None: + base_plan = sp_plan else: base_plan = model.config.base_model_tp_plan or {} diff --git a/src/transformers/models/apertus/configuration_apertus.py b/src/transformers/models/apertus/configuration_apertus.py index 6864d11589f4..aa7ca9fe0676 100644 --- a/src/transformers/models/apertus/configuration_apertus.py +++ b/src/transformers/models/apertus/configuration_apertus.py @@ -54,6 +54,24 @@ class ApertusConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/apertus/modular_apertus.py b/src/transformers/models/apertus/modular_apertus.py index a901c3fba0fa..394dbaf7ea98 100644 --- a/src/transformers/models/apertus/modular_apertus.py +++ b/src/transformers/models/apertus/modular_apertus.py @@ -72,6 +72,24 @@ class ApertusConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/arcee/configuration_arcee.py b/src/transformers/models/arcee/configuration_arcee.py index 0f6e8aa9034b..0b5dc9c89671 100644 --- a/src/transformers/models/arcee/configuration_arcee.py +++ b/src/transformers/models/arcee/configuration_arcee.py @@ -54,6 +54,21 @@ class ArceeConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/arcee/modular_arcee.py b/src/transformers/models/arcee/modular_arcee.py index 316bb90db03a..0c43307747ab 100644 --- a/src/transformers/models/arcee/modular_arcee.py +++ b/src/transformers/models/arcee/modular_arcee.py @@ -58,6 +58,21 @@ class ArceeConfig(LlamaConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } vocab_size: int = 32000 hidden_size: int = 2560 diff --git a/src/transformers/models/aria/configuration_aria.py b/src/transformers/models/aria/configuration_aria.py index d240f4d5bba0..608760699bd2 100644 --- a/src/transformers/models/aria/configuration_aria.py +++ b/src/transformers/models/aria/configuration_aria.py @@ -50,6 +50,22 @@ class AriaTextConfig(PreTrainedConfig): "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/cohere/configuration_cohere.py b/src/transformers/models/cohere/configuration_cohere.py index 678b03eb8894..a365a51e2c8c 100644 --- a/src/transformers/models/cohere/configuration_cohere.py +++ b/src/transformers/models/cohere/configuration_cohere.py @@ -59,6 +59,24 @@ class CohereConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/cohere2/configuration_cohere2.py b/src/transformers/models/cohere2/configuration_cohere2.py index 715749c51e0f..54733a7f297f 100644 --- a/src/transformers/models/cohere2/configuration_cohere2.py +++ b/src/transformers/models/cohere2/configuration_cohere2.py @@ -58,6 +58,22 @@ class Cohere2Config(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/cohere2/modular_cohere2.py b/src/transformers/models/cohere2/modular_cohere2.py index dd7421f320dd..f46085ebaaac 100644 --- a/src/transformers/models/cohere2/modular_cohere2.py +++ b/src/transformers/models/cohere2/modular_cohere2.py @@ -79,6 +79,22 @@ class Cohere2Config(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/cwm/configuration_cwm.py b/src/transformers/models/cwm/configuration_cwm.py index a8ea587eb41f..006419da527e 100644 --- a/src/transformers/models/cwm/configuration_cwm.py +++ b/src/transformers/models/cwm/configuration_cwm.py @@ -55,6 +55,22 @@ class CwmConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py b/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py index 52f03ec0cc58..85c58ebc3e13 100644 --- a/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py +++ b/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py @@ -69,6 +69,7 @@ class DeepseekV2Config(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = None base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/deepseek_v2/modular_deepseek_v2.py b/src/transformers/models/deepseek_v2/modular_deepseek_v2.py index 30f439499c45..07b18b2b4beb 100644 --- a/src/transformers/models/deepseek_v2/modular_deepseek_v2.py +++ b/src/transformers/models/deepseek_v2/modular_deepseek_v2.py @@ -84,6 +84,8 @@ class DeepseekV2Config(LlamaConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + # MLA attention is not yet compatible with sequence parallelism. + base_model_sp_plan = None model_type = "deepseek_v2" keys_to_ignore_at_inference = ["past_key_values"] diff --git a/src/transformers/models/ernie4_5/configuration_ernie4_5.py b/src/transformers/models/ernie4_5/configuration_ernie4_5.py index f32929da5b70..84cdb4d8a745 100644 --- a/src/transformers/models/ernie4_5/configuration_ernie4_5.py +++ b/src/transformers/models/ernie4_5/configuration_ernie4_5.py @@ -56,6 +56,22 @@ class Ernie4_5Config(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/exaone4/configuration_exaone4.py b/src/transformers/models/exaone4/configuration_exaone4.py index 89ee40135153..ac746be58090 100644 --- a/src/transformers/models/exaone4/configuration_exaone4.py +++ b/src/transformers/models/exaone4/configuration_exaone4.py @@ -69,6 +69,24 @@ class Exaone4Config(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/exaone4/modular_exaone4.py b/src/transformers/models/exaone4/modular_exaone4.py index cc152edb42dc..213a9c513e71 100644 --- a/src/transformers/models/exaone4/modular_exaone4.py +++ b/src/transformers/models/exaone4/modular_exaone4.py @@ -98,6 +98,24 @@ class Exaone4Config(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/exaone_moe/configuration_exaone_moe.py b/src/transformers/models/exaone_moe/configuration_exaone_moe.py index 1f948a9eb5fb..dc9329882734 100644 --- a/src/transformers/models/exaone_moe/configuration_exaone_moe.py +++ b/src/transformers/models/exaone_moe/configuration_exaone_moe.py @@ -75,6 +75,24 @@ class ExaoneMoeConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/gemma/configuration_gemma.py b/src/transformers/models/gemma/configuration_gemma.py index bb676b5d8486..40567ac76b58 100644 --- a/src/transformers/models/gemma/configuration_gemma.py +++ b/src/transformers/models/gemma/configuration_gemma.py @@ -56,6 +56,22 @@ class GemmaConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/gemma/modular_gemma.py b/src/transformers/models/gemma/modular_gemma.py index 9a777a6be8ad..02857b31b2a1 100644 --- a/src/transformers/models/gemma/modular_gemma.py +++ b/src/transformers/models/gemma/modular_gemma.py @@ -75,6 +75,22 @@ class GemmaConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/gemma2/configuration_gemma2.py b/src/transformers/models/gemma2/configuration_gemma2.py index 852eb0e54e63..d0d53315b589 100644 --- a/src/transformers/models/gemma2/configuration_gemma2.py +++ b/src/transformers/models/gemma2/configuration_gemma2.py @@ -60,6 +60,22 @@ class Gemma2Config(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/gemma2/modular_gemma2.py b/src/transformers/models/gemma2/modular_gemma2.py index 191ce9e14401..a001311eeb6a 100644 --- a/src/transformers/models/gemma2/modular_gemma2.py +++ b/src/transformers/models/gemma2/modular_gemma2.py @@ -87,6 +87,22 @@ class Gemma2Config(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/gemma3/configuration_gemma3.py b/src/transformers/models/gemma3/configuration_gemma3.py index f7699c7cd994..24b680a6a3a4 100644 --- a/src/transformers/models/gemma3/configuration_gemma3.py +++ b/src/transformers/models/gemma3/configuration_gemma3.py @@ -67,6 +67,24 @@ class Gemma3TextConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/gemma3/modular_gemma3.py b/src/transformers/models/gemma3/modular_gemma3.py index 1ce5f8fa443c..dd948ee74847 100644 --- a/src/transformers/models/gemma3/modular_gemma3.py +++ b/src/transformers/models/gemma3/modular_gemma3.py @@ -95,6 +95,24 @@ class Gemma3TextConfig(Gemma2Config, PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } 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 950f0ea78b47..e62e745e6f96 100644 --- a/src/transformers/models/gemma3n/configuration_gemma3n.py +++ b/src/transformers/models/gemma3n/configuration_gemma3n.py @@ -88,6 +88,24 @@ class Gemma3nTextConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/gemma3n/modular_gemma3n.py b/src/transformers/models/gemma3n/modular_gemma3n.py index 53a2b066d1d7..20708001ef6c 100644 --- a/src/transformers/models/gemma3n/modular_gemma3n.py +++ b/src/transformers/models/gemma3n/modular_gemma3n.py @@ -115,6 +115,24 @@ class Gemma3nTextConfig(Gemma3TextConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } default_theta = {"global": 1_000_000.0, "local": 10_000.0} vocab_size: int = 262_400 diff --git a/src/transformers/models/glm/configuration_glm.py b/src/transformers/models/glm/configuration_glm.py index d9d07638f6cb..494000e3599a 100644 --- a/src/transformers/models/glm/configuration_glm.py +++ b/src/transformers/models/glm/configuration_glm.py @@ -48,6 +48,21 @@ class GlmConfig(PreTrainedConfig): "layers.*.mlp.gate_up_proj": TPStyle("packed_colwise", "none"), # fused gate/up shards stay local for chunk "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_up_proj": TPStyle("packed_colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/glm4/configuration_glm4.py b/src/transformers/models/glm4/configuration_glm4.py index a91eab20b12c..89328611e2df 100644 --- a/src/transformers/models/glm4/configuration_glm4.py +++ b/src/transformers/models/glm4/configuration_glm4.py @@ -50,6 +50,21 @@ class Glm4Config(PreTrainedConfig): ), # we need to replicate here due to the `chunk` operation "layers.*.mlp.down_proj": TPStyle("vocab", "allreduce"), # input is replicated due to the `chunk` operation } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_up_proj": TPStyle("colwise", "allgather"), # fused gate/up needs full tensor for chunk + "layers.*.mlp.down_proj": TPStyle("vocab", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/gpt_neox/configuration_gpt_neox.py b/src/transformers/models/gpt_neox/configuration_gpt_neox.py index 782bea43357f..0b39d857a1b7 100644 --- a/src/transformers/models/gpt_neox/configuration_gpt_neox.py +++ b/src/transformers/models/gpt_neox/configuration_gpt_neox.py @@ -52,6 +52,16 @@ class GPTNeoXConfig(PreTrainedConfig): "layers.*.mlp.dense_h_to_4h": TPStyle("colwise", "none"), "layers.*.mlp.dense_4h_to_h": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.dense_h_to_4h": TPStyle("colwise", "none"), + "layers.*.mlp.dense_4h_to_h": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_in": (["input_ids"], ["inputs_embeds"]), "emb_dropout": (["inputs_embeds"], ["hidden_states"]), diff --git a/src/transformers/models/granite/configuration_granite.py b/src/transformers/models/granite/configuration_granite.py index 6696c4d33685..64a58e9c3738 100644 --- a/src/transformers/models/granite/configuration_granite.py +++ b/src/transformers/models/granite/configuration_granite.py @@ -56,6 +56,22 @@ class GraniteConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/helium/configuration_helium.py b/src/transformers/models/helium/configuration_helium.py index 8e2f44ffca57..cf9d461f1439 100644 --- a/src/transformers/models/helium/configuration_helium.py +++ b/src/transformers/models/helium/configuration_helium.py @@ -50,6 +50,22 @@ class HeliumConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), 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 ca3d81d225be..f25435300f7d 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 @@ -68,6 +68,22 @@ class HiggsAudioV2Config(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/jais2/configuration_jais2.py b/src/transformers/models/jais2/configuration_jais2.py index 4886c5c45acf..af46f9a134fb 100644 --- a/src/transformers/models/jais2/configuration_jais2.py +++ b/src/transformers/models/jais2/configuration_jais2.py @@ -55,6 +55,21 @@ class Jais2Config(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/jais2/modular_jais2.py b/src/transformers/models/jais2/modular_jais2.py index 90279d121eaa..8a1dc6b0dbb4 100644 --- a/src/transformers/models/jais2/modular_jais2.py +++ b/src/transformers/models/jais2/modular_jais2.py @@ -39,6 +39,21 @@ class Jais2Config(LlamaConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } vocab_size: int = 150272 hidden_size: int = 3328 diff --git a/src/transformers/models/llama/configuration_llama.py b/src/transformers/models/llama/configuration_llama.py index f0d0d8a97197..61d7cc76c84e 100644 --- a/src/transformers/models/llama/configuration_llama.py +++ b/src/transformers/models/llama/configuration_llama.py @@ -56,6 +56,22 @@ class LlamaConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/minimax_m2/configuration_minimax_m2.py b/src/transformers/models/minimax_m2/configuration_minimax_m2.py index 0ab84d885690..286af98c7c53 100644 --- a/src/transformers/models/minimax_m2/configuration_minimax_m2.py +++ b/src/transformers/models/minimax_m2/configuration_minimax_m2.py @@ -59,6 +59,24 @@ class MiniMaxM2Config(PreTrainedConfig): shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, ), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "allgather"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "allgather"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "allgather"), + "layers.*.self_attn.o_proj": TPStyle("vocab", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather_split"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/minimax_m2/modular_minimax_m2.py b/src/transformers/models/minimax_m2/modular_minimax_m2.py index b0994cf4cd89..025ae2c86a1c 100644 --- a/src/transformers/models/minimax_m2/modular_minimax_m2.py +++ b/src/transformers/models/minimax_m2/modular_minimax_m2.py @@ -78,6 +78,24 @@ class MiniMaxM2Config(PreTrainedConfig): shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, ), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "allgather"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "allgather"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "allgather"), + "layers.*.self_attn.o_proj": TPStyle("vocab", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather_split"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/ministral/configuration_ministral.py b/src/transformers/models/ministral/configuration_ministral.py index cdd2074230b7..2f72a02ceefe 100644 --- a/src/transformers/models/ministral/configuration_ministral.py +++ b/src/transformers/models/ministral/configuration_ministral.py @@ -58,6 +58,22 @@ class MinistralConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/ministral3/configuration_ministral3.py b/src/transformers/models/ministral3/configuration_ministral3.py index 82cefa62a805..9c17dce093e1 100644 --- a/src/transformers/models/ministral3/configuration_ministral3.py +++ b/src/transformers/models/ministral3/configuration_ministral3.py @@ -64,6 +64,22 @@ class Ministral3Config(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/mistral/configuration_mistral.py b/src/transformers/models/mistral/configuration_mistral.py index 2b64d139f145..1ef743bd0409 100644 --- a/src/transformers/models/mistral/configuration_mistral.py +++ b/src/transformers/models/mistral/configuration_mistral.py @@ -55,6 +55,22 @@ class MistralConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/nanochat/configuration_nanochat.py b/src/transformers/models/nanochat/configuration_nanochat.py index 0752837e4753..37de40aab631 100644 --- a/src/transformers/models/nanochat/configuration_nanochat.py +++ b/src/transformers/models/nanochat/configuration_nanochat.py @@ -51,6 +51,24 @@ class NanoChatConfig(PretrainedConfig): "layers.*.mlp.fc1": TPStyle("colwise", "none"), "layers.*.mlp.fc2": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } vocab_size: int = 50304 hidden_size: int = 768 diff --git a/src/transformers/models/olmo/configuration_olmo.py b/src/transformers/models/olmo/configuration_olmo.py index 1e626e42d98e..ee2a8fae11ef 100644 --- a/src/transformers/models/olmo/configuration_olmo.py +++ b/src/transformers/models/olmo/configuration_olmo.py @@ -59,6 +59,22 @@ class OlmoConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/olmo2/configuration_olmo2.py b/src/transformers/models/olmo2/configuration_olmo2.py index 96dbdd2908fb..c6861d4e3e4f 100644 --- a/src/transformers/models/olmo2/configuration_olmo2.py +++ b/src/transformers/models/olmo2/configuration_olmo2.py @@ -70,6 +70,24 @@ class Olmo2Config(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/olmo2/modular_olmo2.py b/src/transformers/models/olmo2/modular_olmo2.py index 94ca63e2a7c0..9b74f4cd0154 100644 --- a/src/transformers/models/olmo2/modular_olmo2.py +++ b/src/transformers/models/olmo2/modular_olmo2.py @@ -83,6 +83,24 @@ class Olmo2Config(OlmoConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/olmo3/configuration_olmo3.py b/src/transformers/models/olmo3/configuration_olmo3.py index fb7f445cf079..5675ff02aef1 100644 --- a/src/transformers/models/olmo3/configuration_olmo3.py +++ b/src/transformers/models/olmo3/configuration_olmo3.py @@ -65,6 +65,24 @@ class Olmo3Config(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/olmo3/modular_olmo3.py b/src/transformers/models/olmo3/modular_olmo3.py index fba53df7502c..7eb5e105235d 100644 --- a/src/transformers/models/olmo3/modular_olmo3.py +++ b/src/transformers/models/olmo3/modular_olmo3.py @@ -79,6 +79,24 @@ class Olmo3Config(Olmo2Config): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py index 65e9c2043045..cb081cf9fa7c 100644 --- a/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py +++ b/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py @@ -90,6 +90,22 @@ class OlmoHybridConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/olmoe/configuration_olmoe.py b/src/transformers/models/olmoe/configuration_olmoe.py index 17def985f001..99ed6c9668e3 100644 --- a/src/transformers/models/olmoe/configuration_olmoe.py +++ b/src/transformers/models/olmoe/configuration_olmoe.py @@ -57,6 +57,24 @@ class OlmoeConfig(PreTrainedConfig): shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, ), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "allgather"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "allgather"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "allgather"), + "layers.*.self_attn.o_proj": TPStyle("vocab", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather_split"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } vocab_size: int = 50304 hidden_size: int = 2048 diff --git a/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py b/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py index 4da2e361d0aa..3993cd6f0a87 100644 --- a/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py +++ b/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py @@ -104,6 +104,22 @@ class PaddleOCRTextConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/phi/configuration_phi.py b/src/transformers/models/phi/configuration_phi.py index c04e91344b45..91b846bb5c74 100644 --- a/src/transformers/models/phi/configuration_phi.py +++ b/src/transformers/models/phi/configuration_phi.py @@ -54,6 +54,21 @@ class PhiConfig(PreTrainedConfig): "layers.*.mlp.fc1": TPStyle("colwise", "none"), "layers.*.mlp.fc2": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.dense": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.fc1": TPStyle("colwise", "none"), + "layers.*.mlp.fc2": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "embed_dropout": (["inputs_embeds"], ["inputs_embeds"]), diff --git a/src/transformers/models/phi3/configuration_phi3.py b/src/transformers/models/phi3/configuration_phi3.py index 8e5f50697762..b236a280d390 100644 --- a/src/transformers/models/phi3/configuration_phi3.py +++ b/src/transformers/models/phi3/configuration_phi3.py @@ -57,6 +57,19 @@ class Phi3Config(PreTrainedConfig): ), # we need to replicate here due to the `chunk` operation "layers.*.mlp.down_proj": TPStyle("vocab", "allreduce"), # input is replicated due to the `chunk` operation } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.qkv_proj": TPStyle("colwise", "allgather"), # fused qkv needs full tensor for slicing + "layers.*.self_attn.o_proj": TPStyle("vocab", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_up_proj": TPStyle("colwise", "allgather"), # fused gate/up needs full tensor for chunk + "layers.*.mlp.down_proj": TPStyle("vocab", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py b/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py index 62e0925d0e3e..dc66237fbf0a 100644 --- a/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py +++ b/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py @@ -191,6 +191,19 @@ class Phi4MultimodalConfig(PreTrainedConfig): ), # we need to replicate here due to the `chunk` operation "layers.*.mlp.down_proj": TPStyle("vocab", "allreduce"), # input is replicated due to the `chunk` operation } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.qkv_proj": TPStyle("colwise", "allgather"), # fused qkv needs full tensor for slicing + "layers.*.self_attn.o_proj": TPStyle("vocab", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_up_proj": TPStyle("colwise", "allgather"), # fused gate/up needs full tensor for chunk + "layers.*.mlp.down_proj": TPStyle("vocab", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/qwen2/configuration_qwen2.py b/src/transformers/models/qwen2/configuration_qwen2.py index c48a7639fc5e..633275345908 100644 --- a/src/transformers/models/qwen2/configuration_qwen2.py +++ b/src/transformers/models/qwen2/configuration_qwen2.py @@ -53,6 +53,22 @@ class Qwen2Config(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py b/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py index 4a69ea88e43c..763de442c937 100644 --- a/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py +++ b/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py @@ -68,6 +68,29 @@ class Qwen3MoeConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather_split"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), 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 a56602fa2047..67a2cfd84577 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 @@ -386,6 +386,29 @@ class Qwen3OmniMoeTalkerTextConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather_split"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), 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 ec276a41d8e9..867f0d0a22d9 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 @@ -65,6 +65,29 @@ class Qwen3VLMoeTextConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather_split"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/seed_oss/configuration_seed_oss.py b/src/transformers/models/seed_oss/configuration_seed_oss.py index 7be4662965c1..72f948e5b1cc 100644 --- a/src/transformers/models/seed_oss/configuration_seed_oss.py +++ b/src/transformers/models/seed_oss/configuration_seed_oss.py @@ -54,6 +54,22 @@ class SeedOssConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/smollm3/configuration_smollm3.py b/src/transformers/models/smollm3/configuration_smollm3.py index a4898567d4c0..3c958c865543 100644 --- a/src/transformers/models/smollm3/configuration_smollm3.py +++ b/src/transformers/models/smollm3/configuration_smollm3.py @@ -64,6 +64,22 @@ class SmolLM3Config(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/smollm3/modular_smollm3.py b/src/transformers/models/smollm3/modular_smollm3.py index d62061154d50..0917a4db0446 100644 --- a/src/transformers/models/smollm3/modular_smollm3.py +++ b/src/transformers/models/smollm3/modular_smollm3.py @@ -80,6 +80,22 @@ class SmolLM3Config(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/solar_open/configuration_solar_open.py b/src/transformers/models/solar_open/configuration_solar_open.py index 9e35341a59bc..65bfd946f5ea 100644 --- a/src/transformers/models/solar_open/configuration_solar_open.py +++ b/src/transformers/models/solar_open/configuration_solar_open.py @@ -49,6 +49,24 @@ class SolarOpenConfig(PreTrainedConfig): shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, ), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather_split"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/solar_open/modular_solar_open.py b/src/transformers/models/solar_open/modular_solar_open.py index 8363cfec333e..f81f4bbff3fd 100644 --- a/src/transformers/models/solar_open/modular_solar_open.py +++ b/src/transformers/models/solar_open/modular_solar_open.py @@ -55,6 +55,24 @@ class SolarOpenConfig(Glm4MoeConfig): shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, ), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather_split"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } vocab_size: int = 196608 moe_intermediate_size: int = 1280 diff --git a/src/transformers/models/starcoder2/configuration_starcoder2.py b/src/transformers/models/starcoder2/configuration_starcoder2.py index f55d8d17ab0c..1b7a2879fe3c 100644 --- a/src/transformers/models/starcoder2/configuration_starcoder2.py +++ b/src/transformers/models/starcoder2/configuration_starcoder2.py @@ -53,6 +53,21 @@ class Starcoder2Config(PreTrainedConfig): "layers.*.mlp.c_fc": TPStyle("colwise", "none"), "layers.*.mlp.c_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.c_fc": TPStyle("colwise", "none"), + "layers.*.mlp.c_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/t5gemma/configuration_t5gemma.py b/src/transformers/models/t5gemma/configuration_t5gemma.py index 9a5f93ba8a53..b4b39ba7d743 100644 --- a/src/transformers/models/t5gemma/configuration_t5gemma.py +++ b/src/transformers/models/t5gemma/configuration_t5gemma.py @@ -60,6 +60,22 @@ class T5GemmaModuleConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/t5gemma2/configuration_t5gemma2.py b/src/transformers/models/t5gemma2/configuration_t5gemma2.py index 1f04bd70d9e6..86525e4f090d 100644 --- a/src/transformers/models/t5gemma2/configuration_t5gemma2.py +++ b/src/transformers/models/t5gemma2/configuration_t5gemma2.py @@ -54,6 +54,24 @@ class T5Gemma2TextConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), @@ -227,6 +245,24 @@ class T5Gemma2DecoderConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/vaultgemma/configuration_vaultgemma.py b/src/transformers/models/vaultgemma/configuration_vaultgemma.py index 1fd184568ded..5bcec13cb160 100644 --- a/src/transformers/models/vaultgemma/configuration_vaultgemma.py +++ b/src/transformers/models/vaultgemma/configuration_vaultgemma.py @@ -59,6 +59,22 @@ class VaultGemmaConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/youtu/configuration_youtu.py b/src/transformers/models/youtu/configuration_youtu.py index eaef7650cd34..2eee3b5859c6 100644 --- a/src/transformers/models/youtu/configuration_youtu.py +++ b/src/transformers/models/youtu/configuration_youtu.py @@ -87,6 +87,17 @@ class YoutuConfig(PreTrainedConfig): rope_interleave: bool | None = True attention_bias: bool = False attention_dropout: float | int | None = 0.0 + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } embedding_initializer_range: float | None = None def __post_init__(self, **kwargs): diff --git a/src/transformers/models/youtu/modular_youtu.py b/src/transformers/models/youtu/modular_youtu.py index a4ee903c61e1..617972bf8a11 100644 --- a/src/transformers/models/youtu/modular_youtu.py +++ b/src/transformers/models/youtu/modular_youtu.py @@ -65,6 +65,17 @@ class YoutuConfig(DeepseekV3Config): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } attribute_map = {} vocab_size: int = 128256 diff --git a/train_fsdp_tp.py b/train_fsdp_tp.py index 0232f8b3bc3d..ab0737a9d0b6 100644 --- a/train_fsdp_tp.py +++ b/train_fsdp_tp.py @@ -122,4 +122,4 @@ def build_fixed_batches(dp_rank): if rank == 0: print(f"Saved to {args.save_dir}") - torch.distributed.destroy_process_group() + torch.distributed.destroy_process_group() \ No newline at end of file diff --git a/verify_loading.py b/verify_loading.py index ea008f9626f7..ba0f60f31fd3 100644 --- a/verify_loading.py +++ b/verify_loading.py @@ -39,12 +39,12 @@ torch.cuda.set_device(0) configs = { - "single_gpu": None, - "fsdp": DistributedConfig(fsdp_size=2, fsdp_plan="auto"), - "tp": DistributedConfig(tp_size=2, tp_plan="auto"), - "tp_sp": DistributedConfig(tp_size=2, tp_plan="auto", enable_sequence_parallel=True), - "tp_fsdp": DistributedConfig(tp_size=2, tp_plan="auto", fsdp_size=2, fsdp_plan="auto"), - "tp_sp_fsdp": DistributedConfig(tp_size=2, tp_plan="auto", fsdp_size=2, fsdp_plan="auto", enable_sequence_parallel=True), + "single_gpu": lambda: None, + "fsdp": lambda: DistributedConfig(fsdp_size=2, fsdp_plan="auto"), + "tp": lambda: DistributedConfig(tp_size=2, tp_plan="auto"), + "tp_sp": lambda: DistributedConfig(tp_size=2, tp_plan="auto", enable_sequence_parallel=True), + "tp_fsdp": lambda: DistributedConfig(tp_size=2, tp_plan="auto", fsdp_size=2, fsdp_plan="auto"), + "tp_sp_fsdp": lambda: DistributedConfig(tp_size=2, tp_plan="auto", fsdp_size=2, fsdp_plan="auto", enable_sequence_parallel=True), } tokenizer = AutoTokenizer.from_pretrained(model_id) @@ -64,7 +64,7 @@ def compute_loss(model): # Pad sequence length to a multiple of tp_size so DTensor Shard(1) splits evenly # across ranks in SP mode. Always pad (even for non-TP modes) so that all modes # compute on the same input and losses are directly comparable. - max_tp = max((c.tp_size if c is not None else 1) for c in configs.values()) + max_tp = 2 # all TP configs use tp_size=2 seq_len = input_ids.shape[1] if seq_len % max_tp != 0: pad_len = max_tp - (seq_len % max_tp) @@ -88,7 +88,7 @@ def compute_loss(model): # --- Step 1: Load original model and compute loss --- -model = AutoModelForCausalLM.from_pretrained(model_id, distributed_config=configs[args.mode], dtype=torch.float32) +model = AutoModelForCausalLM.from_pretrained(model_id, distributed_config=configs[args.mode](), dtype=torch.float32) if args.mode == "single_gpu": model = model.to("cuda:0") @@ -116,7 +116,7 @@ def compute_loss(model): torch.cuda.empty_cache() # --- Step 3: Reload from saved checkpoint and compute loss --- -model2 = AutoModelForCausalLM.from_pretrained(save_dir, distributed_config=configs[args.mode], dtype=torch.float32) +model2 = AutoModelForCausalLM.from_pretrained(save_dir, distributed_config=configs[args.mode](), dtype=torch.float32) if args.mode == "single_gpu": model2 = model2.to("cuda:0") From 1a51928a2689c697d643de4aa71de0ef84a7ca2a Mon Sep 17 00:00:00 2001 From: 3outeille Date: Sat, 18 Apr 2026 14:58:52 +0000 Subject: [PATCH 032/116] add deepseek v2 sp plan --- .../deepseek_v2/configuration_deepseek_v2.py | 28 ++++++++++++++++- .../models/deepseek_v2/modular_deepseek_v2.py | 31 +++++++++++++++++-- 2 files changed, 56 insertions(+), 3 deletions(-) diff --git a/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py b/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py index 85c58ebc3e13..61b158613355 100644 --- a/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py +++ b/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py @@ -69,7 +69,33 @@ class DeepseekV2Config(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } - base_model_sp_plan = None + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + # 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": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_b_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.kv_b_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather_split"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), + # Shared experts output must stay Replicate to match experts output + # (summed inside the MoE block, before outer allgather_split handles SP). + "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/deepseek_v2/modular_deepseek_v2.py b/src/transformers/models/deepseek_v2/modular_deepseek_v2.py index 07b18b2b4beb..9bfceb8eff2c 100644 --- a/src/transformers/models/deepseek_v2/modular_deepseek_v2.py +++ b/src/transformers/models/deepseek_v2/modular_deepseek_v2.py @@ -84,8 +84,35 @@ class DeepseekV2Config(LlamaConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } - # MLA attention is not yet compatible with sequence parallelism. - base_model_sp_plan = None + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + # 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": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_b_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.kv_b_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather_split"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), + # 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": TPStyle("colwise", "none"), + "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + "lm_head": TPStyle("colwise", "loss_parallel"), + } model_type = "deepseek_v2" keys_to_ignore_at_inference = ["past_key_values"] From fd3a7221bb895289f76168c45ac01fd6623d65f6 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Sat, 18 Apr 2026 15:53:24 +0000 Subject: [PATCH 033/116] undo sp plan for some tricky models --- .../exaone_moe/configuration_exaone_moe.py | 19 +------------------ .../models/gemma3n/configuration_gemma3n.py | 19 +------------------ .../models/gemma3n/modular_gemma3n.py | 19 +------------------ .../models/nanochat/configuration_nanochat.py | 19 +------------------ .../olmo_hybrid/configuration_olmo_hybrid.py | 17 +---------------- .../models/youtu/configuration_youtu.py | 12 +----------- .../models/youtu/modular_youtu.py | 12 +----------- 7 files changed, 7 insertions(+), 110 deletions(-) diff --git a/src/transformers/models/exaone_moe/configuration_exaone_moe.py b/src/transformers/models/exaone_moe/configuration_exaone_moe.py index dc9329882734..04e25fbd056e 100644 --- a/src/transformers/models/exaone_moe/configuration_exaone_moe.py +++ b/src/transformers/models/exaone_moe/configuration_exaone_moe.py @@ -75,24 +75,7 @@ class ExaoneMoeConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } - base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), - } + base_model_sp_plan = None base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/gemma3n/configuration_gemma3n.py b/src/transformers/models/gemma3n/configuration_gemma3n.py index e62e745e6f96..7d299c44ee92 100644 --- a/src/transformers/models/gemma3n/configuration_gemma3n.py +++ b/src/transformers/models/gemma3n/configuration_gemma3n.py @@ -88,24 +88,7 @@ class Gemma3nTextConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } - base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), - } + base_model_sp_plan = None base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/gemma3n/modular_gemma3n.py b/src/transformers/models/gemma3n/modular_gemma3n.py index 20708001ef6c..09aaf0e9fbd5 100644 --- a/src/transformers/models/gemma3n/modular_gemma3n.py +++ b/src/transformers/models/gemma3n/modular_gemma3n.py @@ -115,24 +115,7 @@ class Gemma3nTextConfig(Gemma3TextConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } - base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), - } + base_model_sp_plan = None default_theta = {"global": 1_000_000.0, "local": 10_000.0} vocab_size: int = 262_400 diff --git a/src/transformers/models/nanochat/configuration_nanochat.py b/src/transformers/models/nanochat/configuration_nanochat.py index 37de40aab631..2f419697e952 100644 --- a/src/transformers/models/nanochat/configuration_nanochat.py +++ b/src/transformers/models/nanochat/configuration_nanochat.py @@ -51,24 +51,7 @@ class NanoChatConfig(PretrainedConfig): "layers.*.mlp.fc1": TPStyle("colwise", "none"), "layers.*.mlp.fc2": TPStyle("rowwise", "allreduce"), } - base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), - } + base_model_sp_plan = None vocab_size: int = 50304 hidden_size: int = 768 diff --git a/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py index cb081cf9fa7c..5a5f2b1d8a33 100644 --- a/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py +++ b/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py @@ -90,22 +90,7 @@ class OlmoHybridConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } - base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), - } + base_model_sp_plan = None base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/youtu/configuration_youtu.py b/src/transformers/models/youtu/configuration_youtu.py index 2eee3b5859c6..9a3592c871e9 100644 --- a/src/transformers/models/youtu/configuration_youtu.py +++ b/src/transformers/models/youtu/configuration_youtu.py @@ -87,17 +87,7 @@ class YoutuConfig(PreTrainedConfig): rope_interleave: bool | None = True attention_bias: bool = False attention_dropout: float | int | None = 0.0 - base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), - } + base_model_sp_plan = None embedding_initializer_range: float | None = None def __post_init__(self, **kwargs): diff --git a/src/transformers/models/youtu/modular_youtu.py b/src/transformers/models/youtu/modular_youtu.py index 617972bf8a11..09ccb4e8063d 100644 --- a/src/transformers/models/youtu/modular_youtu.py +++ b/src/transformers/models/youtu/modular_youtu.py @@ -65,17 +65,7 @@ class YoutuConfig(DeepseekV3Config): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } - base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), - } + base_model_sp_plan = None attribute_map = {} vocab_size: int = 128256 From 253b89ee8a33ed4928bcfb2374f3a64938d2fabd Mon Sep 17 00:00:00 2001 From: 3outeille Date: Mon, 20 Apr 2026 10:47:47 +0000 Subject: [PATCH 034/116] remove lm_head from config --- src/transformers/configuration_utils.py | 2 ++ .../integrations/tensor_parallel.py | 21 +++++------ src/transformers/modeling_utils.py | 14 ++++++-- .../models/afmoe/modeling_afmoe.py | 1 + .../models/afmoe/modular_afmoe.py | 1 + .../models/apertus/configuration_apertus.py | 1 - .../models/apertus/modeling_apertus.py | 1 + .../models/apertus/modular_apertus.py | 1 - .../models/arcee/configuration_arcee.py | 1 - .../models/arcee/modeling_arcee.py | 1 + .../models/arcee/modular_arcee.py | 1 - .../models/aria/configuration_aria.py | 1 - src/transformers/models/aria/modeling_aria.py | 1 + .../audioflamingo3/modeling_audioflamingo3.py | 1 + .../audioflamingo3/modular_audioflamingo3.py | 1 + .../models/bamba/modeling_bamba.py | 1 + .../models/bitnet/modeling_bitnet.py | 1 + .../models/bitnet/modular_bitnet.py | 1 + .../models/cohere/configuration_cohere.py | 1 - .../models/cohere/modeling_cohere.py | 1 + .../models/cohere2/configuration_cohere2.py | 1 - .../models/cohere2/modeling_cohere2.py | 1 + .../models/cohere2/modular_cohere2.py | 1 - src/transformers/models/csm/modeling_csm.py | 1 + src/transformers/models/csm/modular_csm.py | 1 + .../models/cwm/configuration_cwm.py | 1 - src/transformers/models/cwm/modeling_cwm.py | 1 + src/transformers/models/dbrx/modeling_dbrx.py | 1 + src/transformers/models/dbrx/modular_dbrx.py | 1 + .../deepseek_v2/configuration_deepseek_v2.py | 7 ++-- .../deepseek_v2/modeling_deepseek_v2.py | 1 + .../models/deepseek_v2/modular_deepseek_v2.py | 1 - .../deepseek_v3/modeling_deepseek_v3.py | 1 + .../models/diffllama/modeling_diffllama.py | 1 + src/transformers/models/doge/modeling_doge.py | 1 + .../models/dots1/modeling_dots1.py | 1 + src/transformers/models/emu3/modeling_emu3.py | 1 + .../models/ernie4_5/configuration_ernie4_5.py | 1 - .../models/ernie4_5/modeling_ernie4_5.py | 1 + .../ernie4_5_moe/modeling_ernie4_5_moe.py | 1 + .../models/eurobert/modeling_eurobert.py | 1 + .../models/eurobert/modular_eurobert.py | 1 + .../models/exaone4/configuration_exaone4.py | 1 - .../models/exaone4/modeling_exaone4.py | 1 + .../models/exaone4/modular_exaone4.py | 1 - .../exaone_moe/configuration_exaone_moe.py | 18 +++++++++- .../models/exaone_moe/modeling_exaone_moe.py | 1 + .../models/falcon_h1/modeling_falcon_h1.py | 1 + .../models/flex_olmo/modeling_flex_olmo.py | 1 + .../models/gemma/configuration_gemma.py | 1 - .../models/gemma/modeling_gemma.py | 1 + .../models/gemma/modular_gemma.py | 1 - .../models/gemma2/configuration_gemma2.py | 1 - .../models/gemma2/modeling_gemma2.py | 1 + .../models/gemma2/modular_gemma2.py | 1 - .../models/gemma3/configuration_gemma3.py | 1 - .../models/gemma3/modeling_gemma3.py | 1 + .../models/gemma3/modular_gemma3.py | 1 - .../models/gemma3n/modeling_gemma3n.py | 1 + .../models/gemma4/modeling_gemma4.py | 1 + .../models/glm/configuration_glm.py | 1 - src/transformers/models/glm/modeling_glm.py | 1 + .../models/glm4/configuration_glm4.py | 1 - src/transformers/models/glm4/modeling_glm4.py | 1 + .../models/glm4_moe/modeling_glm4_moe.py | 1 + .../glm4_moe_lite/modeling_glm4_moe_lite.py | 1 + .../glm_moe_dsa/modeling_glm_moe_dsa.py | 1 + .../models/glmasr/modeling_glmasr.py | 1 + .../models/gpt_neox/configuration_gpt_neox.py | 1 - .../models/gpt_oss/modeling_gpt_oss.py | 1 + .../models/granite/configuration_granite.py | 1 - .../models/granite/modeling_granite.py | 1 + .../models/granitemoe/modeling_granitemoe.py | 1 + .../modeling_granitemoehybrid.py | 1 + .../modeling_granitemoeshared.py | 1 + .../models/helium/configuration_helium.py | 1 - .../models/helium/modeling_helium.py | 1 + .../configuration_higgs_audio_v2.py | 1 - .../modeling_hunyuan_v1_dense.py | 1 + .../hunyuan_v1_moe/modeling_hunyuan_v1_moe.py | 1 + .../models/jais2/configuration_jais2.py | 1 - .../models/jais2/modeling_jais2.py | 1 + .../models/jais2/modular_jais2.py | 1 - .../models/jamba/modeling_jamba.py | 1 + .../modeling_kyutai_speech_to_text.py | 1 + src/transformers/models/lfm2/modeling_lfm2.py | 1 + .../models/lfm2_moe/modeling_lfm2_moe.py | 1 + .../models/llama/configuration_llama.py | 1 - .../models/llama/modeling_llama.py | 1 + .../models/llama4/modeling_llama4.py | 1 + .../longcat_flash/modeling_longcat_flash.py | 1 + .../models/minimax/modeling_minimax.py | 1 + .../minimax_m2/configuration_minimax_m2.py | 1 - .../models/minimax_m2/modeling_minimax_m2.py | 1 + .../models/minimax_m2/modular_minimax_m2.py | 1 - .../ministral/configuration_ministral.py | 1 - .../models/ministral/modeling_ministral.py | 1 + .../ministral3/configuration_ministral3.py | 1 - .../models/ministral3/modeling_ministral3.py | 1 + .../models/mistral/configuration_mistral.py | 1 - .../models/mistral/modeling_mistral.py | 1 + .../models/mistral4/modeling_mistral4.py | 1 + .../models/mixtral/configuration_mixtral.py | 1 - .../models/mixtral/modeling_mixtral.py | 1 + .../musicflamingo/modeling_musicflamingo.py | 1 + .../models/nanochat/modeling_nanochat.py | 1 + .../models/nanochat/modular_nanochat.py | 1 + .../models/olmo/configuration_olmo.py | 1 - src/transformers/models/olmo/modeling_olmo.py | 1 + .../models/olmo2/configuration_olmo2.py | 1 - .../models/olmo2/modeling_olmo2.py | 1 + .../models/olmo2/modular_olmo2.py | 1 - .../models/olmo3/configuration_olmo3.py | 1 - .../models/olmo3/modeling_olmo3.py | 1 + .../models/olmo3/modular_olmo3.py | 1 - .../olmo_hybrid/configuration_olmo_hybrid.py | 16 ++++++++- .../olmo_hybrid/modeling_olmo_hybrid.py | 1 + .../models/olmoe/configuration_olmoe.py | 1 - .../models/olmoe/modeling_olmoe.py | 1 + .../configuration_paddleocr_vl.py | 1 - .../models/phi/configuration_phi.py | 1 - src/transformers/models/phi/modeling_phi.py | 1 + .../models/phi3/configuration_phi3.py | 1 - src/transformers/models/phi3/modeling_phi3.py | 1 + .../configuration_phi4_multimodal.py | 1 - .../modeling_phi4_multimodal.py | 1 + .../models/phimoe/modeling_phimoe.py | 1 + .../models/qwen2/configuration_qwen2.py | 1 - .../models/qwen2/modeling_qwen2.py | 1 + .../models/qwen2_moe/modeling_qwen2_moe.py | 1 + .../models/qwen2_moe/modular_qwen2_moe.py | 1 + .../models/qwen3/configuration_qwen3.py | 1 - .../models/qwen3/modeling_qwen3.py | 1 + .../models/qwen3_5/modeling_qwen3_5.py | 1 + .../qwen3_5_moe/modeling_qwen3_5_moe.py | 2 ++ .../models/qwen3_5_moe/modular_qwen3_5_moe.py | 1 + .../qwen3_moe/configuration_qwen3_moe.py | 1 - .../models/qwen3_moe/modeling_qwen3_moe.py | 1 + .../models/qwen3_next/modeling_qwen3_next.py | 1 + .../configuration_qwen3_omni_moe.py | 2 -- .../qwen3_omni_moe/modeling_qwen3_omni_moe.py | 2 ++ .../models/qwen3_vl/modeling_qwen3_vl.py | 21 +---------- .../configuration_qwen3_vl_moe.py | 1 - .../qwen3_vl_moe/modeling_qwen3_vl_moe.py | 26 +------------- .../models/seed_oss/configuration_seed_oss.py | 1 - .../models/seed_oss/modeling_seed_oss.py | 1 + .../models/smollm3/configuration_smollm3.py | 1 - .../models/smollm3/modeling_smollm3.py | 1 + .../models/smollm3/modular_smollm3.py | 1 - .../solar_open/configuration_solar_open.py | 35 +++++++++---------- .../models/solar_open/modeling_solar_open.py | 1 + .../models/solar_open/modular_solar_open.py | 1 - .../starcoder2/configuration_starcoder2.py | 1 - .../models/starcoder2/modeling_starcoder2.py | 1 + .../models/t5gemma/configuration_t5gemma.py | 1 - .../models/t5gemma2/configuration_t5gemma2.py | 2 -- .../vaultgemma/configuration_vaultgemma.py | 1 - .../models/vaultgemma/modeling_vaultgemma.py | 1 + .../vibevoice_asr/modeling_vibevoice_asr.py | 1 + .../modeling_voxtral_realtime.py | 1 + .../models/youtu/modeling_youtu.py | 1 + 161 files changed, 174 insertions(+), 142 deletions(-) diff --git a/src/transformers/configuration_utils.py b/src/transformers/configuration_utils.py index 2f993e87d4a4..b990b982fcd7 100755 --- a/src/transformers/configuration_utils.py +++ b/src/transformers/configuration_utils.py @@ -216,6 +216,7 @@ 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_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 @@ -1156,6 +1157,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/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index 59aa7ee760e2..045057975607 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -635,20 +635,15 @@ def apply_tensor_parallel(model, tp_mesh, tp_plan): return model if tp_plan == "auto": - enable_sp = getattr(getattr(model.config, "distributed_config", None), "enable_sequence_parallel", False) - sp_plan = getattr(model.config, "base_model_sp_plan", None) - if enable_sp and sp_plan is not None: - base_plan = sp_plan + distributed_config = getattr(model.config, "distributed_config", None) + sp_requested = getattr(distributed_config, "enable_sequence_parallel", False) + sp_supported = getattr(model.config, "base_model_sp_plan", None) is not None + + enable_sp = sp_requested and sp_supported + if enable_sp: + tp_plan = dict(model._sp_plan or {}) else: - base_plan = model.config.base_model_tp_plan or {} - - # Prefix base model keys (e.g. "layers.*.q_proj" → "model.layers.*.q_proj") - # Top-level keys like "lm_head" are kept as-is. - base_model_prefix = model.base_model_prefix - tp_plan = {} - for k, v in base_plan.items(): - is_top_level = hasattr(model, k.split(".")[0]) - tp_plan[k if is_top_level else f"{base_model_prefix}.{k}"] = v + tp_plan = dict(model._tp_plan or {}) parallelize_plan = {} diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 1629c4ca4d9b..2096d98d7275 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -1147,6 +1147,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 @@ -1286,12 +1290,16 @@ 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 = {}, {}, {} + # 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 {} + self._tp_plan, self._sp_plan, self._ep_plan, self._pp_plan = dict(cls_tp_plan), dict(cls_sp_plan), {}, {} # If current model is a base model, attach `base_model_tp_plan` and `base_model_pp_plan` from config if self.base_model is self: self._pp_plan = self.config.base_model_pp_plan.copy() if self.config.base_model_pp_plan is not None else {} self._tp_plan = self.config.base_model_tp_plan.copy() if self.config.base_model_tp_plan is not None else {} + self._sp_plan = self.config.base_model_sp_plan.copy() if self.config.base_model_sp_plan is not None else {} self._ep_plan = self.config.base_model_ep_plan.copy() if self.config.base_model_ep_plan is not None else {} # Current submodel should register its tied weights self.all_tied_weights_keys = self.get_expanded_tied_weights_keys(all_submodels=False) @@ -1309,6 +1317,8 @@ def post_init(self): self._ep_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()}) if plan := getattr(module, "_tp_plan", None): self._tp_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()}) + if plan := getattr(module, "_sp_plan", None): + self._sp_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()}) if plan := getattr(module, "_pp_plan", None): self._pp_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()}) # Always attach the keys of the children (if the children's config says to NOT tie, then it's empty) diff --git a/src/transformers/models/afmoe/modeling_afmoe.py b/src/transformers/models/afmoe/modeling_afmoe.py index 72366f370d90..bf4e39eb5327 100644 --- a/src/transformers/models/afmoe/modeling_afmoe.py +++ b/src/transformers/models/afmoe/modeling_afmoe.py @@ -617,6 +617,7 @@ def forward( class AfmoeForCausalLM(AfmoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/afmoe/modular_afmoe.py b/src/transformers/models/afmoe/modular_afmoe.py index d07a2f1d2017..30cc67787772 100644 --- a/src/transformers/models/afmoe/modular_afmoe.py +++ b/src/transformers/models/afmoe/modular_afmoe.py @@ -397,6 +397,7 @@ def forward( class AfmoeForCausalLM(LlamaForCausalLM, AfmoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("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 aa7ca9fe0676..f116bf8324a7 100644 --- a/src/transformers/models/apertus/configuration_apertus.py +++ b/src/transformers/models/apertus/configuration_apertus.py @@ -70,7 +70,6 @@ class ApertusConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/apertus/modeling_apertus.py b/src/transformers/models/apertus/modeling_apertus.py index 34bac7c18cf7..96c588160a6d 100644 --- a/src/transformers/models/apertus/modeling_apertus.py +++ b/src/transformers/models/apertus/modeling_apertus.py @@ -424,6 +424,7 @@ def forward( class ApertusForCausalLM(ApertusPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/apertus/modular_apertus.py b/src/transformers/models/apertus/modular_apertus.py index 394dbaf7ea98..2da727db820f 100644 --- a/src/transformers/models/apertus/modular_apertus.py +++ b/src/transformers/models/apertus/modular_apertus.py @@ -88,7 +88,6 @@ class ApertusConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/arcee/configuration_arcee.py b/src/transformers/models/arcee/configuration_arcee.py index 0b5dc9c89671..6c2a75ee2da2 100644 --- a/src/transformers/models/arcee/configuration_arcee.py +++ b/src/transformers/models/arcee/configuration_arcee.py @@ -67,7 +67,6 @@ class ArceeConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/arcee/modeling_arcee.py b/src/transformers/models/arcee/modeling_arcee.py index 06916f082039..625b9154e336 100644 --- a/src/transformers/models/arcee/modeling_arcee.py +++ b/src/transformers/models/arcee/modeling_arcee.py @@ -426,6 +426,7 @@ def forward( class ArceeForCausalLM(ArceePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/arcee/modular_arcee.py b/src/transformers/models/arcee/modular_arcee.py index 0c43307747ab..5703ad4e29dd 100644 --- a/src/transformers/models/arcee/modular_arcee.py +++ b/src/transformers/models/arcee/modular_arcee.py @@ -71,7 +71,6 @@ class ArceeConfig(LlamaConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } vocab_size: int = 32000 diff --git a/src/transformers/models/aria/configuration_aria.py b/src/transformers/models/aria/configuration_aria.py index 608760699bd2..bb5a1abb3c62 100644 --- a/src/transformers/models/aria/configuration_aria.py +++ b/src/transformers/models/aria/configuration_aria.py @@ -64,7 +64,6 @@ class AriaTextConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/aria/modeling_aria.py b/src/transformers/models/aria/modeling_aria.py index 7eebaec97e04..d8112acaa870 100644 --- a/src/transformers/models/aria/modeling_aria.py +++ b/src/transformers/models/aria/modeling_aria.py @@ -760,6 +760,7 @@ def forward( class AriaTextForCausalLM(AriaTextPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: AriaTextConfig): diff --git a/src/transformers/models/audioflamingo3/modeling_audioflamingo3.py b/src/transformers/models/audioflamingo3/modeling_audioflamingo3.py index 1fbbc733c308..e9eb4a987579 100644 --- a/src/transformers/models/audioflamingo3/modeling_audioflamingo3.py +++ b/src/transformers/models/audioflamingo3/modeling_audioflamingo3.py @@ -409,6 +409,7 @@ def forward(self, audio_features): class AudioFlamingo3ForConditionalGeneration(AudioFlamingo3PreTrainedModel, GenerationMixin): _keep_in_fp32_modules_strict = None _tp_plan = None + _sp_plan = None _pp_plan = None def __init__(self, config): diff --git a/src/transformers/models/audioflamingo3/modular_audioflamingo3.py b/src/transformers/models/audioflamingo3/modular_audioflamingo3.py index c325bc85300e..0d948ab23694 100644 --- a/src/transformers/models/audioflamingo3/modular_audioflamingo3.py +++ b/src/transformers/models/audioflamingo3/modular_audioflamingo3.py @@ -143,6 +143,7 @@ def __init__(self, config: AudioFlamingo3Config): ) class AudioFlamingo3ForConditionalGeneration(VoxtralForConditionalGeneration): _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 fd63eb7c58f1..fe9b2e95c942 100644 --- a/src/transformers/models/bamba/modeling_bamba.py +++ b/src/transformers/models/bamba/modeling_bamba.py @@ -1073,6 +1073,7 @@ def _update_mamba_mask(self, attention_mask, past_key_values): class BambaForCausalLM(BambaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/bitnet/modeling_bitnet.py b/src/transformers/models/bitnet/modeling_bitnet.py index 14c1581b250f..27c908d204d1 100644 --- a/src/transformers/models/bitnet/modeling_bitnet.py +++ b/src/transformers/models/bitnet/modeling_bitnet.py @@ -423,6 +423,7 @@ def forward( class BitNetForCausalLM(BitNetPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = None + _sp_plan = None _pp_plan = None def __init__(self, config): diff --git a/src/transformers/models/bitnet/modular_bitnet.py b/src/transformers/models/bitnet/modular_bitnet.py index 7f3a248766b2..c8376f131859 100644 --- a/src/transformers/models/bitnet/modular_bitnet.py +++ b/src/transformers/models/bitnet/modular_bitnet.py @@ -110,6 +110,7 @@ class BitNetModel(LlamaModel): class BitNetForCausalLM(LlamaForCausalLM): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = None + _sp_plan = None _pp_plan = None def forward( diff --git a/src/transformers/models/cohere/configuration_cohere.py b/src/transformers/models/cohere/configuration_cohere.py index a365a51e2c8c..b2605904e84d 100644 --- a/src/transformers/models/cohere/configuration_cohere.py +++ b/src/transformers/models/cohere/configuration_cohere.py @@ -75,7 +75,6 @@ class CohereConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/cohere/modeling_cohere.py b/src/transformers/models/cohere/modeling_cohere.py index 80d50905bde0..0315bf78ab3a 100644 --- a/src/transformers/models/cohere/modeling_cohere.py +++ b/src/transformers/models/cohere/modeling_cohere.py @@ -456,6 +456,7 @@ def forward( class CohereForCausalLM(CoherePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/cohere2/configuration_cohere2.py b/src/transformers/models/cohere2/configuration_cohere2.py index 54733a7f297f..9a0aca93cc20 100644 --- a/src/transformers/models/cohere2/configuration_cohere2.py +++ b/src/transformers/models/cohere2/configuration_cohere2.py @@ -72,7 +72,6 @@ class Cohere2Config(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/cohere2/modeling_cohere2.py b/src/transformers/models/cohere2/modeling_cohere2.py index 743031635387..948958759e41 100644 --- a/src/transformers/models/cohere2/modeling_cohere2.py +++ b/src/transformers/models/cohere2/modeling_cohere2.py @@ -435,6 +435,7 @@ def forward( class Cohere2ForCausalLM(Cohere2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/cohere2/modular_cohere2.py b/src/transformers/models/cohere2/modular_cohere2.py index f46085ebaaac..e6272e1e8e16 100644 --- a/src/transformers/models/cohere2/modular_cohere2.py +++ b/src/transformers/models/cohere2/modular_cohere2.py @@ -93,7 +93,6 @@ class Cohere2Config(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/csm/modeling_csm.py b/src/transformers/models/csm/modeling_csm.py index eb78dca8faf5..8e1daac00015 100644 --- a/src/transformers/models/csm/modeling_csm.py +++ b/src/transformers/models/csm/modeling_csm.py @@ -549,6 +549,7 @@ def forward(self, hidden_states, codebook_indices=None): class CsmDepthDecoderForCausalLM(CsmPreTrainedModel, GenerationMixin): _tied_weights_keys = None _tp_plan = None + _sp_plan = None _pp_plan = None def __init__(self, config): diff --git a/src/transformers/models/csm/modular_csm.py b/src/transformers/models/csm/modular_csm.py index 8ba8bc66dad3..e0c8eebab31c 100644 --- a/src/transformers/models/csm/modular_csm.py +++ b/src/transformers/models/csm/modular_csm.py @@ -274,6 +274,7 @@ def forward(self, hidden_states, codebook_indices=None): class CsmDepthDecoderForCausalLM(LlamaForCausalLM, GenerationMixin): _tied_weights_keys = None _tp_plan = None + _sp_plan = None _pp_plan = None def __init__(self, config): diff --git a/src/transformers/models/cwm/configuration_cwm.py b/src/transformers/models/cwm/configuration_cwm.py index 006419da527e..cee303f7b618 100644 --- a/src/transformers/models/cwm/configuration_cwm.py +++ b/src/transformers/models/cwm/configuration_cwm.py @@ -69,7 +69,6 @@ class CwmConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/cwm/modeling_cwm.py b/src/transformers/models/cwm/modeling_cwm.py index 6e60b4ac31da..f3583a78f8ec 100644 --- a/src/transformers/models/cwm/modeling_cwm.py +++ b/src/transformers/models/cwm/modeling_cwm.py @@ -428,6 +428,7 @@ def forward( class CwmForCausalLM(CwmPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/dbrx/modeling_dbrx.py b/src/transformers/models/dbrx/modeling_dbrx.py index 7951f79b334f..db1d86c24a6b 100644 --- a/src/transformers/models/dbrx/modeling_dbrx.py +++ b/src/transformers/models/dbrx/modeling_dbrx.py @@ -644,6 +644,7 @@ def load_balancing_loss_func( class DbrxForCausalLM(DbrxPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "transformer.wte.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("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 34c1b3b6ac5c..e3eabe9fbe08 100644 --- a/src/transformers/models/dbrx/modular_dbrx.py +++ b/src/transformers/models/dbrx/modular_dbrx.py @@ -432,6 +432,7 @@ def forward( class DbrxForCausalLM(DbrxPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "transformer.wte.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: DbrxConfig): diff --git a/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py b/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py index 61b158613355..b0214bd3ed28 100644 --- a/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py +++ b/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py @@ -82,11 +82,13 @@ class DeepseekV2Config(PreTrainedConfig): "layers.*.post_attention_layernorm": TPStyle("activation", "none"), "layers.*.mlp": TPStyle("module", "allgather_split"), "layers.*.mlp.experts": TPStyle( - "moe_experts", "allreduce", + "moe_experts", + "allreduce", shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, ), # Shared experts output must stay Replicate to match experts output - # (summed inside the MoE block, before outer allgather_split handles SP). + # (they're summed inside the MoE block, before the outer allgather_split + # handles the SP boundary). "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), @@ -94,7 +96,6 @@ class DeepseekV2Config(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/deepseek_v2/modeling_deepseek_v2.py b/src/transformers/models/deepseek_v2/modeling_deepseek_v2.py index 672b14742bb7..0c2eacdb5d01 100644 --- a/src/transformers/models/deepseek_v2/modeling_deepseek_v2.py +++ b/src/transformers/models/deepseek_v2/modeling_deepseek_v2.py @@ -543,6 +543,7 @@ def forward( class DeepseekV2ForCausalLM(DeepseekV2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/deepseek_v2/modular_deepseek_v2.py b/src/transformers/models/deepseek_v2/modular_deepseek_v2.py index 9bfceb8eff2c..7e3fc5d7fa25 100644 --- a/src/transformers/models/deepseek_v2/modular_deepseek_v2.py +++ b/src/transformers/models/deepseek_v2/modular_deepseek_v2.py @@ -111,7 +111,6 @@ class DeepseekV2Config(LlamaConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } model_type = "deepseek_v2" diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index fdf708ee9cfa..8c7be7e1e019 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -636,6 +636,7 @@ def forward( class DeepseekV3ForCausalLM(DeepseekV3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/diffllama/modeling_diffllama.py b/src/transformers/models/diffllama/modeling_diffllama.py index 6c5b12cb6850..21e01341050d 100644 --- a/src/transformers/models/diffllama/modeling_diffllama.py +++ b/src/transformers/models/diffllama/modeling_diffllama.py @@ -662,6 +662,7 @@ def forward( class DiffLlamaForCausalLM(DiffLlamaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/doge/modeling_doge.py b/src/transformers/models/doge/modeling_doge.py index ebb6b4bc992c..531a2758548e 100644 --- a/src/transformers/models/doge/modeling_doge.py +++ b/src/transformers/models/doge/modeling_doge.py @@ -718,6 +718,7 @@ def load_balancing_loss_func( class DogeForCausalLM(DogePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/dots1/modeling_dots1.py b/src/transformers/models/dots1/modeling_dots1.py index e77ec0940223..b937fe12c014 100644 --- a/src/transformers/models/dots1/modeling_dots1.py +++ b/src/transformers/models/dots1/modeling_dots1.py @@ -547,6 +547,7 @@ def forward( class Dots1ForCausalLM(Dots1PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/emu3/modeling_emu3.py b/src/transformers/models/emu3/modeling_emu3.py index ec4890f5caa9..b77e735f092f 100644 --- a/src/transformers/models/emu3/modeling_emu3.py +++ b/src/transformers/models/emu3/modeling_emu3.py @@ -1277,6 +1277,7 @@ def forward( class Emu3ForCausalLM(Emu3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} config: Emu3TextConfig diff --git a/src/transformers/models/ernie4_5/configuration_ernie4_5.py b/src/transformers/models/ernie4_5/configuration_ernie4_5.py index 84cdb4d8a745..facfc55d07f0 100644 --- a/src/transformers/models/ernie4_5/configuration_ernie4_5.py +++ b/src/transformers/models/ernie4_5/configuration_ernie4_5.py @@ -70,7 +70,6 @@ class Ernie4_5Config(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/ernie4_5/modeling_ernie4_5.py b/src/transformers/models/ernie4_5/modeling_ernie4_5.py index e8b33e2ade89..86512c389707 100644 --- a/src/transformers/models/ernie4_5/modeling_ernie4_5.py +++ b/src/transformers/models/ernie4_5/modeling_ernie4_5.py @@ -424,6 +424,7 @@ def forward( class Ernie4_5ForCausalLM(Ernie4_5PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): 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 7e2c863211db..e7b3f38181c1 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 @@ -656,6 +656,7 @@ def load_balancing_loss_func( class Ernie4_5_MoeForCausalLM(Ernie4_5_MoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/eurobert/modeling_eurobert.py b/src/transformers/models/eurobert/modeling_eurobert.py index ca6c39acb238..3d32c7f4bae8 100644 --- a/src/transformers/models/eurobert/modeling_eurobert.py +++ b/src/transformers/models/eurobert/modeling_eurobert.py @@ -410,6 +410,7 @@ def forward( class EuroBertForMaskedLM(EuroBertPreTrainedModel): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("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 588508230dfb..1ae4cc7292af 100644 --- a/src/transformers/models/eurobert/modular_eurobert.py +++ b/src/transformers/models/eurobert/modular_eurobert.py @@ -143,6 +143,7 @@ def forward( class EuroBertForMaskedLM(EuroBertPreTrainedModel): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("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 ac746be58090..652939a9c5e0 100644 --- a/src/transformers/models/exaone4/configuration_exaone4.py +++ b/src/transformers/models/exaone4/configuration_exaone4.py @@ -85,7 +85,6 @@ class Exaone4Config(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/exaone4/modeling_exaone4.py b/src/transformers/models/exaone4/modeling_exaone4.py index 0155d7fa3a9f..888e3c091a8b 100644 --- a/src/transformers/models/exaone4/modeling_exaone4.py +++ b/src/transformers/models/exaone4/modeling_exaone4.py @@ -442,6 +442,7 @@ def forward( class Exaone4ForCausalLM(Exaone4PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/exaone4/modular_exaone4.py b/src/transformers/models/exaone4/modular_exaone4.py index 213a9c513e71..5d4926ed33bf 100644 --- a/src/transformers/models/exaone4/modular_exaone4.py +++ b/src/transformers/models/exaone4/modular_exaone4.py @@ -114,7 +114,6 @@ class Exaone4Config(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/exaone_moe/configuration_exaone_moe.py b/src/transformers/models/exaone_moe/configuration_exaone_moe.py index 04e25fbd056e..874490da3422 100644 --- a/src/transformers/models/exaone_moe/configuration_exaone_moe.py +++ b/src/transformers/models/exaone_moe/configuration_exaone_moe.py @@ -75,7 +75,23 @@ class ExaoneMoeConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } - base_model_sp_plan = None + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/exaone_moe/modeling_exaone_moe.py b/src/transformers/models/exaone_moe/modeling_exaone_moe.py index abe193821de9..d5d0ff38fa60 100644 --- a/src/transformers/models/exaone_moe/modeling_exaone_moe.py +++ b/src/transformers/models/exaone_moe/modeling_exaone_moe.py @@ -565,6 +565,7 @@ def forward( class ExaoneMoeForCausalLM(ExaoneMoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/falcon_h1/modeling_falcon_h1.py b/src/transformers/models/falcon_h1/modeling_falcon_h1.py index 8913db392fb1..41520e97518f 100644 --- a/src/transformers/models/falcon_h1/modeling_falcon_h1.py +++ b/src/transformers/models/falcon_h1/modeling_falcon_h1.py @@ -1168,6 +1168,7 @@ def _update_mamba_mask(self, attention_mask, past_key_values): class FalconH1ForCausalLM(FalconH1PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/flex_olmo/modeling_flex_olmo.py b/src/transformers/models/flex_olmo/modeling_flex_olmo.py index 3f38fb9e328c..fa6b44143646 100644 --- a/src/transformers/models/flex_olmo/modeling_flex_olmo.py +++ b/src/transformers/models/flex_olmo/modeling_flex_olmo.py @@ -599,6 +599,7 @@ def load_balancing_loss_func( class FlexOlmoForCausalLM(FlexOlmoPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/gemma/configuration_gemma.py b/src/transformers/models/gemma/configuration_gemma.py index 40567ac76b58..0bcde94f890f 100644 --- a/src/transformers/models/gemma/configuration_gemma.py +++ b/src/transformers/models/gemma/configuration_gemma.py @@ -70,7 +70,6 @@ class GemmaConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/gemma/modeling_gemma.py b/src/transformers/models/gemma/modeling_gemma.py index c26458cfbc54..e2b6b7667667 100644 --- a/src/transformers/models/gemma/modeling_gemma.py +++ b/src/transformers/models/gemma/modeling_gemma.py @@ -452,6 +452,7 @@ def forward( class GemmaForCausalLM(GemmaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/gemma/modular_gemma.py b/src/transformers/models/gemma/modular_gemma.py index 02857b31b2a1..e975c4c96ffd 100644 --- a/src/transformers/models/gemma/modular_gemma.py +++ b/src/transformers/models/gemma/modular_gemma.py @@ -89,7 +89,6 @@ class GemmaConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/gemma2/configuration_gemma2.py b/src/transformers/models/gemma2/configuration_gemma2.py index d0d53315b589..4b840136836d 100644 --- a/src/transformers/models/gemma2/configuration_gemma2.py +++ b/src/transformers/models/gemma2/configuration_gemma2.py @@ -74,7 +74,6 @@ class Gemma2Config(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/gemma2/modeling_gemma2.py b/src/transformers/models/gemma2/modeling_gemma2.py index 83812655fb12..1685a7f8d3bf 100644 --- a/src/transformers/models/gemma2/modeling_gemma2.py +++ b/src/transformers/models/gemma2/modeling_gemma2.py @@ -478,6 +478,7 @@ def forward( class Gemma2ForCausalLM(Gemma2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/gemma2/modular_gemma2.py b/src/transformers/models/gemma2/modular_gemma2.py index a001311eeb6a..54c62724e526 100644 --- a/src/transformers/models/gemma2/modular_gemma2.py +++ b/src/transformers/models/gemma2/modular_gemma2.py @@ -101,7 +101,6 @@ class Gemma2Config(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/gemma3/configuration_gemma3.py b/src/transformers/models/gemma3/configuration_gemma3.py index 24b680a6a3a4..9b9a06d4ae10 100644 --- a/src/transformers/models/gemma3/configuration_gemma3.py +++ b/src/transformers/models/gemma3/configuration_gemma3.py @@ -83,7 +83,6 @@ class Gemma3TextConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/gemma3/modeling_gemma3.py b/src/transformers/models/gemma3/modeling_gemma3.py index 1f8207f8b534..d3f8e5029f7e 100644 --- a/src/transformers/models/gemma3/modeling_gemma3.py +++ b/src/transformers/models/gemma3/modeling_gemma3.py @@ -593,6 +593,7 @@ def forward( class Gemma3ForCausalLM(Gemma3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} config: Gemma3TextConfig diff --git a/src/transformers/models/gemma3/modular_gemma3.py b/src/transformers/models/gemma3/modular_gemma3.py index dd948ee74847..b034bcfe87a5 100644 --- a/src/transformers/models/gemma3/modular_gemma3.py +++ b/src/transformers/models/gemma3/modular_gemma3.py @@ -111,7 +111,6 @@ class Gemma3TextConfig(Gemma2Config, PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } default_theta = {"global": 1_000_000.0, "local": 10_000.0} diff --git a/src/transformers/models/gemma3n/modeling_gemma3n.py b/src/transformers/models/gemma3n/modeling_gemma3n.py index ac502610284f..ca38d342fedb 100644 --- a/src/transformers/models/gemma3n/modeling_gemma3n.py +++ b/src/transformers/models/gemma3n/modeling_gemma3n.py @@ -1771,6 +1771,7 @@ def project_per_layer_inputs( class Gemma3nForCausalLM(Gemma3nPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} config: Gemma3nTextConfig diff --git a/src/transformers/models/gemma4/modeling_gemma4.py b/src/transformers/models/gemma4/modeling_gemma4.py index 17195d973c9c..91cfe090afeb 100644 --- a/src/transformers/models/gemma4/modeling_gemma4.py +++ b/src/transformers/models/gemma4/modeling_gemma4.py @@ -1701,6 +1701,7 @@ def project_per_layer_inputs( class Gemma4ForCausalLM(Gemma4PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} config: Gemma4TextConfig base_model_prefix = "model" diff --git a/src/transformers/models/glm/configuration_glm.py b/src/transformers/models/glm/configuration_glm.py index 494000e3599a..222d1a8dadfa 100644 --- a/src/transformers/models/glm/configuration_glm.py +++ b/src/transformers/models/glm/configuration_glm.py @@ -61,7 +61,6 @@ class GlmConfig(PreTrainedConfig): "layers.*.mlp.gate_up_proj": TPStyle("packed_colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/glm/modeling_glm.py b/src/transformers/models/glm/modeling_glm.py index ddea84722156..084f8b55a12b 100644 --- a/src/transformers/models/glm/modeling_glm.py +++ b/src/transformers/models/glm/modeling_glm.py @@ -441,6 +441,7 @@ def forward( class GlmForCausalLM(GlmPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/glm4/configuration_glm4.py b/src/transformers/models/glm4/configuration_glm4.py index 89328611e2df..e6f51c996b66 100644 --- a/src/transformers/models/glm4/configuration_glm4.py +++ b/src/transformers/models/glm4/configuration_glm4.py @@ -63,7 +63,6 @@ class Glm4Config(PreTrainedConfig): "layers.*.mlp.gate_up_proj": TPStyle("colwise", "allgather"), # fused gate/up needs full tensor for chunk "layers.*.mlp.down_proj": TPStyle("vocab", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/glm4/modeling_glm4.py b/src/transformers/models/glm4/modeling_glm4.py index e896daac86bc..0edde697243d 100644 --- a/src/transformers/models/glm4/modeling_glm4.py +++ b/src/transformers/models/glm4/modeling_glm4.py @@ -446,6 +446,7 @@ def forward( class Glm4ForCausalLM(Glm4PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/glm4_moe/modeling_glm4_moe.py b/src/transformers/models/glm4_moe/modeling_glm4_moe.py index 7275715df0c5..bed0c0153c59 100644 --- a/src/transformers/models/glm4_moe/modeling_glm4_moe.py +++ b/src/transformers/models/glm4_moe/modeling_glm4_moe.py @@ -579,6 +579,7 @@ def forward( class Glm4MoeForCausalLM(Glm4MoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): 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 b7949d728429..1175490d801b 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 @@ -653,6 +653,7 @@ def forward( class Glm4MoeLiteForCausalLM(Glm4MoeLitePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): 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 e88d0f62ed58..037a1b4a6b41 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 @@ -807,6 +807,7 @@ def forward( class GlmMoeDsaForCausalLM(GlmMoeDsaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/glmasr/modeling_glmasr.py b/src/transformers/models/glmasr/modeling_glmasr.py index aff96cad3217..887856bc2102 100644 --- a/src/transformers/models/glmasr/modeling_glmasr.py +++ b/src/transformers/models/glmasr/modeling_glmasr.py @@ -357,6 +357,7 @@ def forward(self, audio_features): class GlmAsrForConditionalGeneration(GlmAsrPreTrainedModel, GenerationMixin): _keep_in_fp32_modules_strict = None _tp_plan = None + _sp_plan = None _pp_plan = None def __init__(self, config): diff --git a/src/transformers/models/gpt_neox/configuration_gpt_neox.py b/src/transformers/models/gpt_neox/configuration_gpt_neox.py index 0b39d857a1b7..83df236f3741 100644 --- a/src/transformers/models/gpt_neox/configuration_gpt_neox.py +++ b/src/transformers/models/gpt_neox/configuration_gpt_neox.py @@ -60,7 +60,6 @@ class GPTNeoXConfig(PreTrainedConfig): "layers.*.mlp.dense_h_to_4h": TPStyle("colwise", "none"), "layers.*.mlp.dense_4h_to_h": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_in": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/gpt_oss/modeling_gpt_oss.py b/src/transformers/models/gpt_oss/modeling_gpt_oss.py index 02527fb44c44..4f39413d4560 100644 --- a/src/transformers/models/gpt_oss/modeling_gpt_oss.py +++ b/src/transformers/models/gpt_oss/modeling_gpt_oss.py @@ -590,6 +590,7 @@ def load_balancing_loss_func( class GptOssForCausalLM(GptOssPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/granite/configuration_granite.py b/src/transformers/models/granite/configuration_granite.py index 64a58e9c3738..fbaaf32ab37f 100644 --- a/src/transformers/models/granite/configuration_granite.py +++ b/src/transformers/models/granite/configuration_granite.py @@ -70,7 +70,6 @@ class GraniteConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/granite/modeling_granite.py b/src/transformers/models/granite/modeling_granite.py index b6062282fe8a..1621b5c86922 100644 --- a/src/transformers/models/granite/modeling_granite.py +++ b/src/transformers/models/granite/modeling_granite.py @@ -447,6 +447,7 @@ def forward( class GraniteForCausalLM(GranitePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/granitemoe/modeling_granitemoe.py b/src/transformers/models/granitemoe/modeling_granitemoe.py index ee63a27d4e11..7a8ab944bdc5 100644 --- a/src/transformers/models/granitemoe/modeling_granitemoe.py +++ b/src/transformers/models/granitemoe/modeling_granitemoe.py @@ -628,6 +628,7 @@ def load_balancing_loss_func( class GraniteMoeForCausalLM(GraniteMoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: GraniteMoeConfig): diff --git a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py index 7441fc47997c..df81d7eea1cd 100644 --- a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py +++ b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py @@ -1308,6 +1308,7 @@ def load_balancing_loss_func( class GraniteMoeHybridForCausalLM(GraniteMoeHybridPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: GraniteMoeHybridConfig): diff --git a/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py b/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py index a698eae304ac..33e9d0bd144f 100644 --- a/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py +++ b/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py @@ -697,6 +697,7 @@ def load_balancing_loss_func( class GraniteMoeSharedForCausalLM(GraniteMoeSharedPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: GraniteMoeSharedConfig): diff --git a/src/transformers/models/helium/configuration_helium.py b/src/transformers/models/helium/configuration_helium.py index cf9d461f1439..8cec52902ceb 100644 --- a/src/transformers/models/helium/configuration_helium.py +++ b/src/transformers/models/helium/configuration_helium.py @@ -64,7 +64,6 @@ class HeliumConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/helium/modeling_helium.py b/src/transformers/models/helium/modeling_helium.py index e3a8de1b8b28..c008cc067365 100644 --- a/src/transformers/models/helium/modeling_helium.py +++ b/src/transformers/models/helium/modeling_helium.py @@ -425,6 +425,7 @@ def forward( class HeliumForCausalLM(HeliumPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, 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 f25435300f7d..bc33ecd6e8c7 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 @@ -82,7 +82,6 @@ class HiggsAudioV2Config(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), 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 481a644d891b..09727a8304b1 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 @@ -463,6 +463,7 @@ def forward( class HunYuanDenseV1ForCausalLM(HunYuanDenseV1PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, 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 75f4033f69f5..73392ca1fdeb 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 @@ -552,6 +552,7 @@ def forward( class HunYuanMoEV1ForCausalLM(HunYuanMoEV1PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/jais2/configuration_jais2.py b/src/transformers/models/jais2/configuration_jais2.py index af46f9a134fb..b5f03d44cc99 100644 --- a/src/transformers/models/jais2/configuration_jais2.py +++ b/src/transformers/models/jais2/configuration_jais2.py @@ -68,7 +68,6 @@ class Jais2Config(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/jais2/modeling_jais2.py b/src/transformers/models/jais2/modeling_jais2.py index ca1de79c6fbb..93316dca79aa 100644 --- a/src/transformers/models/jais2/modeling_jais2.py +++ b/src/transformers/models/jais2/modeling_jais2.py @@ -399,6 +399,7 @@ def forward( class Jais2ForCausalLM(Jais2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/jais2/modular_jais2.py b/src/transformers/models/jais2/modular_jais2.py index 8a1dc6b0dbb4..a5dc6b0302ba 100644 --- a/src/transformers/models/jais2/modular_jais2.py +++ b/src/transformers/models/jais2/modular_jais2.py @@ -52,7 +52,6 @@ class Jais2Config(LlamaConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } vocab_size: int = 150272 diff --git a/src/transformers/models/jamba/modeling_jamba.py b/src/transformers/models/jamba/modeling_jamba.py index 85c00fe0f600..439fdf2f8cfb 100755 --- a/src/transformers/models/jamba/modeling_jamba.py +++ b/src/transformers/models/jamba/modeling_jamba.py @@ -845,6 +845,7 @@ def load_balancing_loss_func( class JambaForCausalLM(JambaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: JambaConfig): 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 8ddd75155b56..6e408304fa28 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 @@ -876,6 +876,7 @@ def forward( class KyutaiSpeechToTextForConditionalGeneration(KyutaiSpeechToTextPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _keep_in_fp32_modules_strict = ["codec_model"] output_modalities = ("audio", "text") diff --git a/src/transformers/models/lfm2/modeling_lfm2.py b/src/transformers/models/lfm2/modeling_lfm2.py index 8071817aa46c..6279867bf54d 100644 --- a/src/transformers/models/lfm2/modeling_lfm2.py +++ b/src/transformers/models/lfm2/modeling_lfm2.py @@ -542,6 +542,7 @@ def forward( class Lfm2ForCausalLM(Lfm2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py b/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py index 4a1935876cca..122816e4c1e3 100644 --- a/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py +++ b/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py @@ -632,6 +632,7 @@ def forward( class Lfm2MoeForCausalLM(Lfm2MoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/llama/configuration_llama.py b/src/transformers/models/llama/configuration_llama.py index 61d7cc76c84e..1f1d378f8e71 100644 --- a/src/transformers/models/llama/configuration_llama.py +++ b/src/transformers/models/llama/configuration_llama.py @@ -70,7 +70,6 @@ class LlamaConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/llama/modeling_llama.py b/src/transformers/models/llama/modeling_llama.py index 0c1aca1edff8..270ef1b0d228 100644 --- a/src/transformers/models/llama/modeling_llama.py +++ b/src/transformers/models/llama/modeling_llama.py @@ -430,6 +430,7 @@ def forward( class LlamaForCausalLM(LlamaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/llama4/modeling_llama4.py b/src/transformers/models/llama4/modeling_llama4.py index 660dddd91bff..6e10d0085057 100644 --- a/src/transformers/models/llama4/modeling_llama4.py +++ b/src/transformers/models/llama4/modeling_llama4.py @@ -592,6 +592,7 @@ class Llama4ForCausalLM(Llama4PreTrainedModel, GenerationMixin): base_model_prefix = "language_model" _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} config: Llama4TextConfig def __init__(self, config: Llama4TextConfig): diff --git a/src/transformers/models/longcat_flash/modeling_longcat_flash.py b/src/transformers/models/longcat_flash/modeling_longcat_flash.py index a1675db42374..8fb78077d9cc 100644 --- a/src/transformers/models/longcat_flash/modeling_longcat_flash.py +++ b/src/transformers/models/longcat_flash/modeling_longcat_flash.py @@ -652,6 +652,7 @@ def forward( class LongcatFlashForCausalLM(LongcatFlashPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _keys_to_ignore_on_load_unexpected = [r"model\.mtp.*"] diff --git a/src/transformers/models/minimax/modeling_minimax.py b/src/transformers/models/minimax/modeling_minimax.py index fde19bce52dd..53a0cd3e6238 100644 --- a/src/transformers/models/minimax/modeling_minimax.py +++ b/src/transformers/models/minimax/modeling_minimax.py @@ -791,6 +791,7 @@ def load_balancing_loss_func( class MiniMaxForCausalLM(MiniMaxPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/minimax_m2/configuration_minimax_m2.py b/src/transformers/models/minimax_m2/configuration_minimax_m2.py index 286af98c7c53..64cb731bc2c9 100644 --- a/src/transformers/models/minimax_m2/configuration_minimax_m2.py +++ b/src/transformers/models/minimax_m2/configuration_minimax_m2.py @@ -75,7 +75,6 @@ class MiniMaxM2Config(PreTrainedConfig): shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, ), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/minimax_m2/modeling_minimax_m2.py b/src/transformers/models/minimax_m2/modeling_minimax_m2.py index 8597f43f40f7..360f50b80bab 100644 --- a/src/transformers/models/minimax_m2/modeling_minimax_m2.py +++ b/src/transformers/models/minimax_m2/modeling_minimax_m2.py @@ -590,6 +590,7 @@ def load_balancing_loss_func( class MiniMaxM2ForCausalLM(MiniMaxM2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/minimax_m2/modular_minimax_m2.py b/src/transformers/models/minimax_m2/modular_minimax_m2.py index 025ae2c86a1c..cf62aae16d7d 100644 --- a/src/transformers/models/minimax_m2/modular_minimax_m2.py +++ b/src/transformers/models/minimax_m2/modular_minimax_m2.py @@ -94,7 +94,6 @@ class MiniMaxM2Config(PreTrainedConfig): shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, ), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/ministral/configuration_ministral.py b/src/transformers/models/ministral/configuration_ministral.py index 2f72a02ceefe..f6002d51ef1a 100644 --- a/src/transformers/models/ministral/configuration_ministral.py +++ b/src/transformers/models/ministral/configuration_ministral.py @@ -72,7 +72,6 @@ class MinistralConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/ministral/modeling_ministral.py b/src/transformers/models/ministral/modeling_ministral.py index 45d3897db33c..3f76f68bdf3a 100644 --- a/src/transformers/models/ministral/modeling_ministral.py +++ b/src/transformers/models/ministral/modeling_ministral.py @@ -431,6 +431,7 @@ def forward( class MinistralForCausalLM(MinistralPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/ministral3/configuration_ministral3.py b/src/transformers/models/ministral3/configuration_ministral3.py index 9c17dce093e1..83dd9d21e75f 100644 --- a/src/transformers/models/ministral3/configuration_ministral3.py +++ b/src/transformers/models/ministral3/configuration_ministral3.py @@ -78,7 +78,6 @@ class Ministral3Config(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/ministral3/modeling_ministral3.py b/src/transformers/models/ministral3/modeling_ministral3.py index e158c819586e..c9c321d2b82a 100644 --- a/src/transformers/models/ministral3/modeling_ministral3.py +++ b/src/transformers/models/ministral3/modeling_ministral3.py @@ -414,6 +414,7 @@ def forward( class Ministral3ForCausalLM(Ministral3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/mistral/configuration_mistral.py b/src/transformers/models/mistral/configuration_mistral.py index 1ef743bd0409..913e6e3c96dc 100644 --- a/src/transformers/models/mistral/configuration_mistral.py +++ b/src/transformers/models/mistral/configuration_mistral.py @@ -69,7 +69,6 @@ class MistralConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/mistral/modeling_mistral.py b/src/transformers/models/mistral/modeling_mistral.py index 523fe2f34c98..6faeb29c8a8e 100644 --- a/src/transformers/models/mistral/modeling_mistral.py +++ b/src/transformers/models/mistral/modeling_mistral.py @@ -403,6 +403,7 @@ def forward( class MistralForCausalLM(MistralPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/mistral4/modeling_mistral4.py b/src/transformers/models/mistral4/modeling_mistral4.py index 15f4e9f43e98..dbc9adc77fa7 100644 --- a/src/transformers/models/mistral4/modeling_mistral4.py +++ b/src/transformers/models/mistral4/modeling_mistral4.py @@ -642,6 +642,7 @@ def forward( class Mistral4ForCausalLM(Mistral4PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/mixtral/configuration_mixtral.py b/src/transformers/models/mixtral/configuration_mixtral.py index 9e4e22ef770b..e5e5294a0116 100644 --- a/src/transformers/models/mixtral/configuration_mixtral.py +++ b/src/transformers/models/mixtral/configuration_mixtral.py @@ -79,7 +79,6 @@ class MixtralConfig(PreTrainedConfig): }, ), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/mixtral/modeling_mixtral.py b/src/transformers/models/mixtral/modeling_mixtral.py index 1965f7ee8da4..c755959639a5 100644 --- a/src/transformers/models/mixtral/modeling_mixtral.py +++ b/src/transformers/models/mixtral/modeling_mixtral.py @@ -582,6 +582,7 @@ def load_balancing_loss_func( class MixtralForCausalLM(MixtralPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/musicflamingo/modeling_musicflamingo.py b/src/transformers/models/musicflamingo/modeling_musicflamingo.py index adec95bbf3e1..f21905516092 100644 --- a/src/transformers/models/musicflamingo/modeling_musicflamingo.py +++ b/src/transformers/models/musicflamingo/modeling_musicflamingo.py @@ -201,6 +201,7 @@ def apply_rotary_time_emb(hidden_states, cos, sin): class MusicFlamingoForConditionalGeneration(MusicFlamingoPreTrainedModel, GenerationMixin): _keep_in_fp32_modules_strict = None _tp_plan = None + _sp_plan = None _pp_plan = None def __init__(self, config: MusicFlamingoConfig): diff --git a/src/transformers/models/nanochat/modeling_nanochat.py b/src/transformers/models/nanochat/modeling_nanochat.py index b14faff93e67..5431b5492da9 100644 --- a/src/transformers/models/nanochat/modeling_nanochat.py +++ b/src/transformers/models/nanochat/modeling_nanochat.py @@ -434,6 +434,7 @@ def forward( class NanoChatForCausalLM(NanoChatPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/nanochat/modular_nanochat.py b/src/transformers/models/nanochat/modular_nanochat.py index 486ec255089e..b70ab836b62d 100644 --- a/src/transformers/models/nanochat/modular_nanochat.py +++ b/src/transformers/models/nanochat/modular_nanochat.py @@ -200,6 +200,7 @@ def forward( @auto_docstring class NanoChatForCausalLM(Gemma2ForCausalLM): _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} def forward(self, **super_kwargs) -> CausalLMOutputWithPast: r""" diff --git a/src/transformers/models/olmo/configuration_olmo.py b/src/transformers/models/olmo/configuration_olmo.py index ee2a8fae11ef..f67cc37f8bd5 100644 --- a/src/transformers/models/olmo/configuration_olmo.py +++ b/src/transformers/models/olmo/configuration_olmo.py @@ -73,7 +73,6 @@ class OlmoConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/olmo/modeling_olmo.py b/src/transformers/models/olmo/modeling_olmo.py index 47a7e2d78edf..cab9bb3b3e1a 100644 --- a/src/transformers/models/olmo/modeling_olmo.py +++ b/src/transformers/models/olmo/modeling_olmo.py @@ -427,6 +427,7 @@ def forward( class OlmoForCausalLM(OlmoPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/olmo2/configuration_olmo2.py b/src/transformers/models/olmo2/configuration_olmo2.py index c6861d4e3e4f..8b6b745812b3 100644 --- a/src/transformers/models/olmo2/configuration_olmo2.py +++ b/src/transformers/models/olmo2/configuration_olmo2.py @@ -86,7 +86,6 @@ class Olmo2Config(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/olmo2/modeling_olmo2.py b/src/transformers/models/olmo2/modeling_olmo2.py index 502d800ce288..c319c4a3478a 100644 --- a/src/transformers/models/olmo2/modeling_olmo2.py +++ b/src/transformers/models/olmo2/modeling_olmo2.py @@ -431,6 +431,7 @@ def forward( class Olmo2ForCausalLM(Olmo2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/olmo2/modular_olmo2.py b/src/transformers/models/olmo2/modular_olmo2.py index 9b74f4cd0154..a4aac76c44d6 100644 --- a/src/transformers/models/olmo2/modular_olmo2.py +++ b/src/transformers/models/olmo2/modular_olmo2.py @@ -99,7 +99,6 @@ class Olmo2Config(OlmoConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } 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 5675ff02aef1..a317fcdb80c9 100644 --- a/src/transformers/models/olmo3/configuration_olmo3.py +++ b/src/transformers/models/olmo3/configuration_olmo3.py @@ -81,7 +81,6 @@ class Olmo3Config(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/olmo3/modeling_olmo3.py b/src/transformers/models/olmo3/modeling_olmo3.py index 04dbdb27b1c2..daf8db42d9d1 100644 --- a/src/transformers/models/olmo3/modeling_olmo3.py +++ b/src/transformers/models/olmo3/modeling_olmo3.py @@ -435,6 +435,7 @@ def forward( class Olmo3ForCausalLM(Olmo3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/olmo3/modular_olmo3.py b/src/transformers/models/olmo3/modular_olmo3.py index 7eb5e105235d..aa70895ab944 100644 --- a/src/transformers/models/olmo3/modular_olmo3.py +++ b/src/transformers/models/olmo3/modular_olmo3.py @@ -95,7 +95,6 @@ class Olmo3Config(Olmo2Config): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } 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 5a5f2b1d8a33..c78c9bc548b7 100644 --- a/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py +++ b/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py @@ -90,7 +90,21 @@ class OlmoHybridConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } - base_model_sp_plan = None + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py index 092cfb94e6c5..788c73eca593 100644 --- a/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py +++ b/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py @@ -1034,6 +1034,7 @@ def _update_linear_attn_mask(self, attention_mask, past_key_values): class OlmoHybridForCausalLM(OlmoHybridPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/olmoe/configuration_olmoe.py b/src/transformers/models/olmoe/configuration_olmoe.py index 99ed6c9668e3..4bef0fdec27b 100644 --- a/src/transformers/models/olmoe/configuration_olmoe.py +++ b/src/transformers/models/olmoe/configuration_olmoe.py @@ -73,7 +73,6 @@ class OlmoeConfig(PreTrainedConfig): shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, ), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } vocab_size: int = 50304 diff --git a/src/transformers/models/olmoe/modeling_olmoe.py b/src/transformers/models/olmoe/modeling_olmoe.py index 6d648ff9b879..2b9ce11116a2 100644 --- a/src/transformers/models/olmoe/modeling_olmoe.py +++ b/src/transformers/models/olmoe/modeling_olmoe.py @@ -606,6 +606,7 @@ def load_balancing_loss_func( class OlmoeForCausalLM(OlmoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py b/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py index 3993cd6f0a87..cc2257fabef2 100644 --- a/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py +++ b/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py @@ -118,7 +118,6 @@ class PaddleOCRTextConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/phi/configuration_phi.py b/src/transformers/models/phi/configuration_phi.py index 91b846bb5c74..99432092ed4f 100644 --- a/src/transformers/models/phi/configuration_phi.py +++ b/src/transformers/models/phi/configuration_phi.py @@ -67,7 +67,6 @@ class PhiConfig(PreTrainedConfig): "layers.*.mlp.fc1": TPStyle("colwise", "none"), "layers.*.mlp.fc2": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/phi/modeling_phi.py b/src/transformers/models/phi/modeling_phi.py index 700fc1ddfd38..ed33a66731ed 100644 --- a/src/transformers/models/phi/modeling_phi.py +++ b/src/transformers/models/phi/modeling_phi.py @@ -408,6 +408,7 @@ def forward( class PhiForCausalLM(PhiPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/phi3/configuration_phi3.py b/src/transformers/models/phi3/configuration_phi3.py index b236a280d390..38691f62356a 100644 --- a/src/transformers/models/phi3/configuration_phi3.py +++ b/src/transformers/models/phi3/configuration_phi3.py @@ -68,7 +68,6 @@ class Phi3Config(PreTrainedConfig): "layers.*.mlp.gate_up_proj": TPStyle("colwise", "allgather"), # fused gate/up needs full tensor for chunk "layers.*.mlp.down_proj": TPStyle("vocab", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/phi3/modeling_phi3.py b/src/transformers/models/phi3/modeling_phi3.py index e537449ffd56..a549f703b647 100644 --- a/src/transformers/models/phi3/modeling_phi3.py +++ b/src/transformers/models/phi3/modeling_phi3.py @@ -434,6 +434,7 @@ def forward( class Phi3ForCausalLM(Phi3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py b/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py index dc66237fbf0a..03a06bce7094 100644 --- a/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py +++ b/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py @@ -202,7 +202,6 @@ class Phi4MultimodalConfig(PreTrainedConfig): "layers.*.mlp.gate_up_proj": TPStyle("colwise", "allgather"), # fused gate/up needs full tensor for chunk "layers.*.mlp.down_proj": TPStyle("vocab", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py b/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py index 22d2886f6e17..abd1e9b09e7e 100644 --- a/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py +++ b/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py @@ -1598,6 +1598,7 @@ def forward( class Phi4MultimodalForCausalLM(Phi4MultimodalPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/phimoe/modeling_phimoe.py b/src/transformers/models/phimoe/modeling_phimoe.py index 2846fd2aefb0..cc76274c379a 100644 --- a/src/transformers/models/phimoe/modeling_phimoe.py +++ b/src/transformers/models/phimoe/modeling_phimoe.py @@ -774,6 +774,7 @@ def load_balancing_loss_func( class PhimoeForCausalLM(PhimoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/qwen2/configuration_qwen2.py b/src/transformers/models/qwen2/configuration_qwen2.py index 633275345908..dc16b48b56f6 100644 --- a/src/transformers/models/qwen2/configuration_qwen2.py +++ b/src/transformers/models/qwen2/configuration_qwen2.py @@ -67,7 +67,6 @@ class Qwen2Config(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/qwen2/modeling_qwen2.py b/src/transformers/models/qwen2/modeling_qwen2.py index d4e8d1f59d8c..b8bcf9d9cdac 100644 --- a/src/transformers/models/qwen2/modeling_qwen2.py +++ b/src/transformers/models/qwen2/modeling_qwen2.py @@ -418,6 +418,7 @@ def forward( class Qwen2ForCausalLM(Qwen2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py b/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py index 542dddd48fe4..7d3a1d2c5209 100644 --- a/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py +++ b/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py @@ -619,6 +619,7 @@ def load_balancing_loss_func( class Qwen2MoeForCausalLM(Qwen2MoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/qwen2_moe/modular_qwen2_moe.py b/src/transformers/models/qwen2_moe/modular_qwen2_moe.py index eb624179143b..d351dfa764e6 100644 --- a/src/transformers/models/qwen2_moe/modular_qwen2_moe.py +++ b/src/transformers/models/qwen2_moe/modular_qwen2_moe.py @@ -232,6 +232,7 @@ def forward( class Qwen2MoeForCausalLM(MixtralForCausalLM, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/qwen3/configuration_qwen3.py b/src/transformers/models/qwen3/configuration_qwen3.py index 372922204d3e..5aa1288e66a8 100644 --- a/src/transformers/models/qwen3/configuration_qwen3.py +++ b/src/transformers/models/qwen3/configuration_qwen3.py @@ -77,7 +77,6 @@ class Qwen3Config(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/qwen3/modeling_qwen3.py b/src/transformers/models/qwen3/modeling_qwen3.py index beeab1982123..d8b2bf9973d1 100644 --- a/src/transformers/models/qwen3/modeling_qwen3.py +++ b/src/transformers/models/qwen3/modeling_qwen3.py @@ -424,6 +424,7 @@ def forward( class Qwen3ForCausalLM(Qwen3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/qwen3_5/modeling_qwen3_5.py b/src/transformers/models/qwen3_5/modeling_qwen3_5.py index fd4f46095111..06a13e099c21 100644 --- a/src/transformers/models/qwen3_5/modeling_qwen3_5.py +++ b/src/transformers/models/qwen3_5/modeling_qwen3_5.py @@ -1689,6 +1689,7 @@ def forward( class Qwen3_5ForCausalLM(Qwen3_5PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} config: Qwen3_5TextConfig _keys_to_ignore_on_load_unexpected = [r"^mtp.*", r"^model.visual.*"] 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 3dd042042ffe..f3afa5bdc169 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 @@ -1896,6 +1896,7 @@ def load_balancing_loss_func( class Qwen3_5MoeForCausalLM(Qwen3_5MoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} config: Qwen3_5MoeTextConfig _keys_to_ignore_on_load_unexpected = [r"^mtp.*", r"^model.visual.*"] @@ -2003,6 +2004,7 @@ class Qwen3_5MoeForConditionalGeneration(Qwen3_5MoePreTrainedModel, GenerationMi accepts_loss_kwargs = False config: Qwen3_5MoeConfig _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("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 7f0baa834584..428cd54248d9 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 @@ -250,6 +250,7 @@ def __init__(self, config): class Qwen3_5MoeForConditionalGeneration(Qwen3VLMoeForConditionalGeneration): _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("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 763de442c937..f8f057b00ea8 100644 --- a/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py +++ b/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py @@ -89,7 +89,6 @@ class Qwen3MoeConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py b/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py index ceec4a562b7d..e101d4f78273 100644 --- a/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py +++ b/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py @@ -587,6 +587,7 @@ def load_balancing_loss_func( class Qwen3MoeForCausalLM(Qwen3MoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/qwen3_next/modeling_qwen3_next.py b/src/transformers/models/qwen3_next/modeling_qwen3_next.py index a991c45fedac..c6ae090dfa3e 100644 --- a/src/transformers/models/qwen3_next/modeling_qwen3_next.py +++ b/src/transformers/models/qwen3_next/modeling_qwen3_next.py @@ -1074,6 +1074,7 @@ def load_balancing_loss_func( class Qwen3NextForCausalLM(Qwen3NextPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, 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 67a2cfd84577..bc6d86ae7edb 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 @@ -292,7 +292,6 @@ class Qwen3OmniMoeTalkerCodePredictorConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -407,7 +406,6 @@ class Qwen3OmniMoeTalkerTextConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), 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 b3e30d387246..e1a53909071d 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 @@ -2622,6 +2622,7 @@ def get_input_embeddings(self): class Qwen3OmniMoeTalkerCodePredictorModelForConditionalGeneration(Qwen3OmniMoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} config_class = Qwen3OmniMoeTalkerCodePredictorConfig base_model_prefix = "talker.code_predictor" @@ -3006,6 +3007,7 @@ def get_input_embeddings(self): class Qwen3OmniMoeTalkerForConditionalGeneration(Qwen3OmniMoeThinkerTextPreTrainedModel, GenerationMixin): _tied_weights_keys = {"codec_head": "model.codec_embedding.weight"} _tp_plan = {"codec_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"codec_head": (["hidden_states"], ["logits"])} config_class = Qwen3OmniMoeTalkerConfig base_model_prefix = "talker" diff --git a/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py b/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py index 73678ee8c736..43946f06bf1e 100644 --- a/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py +++ b/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py @@ -31,7 +31,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func +from ...integrations import use_kernel_forward_from_hub, use_kernelized_func from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -122,7 +122,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: def rotate_half(x): - """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) @@ -406,25 +405,7 @@ def extra_repr(self): return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" -@use_kernel_func_from_hub("rotary_pos_emb") def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): - """Applies Rotary Position Embedding to the query and key tensors. - - Args: - q (`torch.Tensor`): The query tensor. - k (`torch.Tensor`): The key tensor. - cos (`torch.Tensor`): The cosine part of the rotary embedding. - sin (`torch.Tensor`): The sine part of the rotary embedding. - unsqueeze_dim (`int`, *optional*, defaults to 1): - The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and - sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note - that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and - k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes - cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have - the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. - Returns: - `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. - """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) q_embed = (q * cos) + (rotate_half(q) * sin) 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 867f0d0a22d9..68fec77811af 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 @@ -86,7 +86,6 @@ class Qwen3VLMoeTextConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py b/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py index ce405683fc94..679a4357e305 100644 --- a/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py +++ b/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py @@ -31,12 +31,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import ( - use_experts_implementation, - use_kernel_forward_from_hub, - use_kernel_func_from_hub, - use_kernelized_func, -) +from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -145,7 +140,6 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: def rotate_half(x): - """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) @@ -188,25 +182,7 @@ def eager_attention_forward( return attn_output, attn_weights -@use_kernel_func_from_hub("rotary_pos_emb") def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): - """Applies Rotary Position Embedding to the query and key tensors. - - Args: - q (`torch.Tensor`): The query tensor. - k (`torch.Tensor`): The key tensor. - cos (`torch.Tensor`): The cosine part of the rotary embedding. - sin (`torch.Tensor`): The sine part of the rotary embedding. - unsqueeze_dim (`int`, *optional*, defaults to 1): - The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and - sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note - that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and - k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes - cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have - the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. - Returns: - `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. - """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) q_embed = (q * cos) + (rotate_half(q) * sin) diff --git a/src/transformers/models/seed_oss/configuration_seed_oss.py b/src/transformers/models/seed_oss/configuration_seed_oss.py index 72f948e5b1cc..5902cbac4ae7 100644 --- a/src/transformers/models/seed_oss/configuration_seed_oss.py +++ b/src/transformers/models/seed_oss/configuration_seed_oss.py @@ -68,7 +68,6 @@ class SeedOssConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/seed_oss/modeling_seed_oss.py b/src/transformers/models/seed_oss/modeling_seed_oss.py index 5c144edaa94c..b37084599b94 100644 --- a/src/transformers/models/seed_oss/modeling_seed_oss.py +++ b/src/transformers/models/seed_oss/modeling_seed_oss.py @@ -431,6 +431,7 @@ def forward( class SeedOssForCausalLM(SeedOssPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/smollm3/configuration_smollm3.py b/src/transformers/models/smollm3/configuration_smollm3.py index 3c958c865543..9b5a113b7c2f 100644 --- a/src/transformers/models/smollm3/configuration_smollm3.py +++ b/src/transformers/models/smollm3/configuration_smollm3.py @@ -78,7 +78,6 @@ class SmolLM3Config(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/smollm3/modeling_smollm3.py b/src/transformers/models/smollm3/modeling_smollm3.py index 82ad64435bd8..9623ef48a659 100644 --- a/src/transformers/models/smollm3/modeling_smollm3.py +++ b/src/transformers/models/smollm3/modeling_smollm3.py @@ -447,6 +447,7 @@ def forward( class SmolLM3ForCausalLM(SmolLM3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/smollm3/modular_smollm3.py b/src/transformers/models/smollm3/modular_smollm3.py index 0917a4db0446..89f0c813f3f9 100644 --- a/src/transformers/models/smollm3/modular_smollm3.py +++ b/src/transformers/models/smollm3/modular_smollm3.py @@ -94,7 +94,6 @@ class SmolLM3Config(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/solar_open/configuration_solar_open.py b/src/transformers/models/solar_open/configuration_solar_open.py index 65bfd946f5ea..e0e316dbd604 100644 --- a/src/transformers/models/solar_open/configuration_solar_open.py +++ b/src/transformers/models/solar_open/configuration_solar_open.py @@ -49,24 +49,6 @@ class SolarOpenConfig(PreTrainedConfig): shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, ), } - base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather_split"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), - "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), - } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), @@ -102,6 +84,23 @@ 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": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather_split"), + "layers.*.mlp.experts": TPStyle( + "moe_experts", + "allreduce", + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), + "norm": TPStyle("activation", "none"), + } 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 cb3f35fb3700..ac114dbbb82b 100644 --- a/src/transformers/models/solar_open/modeling_solar_open.py +++ b/src/transformers/models/solar_open/modeling_solar_open.py @@ -554,6 +554,7 @@ def forward( class SolarOpenForCausalLM(SolarOpenPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/solar_open/modular_solar_open.py b/src/transformers/models/solar_open/modular_solar_open.py index f81f4bbff3fd..dcf766bbb7aa 100644 --- a/src/transformers/models/solar_open/modular_solar_open.py +++ b/src/transformers/models/solar_open/modular_solar_open.py @@ -71,7 +71,6 @@ class SolarOpenConfig(Glm4MoeConfig): shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, ), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } vocab_size: int = 196608 diff --git a/src/transformers/models/starcoder2/configuration_starcoder2.py b/src/transformers/models/starcoder2/configuration_starcoder2.py index 1b7a2879fe3c..45c349840963 100644 --- a/src/transformers/models/starcoder2/configuration_starcoder2.py +++ b/src/transformers/models/starcoder2/configuration_starcoder2.py @@ -66,7 +66,6 @@ class Starcoder2Config(PreTrainedConfig): "layers.*.mlp.c_fc": TPStyle("colwise", "none"), "layers.*.mlp.c_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/starcoder2/modeling_starcoder2.py b/src/transformers/models/starcoder2/modeling_starcoder2.py index 26fcc4e8b435..cbfb88c2db48 100644 --- a/src/transformers/models/starcoder2/modeling_starcoder2.py +++ b/src/transformers/models/starcoder2/modeling_starcoder2.py @@ -411,6 +411,7 @@ def forward( class Starcoder2ForCausalLM(Starcoder2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/t5gemma/configuration_t5gemma.py b/src/transformers/models/t5gemma/configuration_t5gemma.py index b4b39ba7d743..ba4e2282cb70 100644 --- a/src/transformers/models/t5gemma/configuration_t5gemma.py +++ b/src/transformers/models/t5gemma/configuration_t5gemma.py @@ -74,7 +74,6 @@ class T5GemmaModuleConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/t5gemma2/configuration_t5gemma2.py b/src/transformers/models/t5gemma2/configuration_t5gemma2.py index 86525e4f090d..38f88e669360 100644 --- a/src/transformers/models/t5gemma2/configuration_t5gemma2.py +++ b/src/transformers/models/t5gemma2/configuration_t5gemma2.py @@ -70,7 +70,6 @@ class T5Gemma2TextConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -261,7 +260,6 @@ class T5Gemma2DecoderConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/vaultgemma/configuration_vaultgemma.py b/src/transformers/models/vaultgemma/configuration_vaultgemma.py index 5bcec13cb160..3b535d51bfa9 100644 --- a/src/transformers/models/vaultgemma/configuration_vaultgemma.py +++ b/src/transformers/models/vaultgemma/configuration_vaultgemma.py @@ -73,7 +73,6 @@ class VaultGemmaConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), "norm": TPStyle("activation", "none"), - "lm_head": TPStyle("colwise", "loss_parallel"), } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/vaultgemma/modeling_vaultgemma.py b/src/transformers/models/vaultgemma/modeling_vaultgemma.py index 60122bfcb923..901d585660e8 100644 --- a/src/transformers/models/vaultgemma/modeling_vaultgemma.py +++ b/src/transformers/models/vaultgemma/modeling_vaultgemma.py @@ -469,6 +469,7 @@ def forward( class VaultGemmaForCausalLM(VaultGemmaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/vibevoice_asr/modeling_vibevoice_asr.py b/src/transformers/models/vibevoice_asr/modeling_vibevoice_asr.py index 703bb6ca5130..768bce98f0d0 100644 --- a/src/transformers/models/vibevoice_asr/modeling_vibevoice_asr.py +++ b/src/transformers/models/vibevoice_asr/modeling_vibevoice_asr.py @@ -257,6 +257,7 @@ def _init_weights(self, module): class VibeVoiceAsrForConditionalGeneration(VibeVoiceAsrPreTrainedModel, GenerationMixin): _keep_in_fp32_modules_strict = None _tp_plan = None + _sp_plan = None _pp_plan = None def __init__(self, config: VibeVoiceAsrConfig): diff --git a/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py b/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py index 1dafcb67da56..633aa22edb87 100644 --- a/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py +++ b/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py @@ -825,6 +825,7 @@ def forward( class VoxtralRealtimeTextForCausalLM(VoxtralRealtimeTextPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/youtu/modeling_youtu.py b/src/transformers/models/youtu/modeling_youtu.py index 07e7c639f979..c9ea544decf2 100644 --- a/src/transformers/models/youtu/modeling_youtu.py +++ b/src/transformers/models/youtu/modeling_youtu.py @@ -535,6 +535,7 @@ def forward( class YoutuForCausalLM(YoutuPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} + _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): From 3ff1fee069bdf7c8574fd41a50fdcfbe6d24182b Mon Sep 17 00:00:00 2001 From: 3outeille Date: Mon, 20 Apr 2026 15:25:38 +0000 Subject: [PATCH 035/116] first pass of refactoring dtensor shard operator --- src/transformers/core_model_loading.py | 443 ++++++++++++++----------- tests/utils/test_core_model_loading.py | 80 ++--- 2 files changed, 283 insertions(+), 240 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index bd33512a0296..7871b825f355 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -855,230 +855,273 @@ def _job(): class DtensorShardOperation: - """Extracts the local shard from a full checkpoint tensor based on this rank's DTensor placements.""" - - # tensor_dim -> list of (start, end) index ranges this rank owns - DimRanges = dict[int, list[tuple[int, int]]] + """Read only this rank's local shard out of a checkpoint tensor. + + Input (source) + A torch.Tensor or safetensors slice. Its shape is one of: + (a) the full parameter (source.shape == param.shape); + (b) one expert of a packed MoE param (param.shape = (E, *rest), + source.shape = rest, tensor_idx identifies which expert); + (c) one half of a packed fused weight (e.g. w1 alone, before it + is concatenated with w3 into gate_up_proj). + Cases (b) and (c) can co-occur (MoE expert stored as pre-pack w1/w3). + + Output + A torch.Tensor holding only the slice(s) this rank owns, moved to + device/dtype. Returns None when this rank does not own the expert + (case b) and the expert axis is sharded. + + DTensor placements (one per mesh dim; mesh is e.g. [FSDP, TP]) + Replicate — every rank along this mesh axis has the full axis. + Shard(d) — cut tensor dim d into contiguous chunks along this + mesh axis. + _StridedShard(d, split_factor=k) + — cut tensor dim d into k groups first, then + contiguously shard within each group. Each rank's + slice on dim d is therefore k disjoint intervals + (concatenated after read). This is how a packed + axis like [Q | K | V] gets sharded so every rank + owns half of Q *and* half of K *and* half of V. + + Algorithm + For each source dim, build a list of (start, end) intervals this + rank owns by folding the placements: + - Replicate → no change. + - Shard → one contiguous sub-interval per interval. + - _StridedShard → split_factor disjoint sub-intervals. + Then slice source with those intervals; if one dim has multiple + intervals, concatenate the pieces along it. + + Shape mismatches + - Expert source (param.ndim == source.ndim + 1, tensor_idx set): + if the expert axis is sharded, return None when unowned; else + drop the expert placement and continue. + - Pre-pack source (same ndim mismatch): the packed axis does not exist + in the source yet, so _StridedShard on it degrades to a plain + contiguous cut — the WeightConverter recreates packing later. + + Fallback + When Shard and _StridedShard share a tensor dim, interval arithmetic + cannot express the reorder: materialize the full tensor and call + placement._split_tensor per mesh dim instead. + """ def __init__(self, param: DTensor): self.device_mesh = param.device_mesh - self.param = param self.placements = tuple(param.placements) + self.param_shape = tuple(param.shape) + self.param_ndim = param.ndim local_shape, _ = compute_local_shape_and_global_offset(param.shape, self.device_mesh, self.placements) self.local_shape = tuple(local_shape) def shard_tensor( - self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None + self, source: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None ) -> torch.Tensor | None: - """Return the local shard of ``param`` for this rank, dispatching to the appropriate strategy.""" - # Find sharding placements (keep Shard and _StridedShard only) - # _StridedShard.is_shard() returns False in PyTorch, so we also check for - # the ``dim`` attribute that both Shard and _StridedShard have. - sharding_placements = [ - (i, p) - for i, p in enumerate(self.placements) + source_shape = list(source.shape) if isinstance(source, torch.Tensor) else source.get_shape() + + # ------------------------------------------------------------------ + # What is the source? + # (a) full weight source.ndim == param.ndim + # (b) one expert source.ndim == param.ndim - 1, leading expert axis dropped + # (c) one half of a pack source is smaller on the packed axis (w1 before gate_up concat) + # (b) and (c) can co-occur (MoE expert stored as pre-pack w1/w3). + # ------------------------------------------------------------------ + source_missing_leading_axis = self.param_ndim > len(source_shape) + + # ------------------------------------------------------------------ + # Collect placements that actually split a dim. + # _StridedShard.is_shard() returns False in PyTorch, so also accept + # any non-Replicate placement that exposes a `dim` attribute. + # ------------------------------------------------------------------ + placements = [ + (mesh_dim, p) + for mesh_dim, p in enumerate(self.placements) if p.is_shard() or (hasattr(p, "dim") and not p.is_replicate()) ] - param_shape = list(param.shape) if isinstance(param, torch.Tensor) else param.get_shape() - - # [A] No sharding placements -> Return full copy - if not sharding_placements: - return param[...].to(device=device, dtype=dtype) - - if tensor_idx is not None and len(self.param.shape) == len(param_shape) + 1: - # [B] Expert path: shard on expert dimension (dim 0). - # When dim 0 is the only sharding placement, return the full expert or - # skip it. When TP also shards an inner dim, keep applying the remaining - # placements to the owned expert tensor. - has_expert_sharding = any(self._normalize_param_dim(p.dim) == 0 for _, p in sharding_placements) - if has_expert_sharding: - # [B2] This rank doesn't own the expert tensor -> skip it - if not self._owns_local_expert(tensor_idx): - return None - inner_placements = [(i, p) for i, p in sharding_placements if self._normalize_param_dim(p.dim) != 0] - # [B3] Not composed with TP placements -> return full copy of the expert tensor - if not inner_placements: - return param[...].to(device=device, dtype=dtype) - # [B4] Composed with TP placements -> shard the expert's inner dims - return self._shard_nd(param, inner_placements, param_shape, device, dtype) - - # [B1] has_expert_sharding=False -> fall through to _shard_nd - return self._shard_nd(param, sharding_placements, param_shape, device, dtype) - - def _shard_nd(self, param, sharding_placements, param_shape, device, dtype): - """Handle multi-dimensional sharding, choosing the best strategy.""" - # [C1] Column Parallel when composed with FSDP. We choose the easier path but maybe we should do a better one? - if not self._can_shard_on_read(sharding_placements): - return self._materialize_and_split(param, sharding_placements, device, dtype) - - # [C2] All sharding placements are plain Shard on different dims -> single contiguous slice - has_strided = any(not p.is_shard() for _, p in sharding_placements) - if not has_strided: - local_shape, global_offset = compute_local_shape_and_global_offset( - self.param.shape, self.device_mesh, self.placements - ) - slices = [slice(None)] * len(param_shape) - for _, placement in sharding_placements: - dim = self._checkpoint_dim(placement.dim, param_shape) - offset = global_offset[placement.dim] - slices[dim] = slice(offset, offset + local_shape[placement.dim]) - return param[tuple(slices)].to(device=device, dtype=dtype) - - # [C3] At least one _StridedShard (no same-dim conflict) -> _compute_dim_ranges + _slice_and_read - dim_ranges = self._compute_dim_ranges(sharding_placements, param_shape) - return self._slice_and_read(param, param_shape, dim_ranges, device, dtype) - - def _can_shard_on_read(self, sharding_placements) -> bool: - """Check whether range-based shard-on-read is feasible. - - Returns ``False`` when a ``_StridedShard`` and another placement share the - same tensor dimension — the strided reorder can't be composed via range - arithmetic because ``Shard`` would need to cut across the concatenated - result of ``_StridedShard``'s disjoint ranges. - """ - dims_seen: dict[int, bool] = {} # dim -> has_strided - for _, placement in sharding_placements: - dim = placement.dim - is_strided = not placement.is_shard() - if dim in dims_seen and (is_strided or dims_seen[dim]): - logger.debug( - "Cannot shard-on-read: dim %d has both Shard and _StridedShard placements, " - "falling back to materialize-then-split.", - dim, + if not placements: + return source[...].to(device=device, dtype=dtype) # no sharding → full copy + + # ------------------------------------------------------------------ + # Case (b): resolve the expert axis up front. + # + # The expert axis is param dim 0; it does not exist in the source + # (the source is one single expert, named by tensor_idx). Example: + # param = (E=4, H=8, I=4) with placements [Shard(0), Shard(1)] + # source = (H=8, I=4) for experts.2.w1.weight, tensor_idx=2 + # + # Shard(0) on the expert axis is really asking "which experts do I + # own?" — it splits experts across ranks, not values inside one + # expert. _owns_expert(tensor_idx) answers that question: + # - not owned → this whole file is for other ranks → return None + # - owned → Shard(0) is now fully handled; remove it from + # `placements` so we don't try to slice it again. + # + # Remaining placements (here: Shard(1)) are inner — they still + # need to shard source's dims. We fall through to the generic loop. + # ------------------------------------------------------------------ + source_is_one_expert = tensor_idx is not None and self.param_ndim == len(source_shape) + 1 + if source_is_one_expert and any(self._norm_dim(p.dim) == 0 for _, p in placements): + if not self._owns_expert(tensor_idx): + return None + placements = [(mesh_dim, p) for mesh_dim, p in placements if self._norm_dim(p.dim) != 0] + if not placements: + # Expert axis was the only sharding → keep the whole expert tensor. + return source[...].to(device=device, dtype=dtype) + + # ------------------------------------------------------------------ + # Cases (a) and (c): generic interval loop. + # (a) source IS the full weight → placement.dim maps 1:1 to a + # source dim, interval math applies directly. + # (c) source is one pack-half (same path as (a), but interleaved + # _StridedShard degrades to contiguous — handled inside the + # loop because the packed axis does not exist in source yet). + # First, a conflict check: Shard + _StridedShard on the same source + # dim can't be composed via interval math (the strided reorder + # would be cut across by Shard). Fall back to load-then-split. + # ------------------------------------------------------------------ + shard_dims: set[int] = set() + strided_dims: set[int] = set() + for _, p in placements: + source_dim = self._source_dim(p.dim, source_shape) + (shard_dims if p.is_shard() else strided_dims).add(source_dim) + if shard_dims & strided_dims: + return self._materialize_and_split(source, placements, device, dtype) + + # ------------------------------------------------------------------ + # For each placement, narrow the intervals on its source dim: + # Shard → one contiguous sub-interval per existing one + # _StridedShard → split_factor disjoint sub-intervals + # Source dims that nobody shards stay at [(0, size)]. Two placements + # on the same dim fold naturally (nested cut) because each call + # narrows the list produced by the previous one. + # ------------------------------------------------------------------ + intervals: list[list[tuple[int, int]]] = [[(0, size)] for size in source_shape] + for mesh_dim, placement in placements: + source_dim = self._source_dim(placement.dim, source_shape) + sub_mesh = self._get_sub_mesh(mesh_dim) + rank, world_size = sub_mesh.get_local_rank(), sub_mesh.size() + + # Case (c): the packed axis doesn't exist in source yet, so an + # interleaved (_StridedShard) placement has no groups to reorder + # and degrades to a plain contiguous cut. The WeightConverter + # recreates the packed layout later. + is_interleaved = not placement.is_shard() and not source_missing_leading_axis + if is_interleaved: + intervals[source_dim] = _strided_intervals( + intervals[source_dim], rank, world_size, placement.split_factor ) - return False - dims_seen[dim] = is_strided - return True - - def _materialize_and_split(self, param, sharding_placements, device, dtype): - """Fallback: load the full tensor, then iteratively split per mesh dim.""" - tensor = param[...] if not isinstance(param, torch.Tensor) else param - for mesh_dim_idx, placement in sharding_placements: - sub_mesh = self._get_sub_mesh(mesh_dim_idx) - rank = sub_mesh.get_local_rank() - shards, _ = placement._split_tensor(tensor, sub_mesh.size(), with_padding=False, contiguous=True) - tensor = shards[rank] - return tensor.to(device=device, dtype=dtype) + else: + intervals[source_dim] = _contiguous_intervals(intervals[source_dim], rank, world_size) + + return self._slice_and_cat(source, intervals, device, dtype) - def _compute_dim_ranges(self, sharding_placements, param_shape) -> DtensorShardOperation.DimRanges: - """Compute per-dimension index ranges for this rank. + def _slice_and_cat(self, source, intervals, device, dtype): + """Read `source` with per-source-dim intervals; concat along the sole multi-interval dim if any. - Each sharding placement narrows the ranges on its tensor dimension: - - ``Shard``: one contiguous sub-range per previous range. - - ``_StridedShard``: multiple disjoint sub-ranges (one per split-factor group). + Each entry of ``intervals`` holds the (start, end) pieces this rank + owns on that source dim. At most one dim may have more than one + piece (from _StridedShard); two such dims would require a 2D outer + product of reads and are rejected. """ - dim_ranges: DtensorShardOperation.DimRanges = {} - for mesh_dim_idx, placement in sharding_placements: - sub_mesh = self._get_sub_mesh(mesh_dim_idx) - rank = sub_mesh.get_local_rank() - world_size = sub_mesh.size() - dim = self._checkpoint_dim(placement.dim, param_shape) - prev_ranges = dim_ranges.get(dim, [(0, param_shape[dim])]) - - if placement.is_shard(): - new_ranges = self._contiguous_ranges(prev_ranges, rank, world_size) - elif self._source_tensor_needs_packing(param_shape): - # [C3a] _StridedShard only makes sense once the packed axis exists. While - # loading pre-packed source tensors (e.g. w1/w3 before gate_up_proj - # concatenation), take the contiguous chunk for this rank and let the - # WeightConverter recreate the packed layout afterward. - new_ranges = self._contiguous_ranges(prev_ranges, rank, world_size) - else: - # [C3b] Normal strided -> disjoint ranges + cat - new_ranges = self._strided_ranges(prev_ranges, rank, world_size, placement.split_factor) - dim_ranges[dim] = new_ranges - return dim_ranges + 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 _materialize_and_split(self, source, placements, device, dtype): + """Fallback: load the full tensor, split it once per mesh dim using each placement's own rule.""" + tensor = source if isinstance(source, torch.Tensor) else source[...] + for mesh_dim, placement in placements: + sub_mesh = self._get_sub_mesh(mesh_dim) + shards, _ = placement._split_tensor(tensor, sub_mesh.size(), with_padding=False, contiguous=True) + tensor = shards[sub_mesh.get_local_rank()] + return tensor.to(device=device, dtype=dtype) - def _slice_and_read(self, param, param_shape, dim_ranges: DtensorShardOperation.DimRanges, device, dtype): - """Build slices from computed ranges and read from the tensor. + def _owns_expert(self, expert_idx: int) -> bool: + """True when this rank's shard of the expert axis (param dim 0) contains expert_idx.""" + _, offsets = compute_local_shape_and_global_offset( + torch.Size(self.param_shape), self.device_mesh, self.placements + ) + first_owned_expert = offsets[0] + return first_owned_expert <= expert_idx < first_owned_expert + self.local_shape[0] - At most one dim can have multiple disjoint ranges (from ``_StridedShard``). - If so, read each disjoint range separately and concatenate. - """ - concat_dim = None - concat_ranges = None - base_slices = [slice(None)] * len(param_shape) - for dim, ranges in dim_ranges.items(): - if len(ranges) == 1: - base_slices[dim] = slice(ranges[0][0], ranges[0][1]) - elif len(ranges) > 1: - if concat_dim is not None: - raise ValueError("Shard-on-read only supports disjoint ranges on a single checkpoint dimension.") - concat_dim = dim - concat_ranges = ranges - - if concat_dim is None: - return param[tuple(base_slices)].to(device=device, dtype=dtype) - - pieces = [] - for start, end in concat_ranges: - slices = list(base_slices) - slices[concat_dim] = slice(start, end) - pieces.append(param[tuple(slices)]) - return torch.cat(pieces, dim=concat_dim).to(device=device, dtype=dtype) - - # ------------------------------------------------------------------ - # Utilities - # ------------------------------------------------------------------ - - def _contiguous_ranges( - self, prev_ranges: list[tuple[int, int]], rank: int, world_size: int - ) -> list[tuple[int, int]]: - """Narrow each range by picking the ``rank``-th contiguous chunk (``Shard`` semantics).""" - new_ranges = [] - for prev_start, prev_end in prev_ranges: - shard_size, offset = Shard.local_shard_size_and_offset(prev_end - prev_start, world_size, rank) - if shard_size > 0: - new_ranges.append((prev_start + offset, prev_start + offset + shard_size)) - return new_ranges - - def _strided_ranges( - self, prev_ranges: list[tuple[int, int]], rank: int, world_size: int, split_factor: int - ) -> list[tuple[int, int]]: - """Narrow each range using ``_StridedShard`` semantics. - - Divides each range into ``split_factor`` groups, then within each group - picks the ``rank``-th chunk of ``world_size`` equal pieces. - """ - new_ranges = [] - for prev_start, prev_end in prev_ranges: - group_size = math.ceil((prev_end - prev_start) / split_factor) - for g in range(split_factor): - g_start = prev_start + g * group_size - g_end = min(g_start + group_size, prev_end) - if g_end <= g_start: - continue - shard_size, offset = Shard.local_shard_size_and_offset(g_end - g_start, world_size, rank) - if shard_size > 0: - new_ranges.append((g_start + offset, g_start + offset + shard_size)) - return new_ranges - - def _get_sub_mesh(self, mesh_dim_idx: int): - """Return the 1-D sub-mesh for ``mesh_dim_idx``.""" + def _get_sub_mesh(self, mesh_dim: int): + """1-D sub-mesh along a single mesh axis (e.g. just the TP axis).""" if self.device_mesh.ndim > 1: - return self.device_mesh[self.device_mesh.mesh_dim_names[mesh_dim_idx]] + return self.device_mesh[self.device_mesh.mesh_dim_names[mesh_dim]] return self.device_mesh - def _normalize_param_dim(self, dim: int) -> int: - return dim if dim >= 0 else self.param.ndim + dim + def _norm_dim(self, dim: int) -> int: + """Normalize a (possibly negative) param dim to a non-negative index.""" + return dim if dim >= 0 else self.param_ndim + dim + + def _source_dim(self, placement_dim: int, source_shape) -> int: + """Map a placement's param dim onto the corresponding source-tensor dim. - def _checkpoint_dim(self, placement_dim: int, param_shape) -> int: - """Map a placement dim from the DTensor shape to the checkpoint tensor shape.""" - dim = self._normalize_param_dim(placement_dim) - ndim_diff = self.param.ndim - len(param_shape) - if ndim_diff > 0 and dim >= ndim_diff: - dim -= ndim_diff + When the source has fewer dims than the param (expert or pre-pack + source), leading axes are absent: any placement dim past the missing + prefix shifts down by the difference. + """ + dim = self._norm_dim(placement_dim) + missing_leading_dims = self.param_ndim - len(source_shape) + if missing_leading_dims > 0 and dim >= missing_leading_dims: + dim -= missing_leading_dims return dim - def _owns_local_expert(self, tensor_idx: int) -> bool: - _, offsets = compute_local_shape_and_global_offset(self.param.shape, self.device_mesh, self.placements) - return offsets[0] <= tensor_idx < offsets[0] + self.local_shape[0] - def _source_tensor_needs_packing(self, param_shape) -> bool: - # A single source tensor still missing the leading expert axis is being - # converted into a packed expert parameter. In that case _StridedShard's - # split groups do not exist yet. - return self.param.ndim == len(param_shape) + 1 +def _contiguous_intervals( + intervals: list[tuple[int, int]], rank: int, world_size: int +) -> list[tuple[int, int]]: + """Narrow each interval to this rank's contiguous sub-interval (Shard semantics).""" + narrowed = [] + for start, end in intervals: + size, offset = Shard.local_shard_size_and_offset(end - start, world_size, rank) + if size > 0: + narrowed.append((start + offset, start + offset + size)) + return narrowed + + +def _strided_intervals( + intervals: list[tuple[int, int]], rank: int, world_size: int, split_factor: int +) -> list[tuple[int, int]]: + """Split each interval into `split_factor` groups, then shard contiguously within each group. + + Produces one sub-interval per group (up to `split_factor` disjoint pieces + per input interval), matching _StridedShard's packed-axis layout. + """ + 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 dot_natural_key(s: str): diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py index d68ad2621437..c61c7413d31f 100644 --- a/tests/utils/test_core_model_loading.py +++ b/tests/utils/test_core_model_loading.py @@ -265,9 +265,8 @@ def _make_dtensor_shard_op(mesh, placements, param_shape, local_shape): op = object.__new__(DtensorShardOperation) op.device_mesh = mesh op.placements = tuple(placements) - ns = SimpleNamespace(shape=torch.Size(param_shape), ndim=len(param_shape)) - ns.dim = lambda: len(param_shape) - op.param = ns + op.param_shape = tuple(param_shape) + op.param_ndim = len(param_shape) op.local_shape = tuple(local_shape) return op @@ -876,31 +875,32 @@ def test_ernie4_5_vl_moe_conversion_reversed(self): class TestDtensorShardOperation(unittest.TestCase): - """Unit tests for DtensorShardOperation.shard_tensor — one test per code path. - - Branch coverage map (labels [A]–[C3b] match comments in core_model_loading.py): - - shard_tensor() - ├── [A] no sharding placements → full copy [test_no_shard_returns_full_tensor] - ├── [B] expert path (tensor_idx set, ndim mismatch) - │ ├── [B1] has_expert_sharding=False → fall through to C [test_expert_shaped_tp_only_no_expert_sharding] - │ ├── [B2] not owns_local_expert → None [test_expert_filtering] - │ ├── [B3] owned, no inner placements → full copy [test_expert_filtering] - │ └── [B4] owned, with inner placements → _shard_nd [test_expert_filtering_preserves_inner_sharding] - └── [C] _shard_nd() - ├── [C1] _can_shard_on_read=False → _materialize_and_split [test_nd_strided_plus_shard_same_dim_fallback] - ├── [C2] has_strided=False → contiguous slice - │ ├── 1D mesh [test_1d_shard_fast_path] - │ ├── 2D mesh [test_nd_contiguous_single_slice] - │ ├── negative dim [test_negative_dim_normalizes_correctly] - │ └── uneven division [test_contiguous_shard_uneven_division] - └── [C3] has_strided=True → _compute_dim_ranges + _slice_and_read - ├── [C3a] _source_tensor_needs_packing → contiguous [test_prepacked_strided_shard_uses_contiguous_source_slice] - └── [C3b] _StridedShard → _strided_ranges [test_nd_strided_shard_disjoint_ranges] - - _slice_and_read (tested directly) - ├── all single ranges → simple slice [test_slice_and_read_all_single_ranges] - └── two multi-range dims → ValueError [test_slice_and_read_raises_on_two_multi_range_dims] + """Unit tests for DtensorShardOperation.shard_tensor. + + The class handles three source shapes relative to the DTensor param: + (a) full weight — interval math applies directly. + (b) one expert of a packed param — skip if not owned, else drop the + expert placement and continue. + (c) one pack-half — interleaved (_StridedShard) on the + missing packed axis degrades to a + plain contiguous cut. + A Shard + _StridedShard on the same source dim falls back to + materialize-then-split. + + Covered paths: + - no sharding / replicate-only [test_no_shard_returns_full_tensor] + - single-placement contiguous shard [test_1d_single_shard, + test_negative_dim_normalizes_correctly, + test_contiguous_shard_uneven_division] + - multi-placement contiguous shard [test_nd_contiguous_single_slice] + - strided shard producing disjoint intervals [test_nd_strided_shard_disjoint_ranges] + - same-dim conflict → materialize+split fallback [test_nd_strided_plus_shard_same_dim_fallback] + - per-expert, no expert-axis sharding [test_expert_shaped_tp_only_no_expert_sharding] + - per-expert, expert axis sharded (skip/own) [test_expert_filtering] + - per-expert + inner TP [test_expert_filtering_preserves_inner_sharding] + - per-expert + pre-pack axis (strided→contig) [test_prepacked_strided_shard_uses_contiguous_source_slice] + - internal _slice_and_cat [test_slice_and_cat_all_single_ranges, + test_slice_and_cat_raises_on_two_multi_range_dims] """ def test_no_shard_returns_full_tensor(self): @@ -910,8 +910,8 @@ def test_no_shard_returns_full_tensor(self): tensor = torch.arange(16).reshape(4, 4).float() torch.testing.assert_close(op.shard_tensor(tensor), tensor) - def test_1d_shard_fast_path(self): - # TODO(3outeille): double check fast path + def test_1d_single_shard(self): + """1D mesh with a single Shard(0) → contiguous split across ranks.""" tensor = torch.arange(16).reshape(4, 4).float() for rank, expected in [(0, tensor[:2]), (1, tensor[2:])]: mesh = FakeMesh(shape=(2,), rank=rank) @@ -975,9 +975,9 @@ def test_prepacked_strided_shard_uses_contiguous_source_slice(self): torch.testing.assert_close(op.shard_tensor(tensor, tensor_idx=0), expected, msg=f"rank {rank}") def test_expert_shaped_tp_only_no_expert_sharding(self): - """Expert-shaped param with TP on dim 1 but no expert sharding on dim 0 → regular _shard_nd path.""" + """Expert-shaped param with TP on dim 1 but no expert sharding on dim 0 → plain interval loop.""" tensor = torch.arange(8).reshape(4, 2).float() - # Shard(1) on 3D param maps to dim 0 of the 2D checkpoint tensor (ndim_diff=1) + # Shard(1) on a 3D param maps to source dim 0 (source is missing the leading expert axis). for rank, expected in [(0, tensor[:2]), (1, tensor[2:])]: mesh = FakeMesh(shape=(2,), rank=rank) op = _make_dtensor_shard_op(mesh, [Shard(1)], param_shape=(4, 4, 2), local_shape=(4, 2, 2)) @@ -1028,23 +1028,23 @@ def test_contiguous_shard_uneven_division(self): op = _make_dtensor_shard_op(mesh, [Shard(0)], param_shape=(5, 4), local_shape=(local_rows, 4)) torch.testing.assert_close(op.shard_tensor(tensor), expected[rank], msg=f"rank {rank}") - def test_slice_and_read_all_single_ranges(self): - """When every dim has exactly one range, _slice_and_read takes the simple slice path (no concat).""" + def test_slice_and_cat_all_single_ranges(self): + """When every dim has exactly one interval, _slice_and_cat does a single slice read (no concat).""" 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, 4)) - dim_ranges = {0: [(0, 4)], 1: [(2, 6)]} - result = op._slice_and_read(tensor, [8, 8], dim_ranges, None, None) + intervals = [[(0, 4)], [(2, 6)]] + result = op._slice_and_cat(tensor, intervals, None, None) torch.testing.assert_close(result, tensor[0:4, 2:6]) - def test_slice_and_read_raises_on_two_multi_range_dims(self): - """Multiple disjoint ranges on two different dims → ValueError.""" + def test_slice_and_cat_raises_on_two_multi_range_dims(self): + """Multiple disjoint intervals on two different dims → ValueError.""" 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, 4)) - dim_ranges = {0: [(0, 2), (4, 6)], 1: [(0, 2), (4, 6)]} + intervals = [[(0, 2), (4, 6)], [(0, 2), (4, 6)]] with self.assertRaises(ValueError): - op._slice_and_read(tensor, [8, 8], dim_ranges, None, None) + op._slice_and_cat(tensor, intervals, None, None) class TestConversionMapping(unittest.TestCase): From 4d96b2dc4212e60b7eb2493bb5a02be443ef847f Mon Sep 17 00:00:00 2001 From: 3outeille Date: Mon, 20 Apr 2026 17:25:14 +0000 Subject: [PATCH 036/116] better refacto --- src/transformers/core_model_loading.py | 174 ++++++++++++++----------- tests/utils/test_core_model_loading.py | 10 +- 2 files changed, 104 insertions(+), 80 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 7871b825f355..7e8b2a780691 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -887,7 +887,10 @@ class DtensorShardOperation: For each source dim, build a list of (start, end) intervals this rank owns by folding the placements: - Replicate → no change. - - Shard → one contiguous sub-interval per interval. + - Shard → treat the current list as one flat logical + sequence and take this rank's contiguous + chunk (may span multiple sub-intervals when + a prior _StridedShard left disjoint pieces). - _StridedShard → split_factor disjoint sub-intervals. Then slice source with those intervals; if one dim has multiple intervals, concatenate the pieces along it. @@ -899,11 +902,6 @@ class DtensorShardOperation: - Pre-pack source (same ndim mismatch): the packed axis does not exist in the source yet, so _StridedShard on it degrades to a plain contiguous cut — the WeightConverter recreates packing later. - - Fallback - When Shard and _StridedShard share a tensor dim, interval arithmetic - cannot express the reorder: materialize the full tensor and call - placement._split_tensor per mesh dim instead. """ def __init__(self, param: DTensor): @@ -970,30 +968,41 @@ def shard_tensor( # ------------------------------------------------------------------ # Cases (a) and (c): generic interval loop. - # (a) source IS the full weight → placement.dim maps 1:1 to a - # source dim, interval math applies directly. - # (c) source is one pack-half (same path as (a), but interleaved - # _StridedShard degrades to contiguous — handled inside the - # loop because the packed axis does not exist in source yet). - # First, a conflict check: Shard + _StridedShard on the same source - # dim can't be composed via interval math (the strided reorder - # would be cut across by Shard). Fall back to load-then-split. - # ------------------------------------------------------------------ - shard_dims: set[int] = set() - strided_dims: set[int] = set() - for _, p in placements: - source_dim = self._source_dim(p.dim, source_shape) - (shard_dims if p.is_shard() else strided_dims).add(source_dim) - if shard_dims & strided_dims: - return self._materialize_and_split(source, placements, device, dtype) - - # ------------------------------------------------------------------ - # For each placement, narrow the intervals on its source dim: - # Shard → one contiguous sub-interval per existing one - # _StridedShard → split_factor disjoint sub-intervals - # Source dims that nobody shards stay at [(0, size)]. Two placements - # on the same dim fold naturally (nested cut) because each call - # narrows the list produced by the previous one. + # + # Build, per source dim, the list of (start, end) index ranges this + # rank owns. We start with "the whole axis" on every dim and then + # apply each placement, which narrows the list on its source dim. + # + # Shard(d) → treat the current list on dim d as one flat + # logical sequence and take this rank's + # contiguous chunk of it. If the prior result + # was a single interval, the output is still + # one interval (standard Shard). If it was + # multiple disjoint intervals (left behind by + # a _StridedShard on the same dim), the chunk + # may span several of them — handled by + # _contiguous_intervals' flat-walk logic. + # _StridedShard → cut each (start, end) into split_factor groups + # and keep this rank's contiguous sub-range of + # each group. List length grows by split_factor; + # _slice_and_cat will read the pieces and + # concatenate them. + # + # Example — 2×2 mesh [FSDP, TP], param (8, 8), rank coord (0, 0), + # placements [Shard(0), Shard(1)]: + # initial → dim 0: [(0, 8)] dim 1: [(0, 8)] + # after Shard(0)→ dim 0: [(0, 4)] dim 1: [(0, 8)] + # after Shard(1)→ dim 0: [(0, 4)] dim 1: [(0, 4)] + # slice result → source[0:4, 0:4] + # + # Example — [_StridedShard(0, sf=2), Shard(0)] on size-4 dim, 2×2 + # mesh. _StridedShard leaves rank 0 with [(0, 1), (2, 3)] (rows 0 + # and 2 as a flat 2-row sequence). Shard then takes rank 0's half + # of that flat sequence → [(0, 1)]; rank 1 gets [(2, 3)]. + # + # Two placements on the same tensor dim (e.g. [Shard(0), Shard(0)] + # when two mesh dims both split rows — 2-level FSDP) fold naturally: + # the second call narrows the list produced by the first. # ------------------------------------------------------------------ intervals: list[list[tuple[int, int]]] = [[(0, size)] for size in source_shape] for mesh_dim, placement in placements: @@ -1001,17 +1010,19 @@ def shard_tensor( sub_mesh = self._get_sub_mesh(mesh_dim) rank, world_size = sub_mesh.get_local_rank(), sub_mesh.size() - # Case (c): the packed axis doesn't exist in source yet, so an - # interleaved (_StridedShard) placement has no groups to reorder - # and degrades to a plain contiguous cut. The WeightConverter - # recreates the packed layout later. + # Case (c): _StridedShard only makes sense when the packed axis + # actually exists in the source. If the source is a pre-pack + # half (w1 alone, before gate_up concat), the split_factor + # groups haven't been laid out yet — we just take a plain + # contiguous cut on this half, and the WeightConverter will + # rebuild the packed layout afterwards. is_interleaved = not placement.is_shard() and not source_missing_leading_axis if is_interleaved: - intervals[source_dim] = _strided_intervals( + intervals[source_dim] = self._strided_intervals( intervals[source_dim], rank, world_size, placement.split_factor ) else: - intervals[source_dim] = _contiguous_intervals(intervals[source_dim], rank, world_size) + intervals[source_dim] = self._contiguous_intervals(intervals[source_dim], rank, world_size) return self._slice_and_cat(source, intervals, device, dtype) @@ -1049,15 +1060,6 @@ def _slice_and_cat(self, source, intervals, device, dtype): pieces_read.append(source[tuple(piece_slices)]) return torch.cat(pieces_read, dim=multi_interval_dim).to(device=device, dtype=dtype) - def _materialize_and_split(self, source, placements, device, dtype): - """Fallback: load the full tensor, split it once per mesh dim using each placement's own rule.""" - tensor = source if isinstance(source, torch.Tensor) else source[...] - for mesh_dim, placement in placements: - sub_mesh = self._get_sub_mesh(mesh_dim) - shards, _ = placement._split_tensor(tensor, sub_mesh.size(), with_padding=False, contiguous=True) - tensor = shards[sub_mesh.get_local_rank()] - return tensor.to(device=device, dtype=dtype) - def _owns_expert(self, expert_idx: int) -> bool: """True when this rank's shard of the expert axis (param dim 0) contains expert_idx.""" _, offsets = compute_local_shape_and_global_offset( @@ -1090,38 +1092,60 @@ def _source_dim(self, placement_dim: int, source_shape) -> int: return dim -def _contiguous_intervals( - intervals: list[tuple[int, int]], rank: int, world_size: int -) -> list[tuple[int, int]]: - """Narrow each interval to this rank's contiguous sub-interval (Shard semantics).""" - narrowed = [] - for start, end in intervals: - size, offset = Shard.local_shard_size_and_offset(end - start, world_size, rank) - if size > 0: - narrowed.append((start + offset, start + offset + size)) - return narrowed - + def _contiguous_intervals( + self, intervals: list[tuple[int, int]], rank: int, world_size: int + ) -> list[tuple[int, int]]: + """Shard semantics: treat `intervals` as one flat logical sequence and take this rank's contiguous chunk of it. -def _strided_intervals( - intervals: list[tuple[int, int]], rank: int, world_size: int, split_factor: int -) -> list[tuple[int, int]]: - """Split each interval into `split_factor` groups, then shard contiguously within each group. - - Produces one sub-interval per group (up to `split_factor` disjoint pieces - per input interval), matching _StridedShard's packed-axis layout. - """ - 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: + For a single input interval this is a plain per-interval cut. For + multiple intervals (left by a prior _StridedShard on the same dim), + the chunk can cross range boundaries — walk the intervals and emit + the sub-ranges that fall inside [my_offset, my_offset + my_size) in + the flat view. + """ + 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 - 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 + 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 _strided_intervals( + self, intervals: list[tuple[int, int]], rank: int, world_size: int, split_factor: int + ) -> list[tuple[int, int]]: + """Split each interval into `split_factor` groups, then shard contiguously within each group. + + Produces one sub-interval per group (up to `split_factor` disjoint pieces + per input interval), matching _StridedShard's packed-axis layout. + """ + 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 dot_natural_key(s: str): diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py index c61c7413d31f..750dd04a63b4 100644 --- a/tests/utils/test_core_model_loading.py +++ b/tests/utils/test_core_model_loading.py @@ -884,8 +884,8 @@ class TestDtensorShardOperation(unittest.TestCase): (c) one pack-half — interleaved (_StridedShard) on the missing packed axis degrades to a plain contiguous cut. - A Shard + _StridedShard on the same source dim falls back to - materialize-then-split. + Shard after _StridedShard on the same dim composes via a flat-view cut + in _contiguous_intervals. Covered paths: - no sharding / replicate-only [test_no_shard_returns_full_tensor] @@ -894,7 +894,7 @@ class TestDtensorShardOperation(unittest.TestCase): test_contiguous_shard_uneven_division] - multi-placement contiguous shard [test_nd_contiguous_single_slice] - strided shard producing disjoint intervals [test_nd_strided_shard_disjoint_ranges] - - same-dim conflict → materialize+split fallback [test_nd_strided_plus_shard_same_dim_fallback] + - Shard after _StridedShard on same dim [test_nd_strided_plus_shard_same_dim] - per-expert, no expert-axis sharding [test_expert_shaped_tp_only_no_expert_sharding] - per-expert, expert axis sharded (skip/own) [test_expert_filtering] - per-expert + inner TP [test_expert_filtering_preserves_inner_sharding] @@ -947,8 +947,8 @@ def test_nd_strided_shard_disjoint_ranges(self): ) torch.testing.assert_close(op.shard_tensor(tensor), expected[rank], msg=f"rank {rank}") - def test_nd_strided_plus_shard_same_dim_fallback(self): - """_StridedShard + Shard on same dim → materialize-then-split fallback.""" + def test_nd_strided_plus_shard_same_dim(self): + """_StridedShard + Shard on the same dim: Shard takes a flat-view slice of the strided output.""" tensor = torch.arange(16).reshape(4, 4).float() expected = {0: tensor[[0]], 1: tensor[[2]], 2: tensor[[1]], 3: tensor[[3]]} for rank in range(4): From 04521bfe0d37e78e228a4e3d40187f58eed6e1e7 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 22 Apr 2026 10:51:18 +0000 Subject: [PATCH 037/116] batter explanation of DtensorShardOperation --- src/transformers/core_model_loading.py | 351 +++++++++++++------------ 1 file changed, 176 insertions(+), 175 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 7e8b2a780691..2450b574d0af 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -856,52 +856,47 @@ def _job(): class DtensorShardOperation: """Read only this rank's local shard out of a checkpoint tensor. - - Input (source) - A torch.Tensor or safetensors slice. Its shape is one of: - (a) the full parameter (source.shape == param.shape); - (b) one expert of a packed MoE param (param.shape = (E, *rest), - source.shape = rest, tensor_idx identifies which expert); - (c) one half of a packed fused weight (e.g. w1 alone, before it - is concatenated with w3 into gate_up_proj). - Cases (b) and (c) can co-occur (MoE expert stored as pre-pack w1/w3). - - Output - A torch.Tensor holding only the slice(s) this rank owns, moved to - device/dtype. Returns None when this rank does not own the expert - (case b) and the expert axis is sharded. - - DTensor placements (one per mesh dim; mesh is e.g. [FSDP, TP]) - Replicate — every rank along this mesh axis has the full axis. - Shard(d) — cut tensor dim d into contiguous chunks along this - mesh axis. - _StridedShard(d, split_factor=k) - — cut tensor dim d into k groups first, then - contiguously shard within each group. Each rank's - slice on dim d is therefore k disjoint intervals - (concatenated after read). This is how a packed - axis like [Q | K | V] gets sharded so every rank - owns half of Q *and* half of K *and* half of V. - - Algorithm - For each source dim, build a list of (start, end) intervals this - rank owns by folding the placements: - - Replicate → no change. - - Shard → treat the current list as one flat logical - sequence and take this rank's contiguous - chunk (may span multiple sub-intervals when - a prior _StridedShard left disjoint pieces). - - _StridedShard → split_factor disjoint sub-intervals. - Then slice source with those intervals; if one dim has multiple - intervals, concatenate the pieces along it. - - Shape mismatches - - Expert source (param.ndim == source.ndim + 1, tensor_idx set): - if the expert axis is sharded, return None when unowned; else - drop the expert placement and continue. - - Pre-pack source (same ndim mismatch): the packed axis does not exist - in the source yet, so _StridedShard on it degrades to a plain - contiguous cut — the WeightConverter recreates packing later. + We first need to classify the source tensor (checkpoint tensor) relative to + the destination `param` (the Dtensor sharded parameter we are loading into). + This decides how we must slice it before writing into the local + DTensor shard. + + (a) Full weight — source.ndim == param.ndim + The checkpoint tensor has the same rank as the model param. + Example: param = (8, 8), source = (8, 8). + → Apply placements directly via the generic interval loop. + + (b) One expert — source.ndim == param.ndim - 1 + MoE models stack experts along a leading axis (E, ...) in the + model, but checkpoints store each expert in its own file + (e.g. `experts.2.w1.weight`). The source is missing that + leading expert axis; `tensor_idx` names which expert it is. + Example: param = (E=4, H=8, I=4), source = (H=8, I=4), + tensor_idx = 2. + → If Shard(0) (the expert axis) is in placements, first decide + whether this rank owns expert `tensor_idx`: + - not owned: return None (file belongs to other ranks); + - owned: drop Shard(0) from placements and fall through + to the generic loop for the remaining inner + dims (e.g. Shard(1) on H). + + (c) One half of a pack — source smaller than param on the packed axis + Some weights are concatenated at load time (gate+up → gate_up, + Q+K+V → qkv). The checkpoint stores each half separately + (e.g. `w1` before concat with `w3`), so source is smaller than + param along the packed axis. + Example: param = (2H, D), source = (H, D) for just `w1`. + → _StridedShard on the packed axis cannot stride over a pack + that does not exist yet, so it degrades to a plain contiguous + cut here. The WeightConverter re-creates the packing later by + concatenating the per-half results. + + (b) + (c) co-occurring + MoE models where each expert is *also* stored pre-pack — e.g. + `experts.2.w1.weight` is one expert (b) AND only half of the + gate_up pack (c). Handled in order: resolve the expert axis + first (b), then run the generic interval loop over whatever + placements remain (c behavior falls out naturally). """ def __init__(self, param: DTensor): @@ -915,22 +910,41 @@ def __init__(self, param: DTensor): def shard_tensor( self, source: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None ) -> torch.Tensor | None: + # Idea + # ---- + # Mesh convention: [FSDP, TP]; FSDP shards dim 0; TP shards dim 0 + # (ColwiseParallel) or dim 1 (RowwiseParallel). + # + # Two placement layouts show up in practice: + # + # ── Case 1: FSDP and TP on different dims — [Shard(0), Shard(1)] ── + # Arises on row-parallel layers (o_proj, down_proj). No collision: + # + # 2×2 mesh, param (8, 8), rank (0, 0): + # initial → dim 0: [(0, 8)] dim 1: [(0, 8)] + # after Shard(0) → dim 0: [(0, 4)] dim 1: [(0, 8)] + # after Shard(1) → dim 0: [(0, 4)] dim 1: [(0, 4)] + # → source[0:4, 0:4] + # + # ── Case 2: FSDP and TP on the same dim — [_StridedShard(0, sf), + # Shard(0)] ──────────────────────────────────────────────────── + # Arises on column-parallel layers (q/k/v/qkv_proj, gate/up/gate_up + # _proj, lm_head). Both want dim 0 → collision; the stride resolves + # it so TP still sees contiguous chunks: + # + # 2×2 mesh, size-4 dim 0, sf=2: + # after _StridedShard(0, sf=2): split (0,4) into 2 groups + # (0,2) and (2,4); each FSDP rank keeps its half of each → + # FSDP 0 → [(0,1),(2,3)] FSDP 1 → [(1,2),(3,4)] + # after Shard(0): view the list flat, take TP rank's half → + # (0,0)→row 0 (0,1)→row 2 (1,0)→row 1 (1,1)→row 3 source_shape = list(source.shape) if isinstance(source, torch.Tensor) else source.get_shape() - # ------------------------------------------------------------------ - # What is the source? - # (a) full weight source.ndim == param.ndim - # (b) one expert source.ndim == param.ndim - 1, leading expert axis dropped - # (c) one half of a pack source is smaller on the packed axis (w1 before gate_up concat) - # (b) and (c) can co-occur (MoE expert stored as pre-pack w1/w3). - # ------------------------------------------------------------------ source_missing_leading_axis = self.param_ndim > len(source_shape) - # ------------------------------------------------------------------ # Collect placements that actually split a dim. # _StridedShard.is_shard() returns False in PyTorch, so also accept # any non-Replicate placement that exposes a `dim` attribute. - # ------------------------------------------------------------------ placements = [ (mesh_dim, p) for mesh_dim, p in enumerate(self.placements) @@ -939,83 +953,43 @@ def shard_tensor( if not placements: return source[...].to(device=device, dtype=dtype) # no sharding → full copy - # ------------------------------------------------------------------ - # Case (b): resolve the expert axis up front. - # - # The expert axis is param dim 0; it does not exist in the source - # (the source is one single expert, named by tensor_idx). Example: - # param = (E=4, H=8, I=4) with placements [Shard(0), Shard(1)] - # source = (H=8, I=4) for experts.2.w1.weight, tensor_idx=2 - # - # Shard(0) on the expert axis is really asking "which experts do I - # own?" — it splits experts across ranks, not values inside one - # expert. _owns_expert(tensor_idx) answers that question: - # - not owned → this whole file is for other ranks → return None - # - owned → Shard(0) is now fully handled; remove it from - # `placements` so we don't try to slice it again. - # - # Remaining placements (here: Shard(1)) are inner — they still - # need to shard source's dims. We fall through to the generic loop. - # ------------------------------------------------------------------ source_is_one_expert = tensor_idx is not None and self.param_ndim == len(source_shape) + 1 if source_is_one_expert and any(self._norm_dim(p.dim) == 0 for _, p in placements): - if not self._owns_expert(tensor_idx): + if not self._owns_expert(tensor_idx): # Case (b) -> Not owned return None + # Case (b) -> Owned, drop expert axis and continue with the generic interval loop placements = [(mesh_dim, p) for mesh_dim, p in placements if self._norm_dim(p.dim) != 0] if not placements: # Expert axis was the only sharding → keep the whole expert tensor. return source[...].to(device=device, dtype=dtype) - - # ------------------------------------------------------------------ - # Cases (a) and (c): generic interval loop. - # - # Build, per source dim, the list of (start, end) index ranges this - # rank owns. We start with "the whole axis" on every dim and then - # apply each placement, which narrows the list on its source dim. - # - # Shard(d) → treat the current list on dim d as one flat - # logical sequence and take this rank's - # contiguous chunk of it. If the prior result - # was a single interval, the output is still - # one interval (standard Shard). If it was - # multiple disjoint intervals (left behind by - # a _StridedShard on the same dim), the chunk - # may span several of them — handled by - # _contiguous_intervals' flat-walk logic. - # _StridedShard → cut each (start, end) into split_factor groups - # and keep this rank's contiguous sub-range of - # each group. List length grows by split_factor; - # _slice_and_cat will read the pieces and - # concatenate them. + + # How this maps to the code below + # ------------------------------- + # - intervals[d] (line: intervals = [[(0, size)] ...]) is the + # per-dim range list, initialized to the whole source axis. + # - Each loop iteration picks one helper to narrow + # intervals[source_dim]. The default rule is: + # Shard → _contiguous_intervals + # _StridedShard → _strided_intervals + # ...but there is one exception: if `source` has fewer dims + # than `param` (cases b or c), _StridedShard *also* falls back + # to _contiguous_intervals. That is what `is_interleaved` + # encodes in a single line. + + # All four combinations: # - # Example — 2×2 mesh [FSDP, TP], param (8, 8), rank coord (0, 0), - # placements [Shard(0), Shard(1)]: - # initial → dim 0: [(0, 8)] dim 1: [(0, 8)] - # after Shard(0)→ dim 0: [(0, 4)] dim 1: [(0, 8)] - # after Shard(1)→ dim 0: [(0, 4)] dim 1: [(0, 4)] - # slice result → source[0:4, 0:4] - # - # Example — [_StridedShard(0, sf=2), Shard(0)] on size-4 dim, 2×2 - # mesh. _StridedShard leaves rank 0 with [(0, 1), (2, 3)] (rows 0 - # and 2 as a flat 2-row sequence). Shard then takes rank 0's half - # of that flat sequence → [(0, 1)]; rank 1 gets [(2, 3)]. - # - # Two placements on the same tensor dim (e.g. [Shard(0), Shard(0)] - # when two mesh dims both split rows — 2-level FSDP) fold naturally: - # the second call narrows the list produced by the first. - # ------------------------------------------------------------------ + # placement type source ndim → branch case + # Shard same as param → _contiguous_intervals (a) + # _StridedShard same as param → _strided_intervals (a) + # Shard less than param → _contiguous_intervals (b)/(c) + # _StridedShard less than param → _contiguous_intervals (b)/(c) + intervals: list[list[tuple[int, int]]] = [[(0, size)] for size in source_shape] for mesh_dim, placement in placements: source_dim = self._source_dim(placement.dim, source_shape) sub_mesh = self._get_sub_mesh(mesh_dim) rank, world_size = sub_mesh.get_local_rank(), sub_mesh.size() - # Case (c): _StridedShard only makes sense when the packed axis - # actually exists in the source. If the source is a pre-pack - # half (w1 alone, before gate_up concat), the split_factor - # groups haven't been laid out yet — we just take a plain - # contiguous cut on this half, and the WeightConverter will - # rebuild the packed layout afterwards. is_interleaved = not placement.is_shard() and not source_missing_leading_axis if is_interleaved: intervals[source_dim] = self._strided_intervals( @@ -1023,9 +997,92 @@ def shard_tensor( ) else: intervals[source_dim] = self._contiguous_intervals(intervals[source_dim], rank, world_size) - + # Finally, read the source with those intervals (concatenating along a multi-interval dim if any). return self._slice_and_cat(source, intervals, device, dtype) + def _strided_intervals( + self, intervals: list[tuple[int, int]], rank: int, world_size: int, split_factor: int + ) -> list[tuple[int, int]]: + """Split each interval into `split_factor` groups, then shard contiguously within each group. + + Produces one sub-interval per group (up to `split_factor` disjoint pieces + per input interval), matching _StridedShard's packed-axis layout. + + Example: + intervals = [(0, 4)], world_size = 2, split_factor = 2 + group 0 = (0, 2), group 1 = (2, 4). + Within each group each rank owns half: + rank 0 → [(0, 1), (2, 3)] (first half of each group) + rank 1 → [(1, 2), (3, 4)] (second half of each group) + """ + 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]]: + """Shard semantics: treat `intervals` as one flat logical sequence and take this rank's contiguous chunk of it. + + For a single input interval this is a plain per-interval cut. For + multiple intervals (left by a prior _StridedShard on the same dim), + the chunk can cross range boundaries — walk the intervals and emit + the sub-ranges that fall inside [my_offset, my_offset + my_size) in + the flat view. + + Examples: + Single interval, plain split: + intervals = [(0, 8)], world_size = 2 + rank 0 → [(0, 4)] + rank 1 → [(4, 8)] + + Multiple intervals from a prior _StridedShard, chunk stays + inside one input interval: + intervals = [(0, 1), (2, 3)], world_size = 2 + flat view = [row 0, row 2], rank 0 owns first half + rank 0 → [(0, 1)] (row 0 only) + rank 1 → [(2, 3)] (row 2 only) + + Multiple intervals, chunk crosses a boundary: + intervals = [(0, 2), (4, 6)], world_size = 2, rank = 0 + flat view = [0,1,2,3] mapped to rows [0,1,4,5]; + rank 0 owns flat positions 0..1 → rows 0..1 + → [(0, 2)] + Same with world_size = 4, rank = 1: + rank 1 owns flat position 1 → row 1 + → [(1, 2)] + """ + 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): """Read `source` with per-source-dim intervals; concat along the sole multi-interval dim if any. @@ -1092,62 +1149,6 @@ def _source_dim(self, placement_dim: int, source_shape) -> int: return dim - def _contiguous_intervals( - self, intervals: list[tuple[int, int]], rank: int, world_size: int - ) -> list[tuple[int, int]]: - """Shard semantics: treat `intervals` as one flat logical sequence and take this rank's contiguous chunk of it. - - For a single input interval this is a plain per-interval cut. For - multiple intervals (left by a prior _StridedShard on the same dim), - the chunk can cross range boundaries — walk the intervals and emit - the sub-ranges that fall inside [my_offset, my_offset + my_size) in - the flat view. - """ - 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 _strided_intervals( - self, intervals: list[tuple[int, int]], rank: int, world_size: int, split_factor: int - ) -> list[tuple[int, int]]: - """Split each interval into `split_factor` groups, then shard contiguously within each group. - - Produces one sub-interval per group (up to `split_factor` disjoint pieces - per input interval), matching _StridedShard's packed-axis layout. - """ - 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 dot_natural_key(s: str): """Sort key for state-dict names: split on ``"."`` and sort digits numerically and strings alphabetically. We emit a tuple at each point to sort ints From f710f0d3f0ff656c178ecd23f02c998959fd8170 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 22 Apr 2026 12:09:58 +0000 Subject: [PATCH 038/116] refactor dtensor test to reflect real world scenario --- tests/utils/test_core_model_loading.py | 332 +++++++++++++++++-------- 1 file changed, 235 insertions(+), 97 deletions(-) diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py index 750dd04a63b4..6da2762f95bf 100644 --- a/tests/utils/test_core_model_loading.py +++ b/tests/utils/test_core_model_loading.py @@ -877,60 +877,116 @@ def test_ernie4_5_vl_moe_conversion_reversed(self): class TestDtensorShardOperation(unittest.TestCase): """Unit tests for DtensorShardOperation.shard_tensor. - The class handles three source shapes relative to the DTensor param: - (a) full weight — interval math applies directly. - (b) one expert of a packed param — skip if not owned, else drop the - expert placement and continue. - (c) one pack-half — interleaved (_StridedShard) on the - missing packed axis degrades to a - plain contiguous cut. - Shard after _StridedShard on the same dim composes via a flat-view cut - in _contiguous_intervals. - - Covered paths: - - no sharding / replicate-only [test_no_shard_returns_full_tensor] - - single-placement contiguous shard [test_1d_single_shard, - test_negative_dim_normalizes_correctly, - test_contiguous_shard_uneven_division] - - multi-placement contiguous shard [test_nd_contiguous_single_slice] - - strided shard producing disjoint intervals [test_nd_strided_shard_disjoint_ranges] - - Shard after _StridedShard on same dim [test_nd_strided_plus_shard_same_dim] - - per-expert, no expert-axis sharding [test_expert_shaped_tp_only_no_expert_sharding] - - per-expert, expert axis sharded (skip/own) [test_expert_filtering] - - per-expert + inner TP [test_expert_filtering_preserves_inner_sharding] - - per-expert + pre-pack axis (strided→contig) [test_prepacked_strided_shard_uses_contiguous_source_slice] - - internal _slice_and_cat [test_slice_and_cat_all_single_ranges, - test_slice_and_cat_raises_on_two_multi_range_dims] + Each test mirrors a real transformer-layer scenario that produces one + of the placement patterns shard_tensor must handle. The docstring of + every test walks the interval-narrowing loop step by step so the + expected output is traceable by hand. + + Scenarios Placement pattern + --------------------------------------------------------------------------------------------- + Row-parallel layer (o_proj, down_proj) [Shard(0), Shard(1)] + Column-parallel layer (q/k/v_proj, gate_up_proj, lm_head) [_StridedShard(0, sf=TP), Shard(0)] + _StridedShard alone (TP on input dim with group pattern) [Shard(0), _StridedShard(1, sf=2)] + MoE expert, not owned (mixtral experts) [Shard(0)] on expert dim + MoE expert, owned (mixtral experts) [Shard(0)] on expert dim + MoE expert + inner TP (experts + TP) [Shard(0), Shard(1)] + MoE TP without expert- (experts + TP, expert axis [Shard(1)] on inner dim + axis shard replicated) + Pre-pack half (one of gate_up halves) _StridedShard on missing packed axis + Replicate only (biases, norms) [Replicate()] + + Edge cases: + - uneven shard division (5 rows / 2 ranks) + - negative dim index normalization (Shard(-1)) + + Internal _slice_and_cat helper: + - fast path (one interval per dim) + - rejection of two multi-range dims """ - def test_no_shard_returns_full_tensor(self): - """Replicate-only → full copy.""" - mesh = FakeMesh(shape=(2,), rank=0) - op = _make_dtensor_shard_op(mesh, [Replicate()], param_shape=(4, 4), local_shape=(4, 4)) - tensor = torch.arange(16).reshape(4, 4).float() - torch.testing.assert_close(op.shard_tensor(tensor), tensor) + # -------------------------------------------------------------- + # Row-parallel: FSDP and TP shard different dims (no collision) + # -------------------------------------------------------------- + def test_row_parallel_layer_shards_different_dims(self): + """Row-parallel (o_proj / down_proj) on a 2×2 mesh [FSDP, TP]. - def test_1d_single_shard(self): - """1D mesh with a single Shard(0) → contiguous split across ranks.""" - tensor = torch.arange(16).reshape(4, 4).float() - for rank, expected in [(0, tensor[:2]), (1, tensor[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, msg=f"rank {rank}") + param = Linear.weight, shape (out=8, in=8). + placements = [Shard(0), Shard(1)] — FSDP on output rows, TP on input cols. + + Walk for rank (FSDP=0, TP=0): + init → dim 0: [(0, 8)] dim 1: [(0, 8)] + after Shard(0) → dim 0: [(0, 4)] dim 1: [(0, 8)] + after Shard(1) → dim 0: [(0, 4)] dim 1: [(0, 4)] + → source[0:4, 0:4] - def test_nd_contiguous_single_slice(self): - """nD Shard on different dims → single slice read per rank.""" + Each of the 4 ranks owns a disjoint 4×4 quadrant. + """ tensor = torch.arange(64).reshape(8, 8).float() - expected = {0: tensor[:4, :4], 1: tensor[:4, 4:], 2: tensor[4:, :4], 3: tensor[4:, 4:]} + expected = { + 0: tensor[:4, :4], # (FSDP=0, TP=0) top-left + 1: tensor[:4, 4:], # (FSDP=0, TP=1) top-right + 2: tensor[4:, :4], # (FSDP=1, TP=0) bottom-left + 3: tensor[4:, 4:], # (FSDP=1, TP=1) 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_nd_strided_shard_disjoint_ranges(self): - """_StridedShard on its own dim → multiple slice reads + cat.""" + # -------------------------------------------------------------- + # Column-parallel: FSDP + TP both shard dim 0 (stride resolves) + # -------------------------------------------------------------- + def test_column_parallel_layer_same_dim_collision(self): + """Column-parallel (q/k/v_proj, gate_up_proj, lm_head) on a 2×2 mesh [FSDP, TP]. + + param = Linear.weight, shape (out=4, in=4). + placements = [_StridedShard(0, sf=2), Shard(0)] — both on dim 0. + + Walk for rank (FSDP=0, TP=0): + init → dim 0: [(0, 4)] dim 1: [(0, 4)] + after _StridedShard(0, 2) → dim 0: [(0, 1), (2, 3)] + # groups (0,2) and (2,4); FSDP 0 keeps first half of each → rows {0, 2} + after Shard(0) → dim 0: [(0, 1)] + # view [(0,1),(2,3)] flat → {row 0, row 2}; TP 0 takes first half → row 0 + → source[0:1, :] + + Gather across FSDP: TP 0 sees rows {0,1}, TP 1 sees rows {2,3} — contiguous + chunks as column-parallel kernels require. + """ + tensor = torch.arange(16).reshape(4, 4).float() + expected = { + 0: tensor[[0]], # (FSDP=0, TP=0) + 1: tensor[[2]], # (FSDP=0, TP=1) + 2: tensor[[1]], # (FSDP=1, TP=0) + 3: tensor[[3]], # (FSDP=1, TP=1) + } + 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}") + + # -------------------------------------------------------------- + # _StridedShard alone (on its own dim): multi-interval + concat + # -------------------------------------------------------------- + def test_strided_shard_alone_produces_disjoint_intervals(self): + """_StridedShard on its own dim (different from Shard's dim) yields multi-interval reads. + + param shape (8, 8), placements = [Shard(0), _StridedShard(1, sf=2)]. + + Walk for rank (0, 0): + init → dim 0: [(0, 8)] dim 1: [(0, 8)] + after Shard(0) → dim 0: [(0, 4)] dim 1: [(0, 8)] + after _StridedShard(1, sf=2) → dim 0: [(0, 4)] dim 1: [(0, 2), (4, 6)] + # groups (0,4) and (4,8); rank 0 keeps first half of each → cols {0-1, 4-5} + → _slice_and_cat reads source[:4, 0:2] and source[:4, 4:6], + concatenates along dim 1. + """ tensor = torch.arange(64).reshape(8, 8).float() - # Shard(0) splits rows; _StridedShard(1, split_factor=2) produces disjoint col ranges expected = { 0: torch.cat([tensor[:4, :2], tensor[:4, 4:6]], dim=1), 1: torch.cat([tensor[:4, 2:4], tensor[:4, 6:8]], dim=1), @@ -947,59 +1003,60 @@ def test_nd_strided_shard_disjoint_ranges(self): ) torch.testing.assert_close(op.shard_tensor(tensor), expected[rank], msg=f"rank {rank}") - def test_nd_strided_plus_shard_same_dim(self): - """_StridedShard + Shard on the same dim: Shard takes a flat-view slice of the strided output.""" - tensor = torch.arange(16).reshape(4, 4).float() - expected = {0: tensor[[0]], 1: tensor[[2]], 2: tensor[[1]], 3: tensor[[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}") + # -------------------------------------------------------------- + # MoE experts (case b): source.ndim == param.ndim - 1 + # -------------------------------------------------------------- + def test_moe_expert_not_owned_returns_none(self): + """Expert file belongs to another rank → return None (skip the file). - def test_prepacked_strided_shard_uses_contiguous_source_slice(self): - """Pre-concat w1/w3 tensors should shard contiguously before gate/up packing.""" - tensor = torch.arange(8).reshape(4, 2).float() - for rank, expected in [(0, tensor[:2]), (1, tensor[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(tensor, tensor_idx=0), expected, msg=f"rank {rank}") + param (E=4, H=2, I=2), placements = [Shard(0)] on expert axis. + 2-rank mesh (FSDP=2). Rank 1 owns experts {2, 3} (offset=2, size=2). + Loading expert_idx=0 on rank 1 → the file is for rank 0, skip. + """ + mesh = FakeMesh(shape=(2,), rank=1) + op = _make_dtensor_shard_op(mesh, [Shard(0)], param_shape=(4, 2, 2), local_shape=(2, 2, 2)) + expert_tensor = torch.ones(2, 2) + self.assertIsNone(op.shard_tensor(expert_tensor, tensor_idx=0)) - def test_expert_shaped_tp_only_no_expert_sharding(self): - """Expert-shaped param with TP on dim 1 but no expert sharding on dim 0 → plain interval loop.""" - tensor = torch.arange(8).reshape(4, 2).float() - # Shard(1) on a 3D param maps to source dim 0 (source is missing the leading expert axis). - for rank, expected in [(0, tensor[:2]), (1, tensor[2:])]: - mesh = FakeMesh(shape=(2,), rank=rank) - op = _make_dtensor_shard_op(mesh, [Shard(1)], param_shape=(4, 4, 2), local_shape=(4, 2, 2)) - torch.testing.assert_close(op.shard_tensor(tensor, tensor_idx=0), expected, msg=f"rank {rank}") + def test_moe_expert_owned_without_inner_sharding(self): + """Expert owned and no inner sharding → keep the whole expert tensor. - def test_expert_filtering(self): - """Mixtral-style experts: skip non-owned, return owned.""" + Same param / placements / mesh as above. Loading expert_idx=2 on rank 1: + source_is_one_expert = True + Shard(0) dim normalizes to 0 → enters the expert branch. + _owns_expert(2) = True → drop Shard(0) from placements. + Remaining placements = [] → early return of the full source tensor. + """ mesh = FakeMesh(shape=(2,), rank=1) op = _make_dtensor_shard_op(mesh, [Shard(0)], param_shape=(4, 2, 2), local_shape=(2, 2, 2)) expert_tensor = torch.ones(2, 2) - # rank 1 owns experts 2,3 (offset=2) - self.assertIsNone(op.shard_tensor(expert_tensor, tensor_idx=0)) torch.testing.assert_close(op.shard_tensor(expert_tensor, tensor_idx=2), expert_tensor) - def test_expert_filtering_preserves_inner_sharding(self): - """MoE expert ownership checks should still apply TP sharding on inner dims.""" + def test_moe_expert_owned_with_inner_tp(self): + """MoE expert sharded on expert axis (FSDP) and inner dim (TP). + + param (E=4, H=4, I=2), placements = [Shard(0), Shard(1)]: + Shard(0) on 2-way FSDP → expert-axis shard + Shard(1) on 2-way TP → inner hidden-dim shard + source = (H=4, I=2) for one expert, tensor_idx=1. + + Walk per rank (FSDP, TP) with tensor_idx=1: + expert branch: rank owns expert 1 only if FSDP==0 (offset 0, size 2) + FSDP=1 ranks → return None + FSDP=0 ranks → drop Shard(0); continue with Shard(1) + remaining placements = [(1, Shard(1))] + source_dim = _source_dim(1, [4, 2]) = 1 - missing_leading(1) = 0 + init → dim 0: [(0, 4)] dim 1: [(0, 2)] + after Shard(1)→src 0: + TP=0 rank → dim 0: [(0, 2)] → source[:2] + TP=1 rank → dim 0: [(2, 4)] → source[2:] + """ tensor = torch.arange(8).reshape(4, 2).float() expected = { - 0: tensor[:2], - 1: tensor[2:], - 2: None, - 3: None, + 0: tensor[:2], # (FSDP=0, TP=0) — owns expert 1, inner rows 0-1 + 1: tensor[2:], # (FSDP=0, TP=1) — owns expert 1, inner rows 2-3 + 2: None, # (FSDP=1, TP=0) — does not own expert 1 + 3: None, # (FSDP=1, TP=1) — does not own expert 1 } for rank in range(4): mesh = FakeMesh(shape=(2, 2), rank=rank) @@ -1010,16 +1067,81 @@ def test_expert_filtering_preserves_inner_sharding(self): else: torch.testing.assert_close(shard, expected[rank], msg=f"rank {rank}") - def test_negative_dim_normalizes_correctly(self): - """Shard(-1) on a 2D tensor should shard the last dimension.""" - tensor = torch.arange(16).reshape(4, 4).float() - for rank, expected in [(0, tensor[:, :2]), (1, tensor[:, 2:])]: + def test_moe_expert_tp_only_no_expert_axis_shard(self): + """MoE param with TP on inner dim but no expert-axis sharding. + + param (E=4, H=4, I=2), placements = [Shard(1)] — TP only, 2-rank mesh. + source = (H=4, I=2), tensor_idx=0. + + source_is_one_expert = True, BUT Shard(1).dim normalizes to 1 ≠ 0, + so no placement targets the expert axis → skip the expert branch and + fall into the generic loop with missing_leading_dims=1. + + Walk: + source_dim = _source_dim(1, [4, 2]) = 1 - 1 = 0 + init → dim 0: [(0, 4)] dim 1: [(0, 2)] + after Shard(1)→src 0: + rank 0 → dim 0: [(0, 2)] → source[:2] + rank 1 → dim 0: [(2, 4)] → source[2:] + """ + tensor = torch.arange(8).reshape(4, 2).float() + for rank, expected in [(0, tensor[:2]), (1, tensor[2:])]: mesh = FakeMesh(shape=(2,), rank=rank) - op = _make_dtensor_shard_op(mesh, [Shard(-1)], param_shape=(4, 4), local_shape=(4, 2)) - torch.testing.assert_close(op.shard_tensor(tensor), expected, msg=f"rank {rank}") + op = _make_dtensor_shard_op(mesh, [Shard(1)], param_shape=(4, 4, 2), local_shape=(4, 2, 2)) + torch.testing.assert_close(op.shard_tensor(tensor, tensor_idx=0), expected, msg=f"rank {rank}") + + # -------------------------------------------------------------- + # Pre-pack half (case c): _StridedShard on missing packed axis + # -------------------------------------------------------------- + def test_prepack_half_strided_degrades_to_contiguous(self): + """Pre-concat w1 / w3 tensor: _StridedShard on the packed axis degrades. + + param = packed gate_up per-expert, shape (E=8, 2H=8, D=2). + source = (H=4, D=2) — single w1 half for one expert, tensor_idx=0. + placements = [_StridedShard(dim=1, sf=2)] — would stride the packed 2H dim. + + source_missing_leading_axis = True (source ndim 2 < param ndim 3). + is_interleaved requires same ndim → False. + → _StridedShard falls through to _contiguous_intervals. + The packing is rebuilt later by the WeightConverter's Concatenate op. + + Walk for rank 0 (2-rank mesh): + source_dim = _source_dim(1, [4, 2]) = 1 - 1 = 0 + init → dim 0: [(0, 4)] dim 1: [(0, 2)] + after _StridedShard→src 0 → dim 0: [(0, 2)] dim 1: [(0, 2)] + → source[:2] + """ + tensor = torch.arange(8).reshape(4, 2).float() + for rank, expected in [(0, tensor[:2]), (1, tensor[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(tensor, tensor_idx=0), expected, msg=f"rank {rank}") + # -------------------------------------------------------------- + # Replicate only (biases, norms): no narrowing + # -------------------------------------------------------------- + def test_replicate_only_returns_full_tensor(self): + """All placements are Replicate → placements filter drops everything + → early return with a full-tensor copy.""" + mesh = FakeMesh(shape=(2,), rank=0) + op = _make_dtensor_shard_op(mesh, [Replicate()], param_shape=(4, 4), local_shape=(4, 4)) + tensor = torch.arange(16).reshape(4, 4).float() + torch.testing.assert_close(op.shard_tensor(tensor), tensor) + + # -------------------------------------------------------------- + # Edge cases for _contiguous_intervals + # -------------------------------------------------------------- def test_contiguous_shard_uneven_division(self): - """Shard(0) on 5 rows across 2 ranks → rank 0 gets 3 rows, rank 1 gets 2.""" + """Size-5 axis sharded across 2 ranks: rank 0 gets 3 rows, rank 1 gets 2. + + _contiguous_intervals calls Shard.local_shard_size_and_offset which + rounds up the per-rank share; the last rank takes whatever remains. + """ tensor = torch.arange(20).reshape(5, 4).float() expected = {0: tensor[:3], 1: tensor[3:]} for rank in range(2): @@ -1028,8 +1150,22 @@ def test_contiguous_shard_uneven_division(self): op = _make_dtensor_shard_op(mesh, [Shard(0)], param_shape=(5, 4), local_shape=(local_rows, 4)) torch.testing.assert_close(op.shard_tensor(tensor), expected[rank], msg=f"rank {rank}") - def test_slice_and_cat_all_single_ranges(self): - """When every dim has exactly one interval, _slice_and_cat does a single slice read (no concat).""" + def test_negative_dim_normalization(self): + """Shard(-1) on a 2D tensor shards the last dim (dim 1). + + _norm_dim(-1) with param_ndim=2 → 2 + (-1) = 1. + """ + tensor = torch.arange(16).reshape(4, 4).float() + for rank, expected in [(0, tensor[:, :2]), (1, tensor[:, 2:])]: + mesh = FakeMesh(shape=(2,), rank=rank) + op = _make_dtensor_shard_op(mesh, [Shard(-1)], param_shape=(4, 4), local_shape=(4, 2)) + torch.testing.assert_close(op.shard_tensor(tensor), expected, msg=f"rank {rank}") + + # -------------------------------------------------------------- + # Internal helper: _slice_and_cat + # -------------------------------------------------------------- + def test_slice_and_cat_fast_path_single_interval_per_dim(self): + """Every dim has exactly one interval → fast path: single slice read, no concat.""" 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, 4)) @@ -1037,8 +1173,10 @@ def test_slice_and_cat_all_single_ranges(self): result = op._slice_and_cat(tensor, intervals, None, None) torch.testing.assert_close(result, tensor[0:4, 2:6]) - def test_slice_and_cat_raises_on_two_multi_range_dims(self): - """Multiple disjoint intervals on two different dims → ValueError.""" + def test_slice_and_cat_rejects_two_multi_interval_dims(self): + """Two dims with multiple disjoint ranges would require a 2D outer-product + of reads. Not supported → ValueError. + """ 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, 4)) From a35993c156e3b791aab65857b86cb8ea3ede437c Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 22 Apr 2026 13:18:06 +0000 Subject: [PATCH 039/116] more comments --- src/transformers/core_model_loading.py | 137 +++++++----------- .../integrations/tensor_parallel.py | 17 ++- 2 files changed, 65 insertions(+), 89 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 2450b574d0af..8ef35d46e7a9 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -881,22 +881,18 @@ class DtensorShardOperation: dims (e.g. Shard(1) on H). (c) One half of a pack — source smaller than param on the packed axis - Some weights are concatenated at load time (gate+up → gate_up, - Q+K+V → qkv). The checkpoint stores each half separately - (e.g. `w1` before concat with `w3`), so source is smaller than - param along the packed axis. + Some params are built by concatenating two checkpoint tensors + (gate+up → gate_up, Q+K+V → qkv). Each half is loaded on its own, + so source is smaller than param on that axis. Example: param = (2H, D), source = (H, D) for just `w1`. - → _StridedShard on the packed axis cannot stride over a pack - that does not exist yet, so it degrades to a plain contiguous - cut here. The WeightConverter re-creates the packing later by - concatenating the per-half results. + → _StridedShard can't stride a pack that doesn't exist yet, so on + the packed axis it falls back to a plain contiguous cut. The + WeightConverter concatenates the halves afterwards. (b) + (c) co-occurring - MoE models where each expert is *also* stored pre-pack — e.g. - `experts.2.w1.weight` is one expert (b) AND only half of the - gate_up pack (c). Handled in order: resolve the expert axis - first (b), then run the generic interval loop over whatever - placements remain (c behavior falls out naturally). + MoE checkpoint that is both per-expert and pre-pack (e.g. + `experts.2.w1.weight`). Resolve the expert axis first (b); the + generic loop then handles the remaining dims with (c) behavior. """ def __init__(self, param: DTensor): @@ -910,34 +906,6 @@ def __init__(self, param: DTensor): def shard_tensor( self, source: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None ) -> torch.Tensor | None: - # Idea - # ---- - # Mesh convention: [FSDP, TP]; FSDP shards dim 0; TP shards dim 0 - # (ColwiseParallel) or dim 1 (RowwiseParallel). - # - # Two placement layouts show up in practice: - # - # ── Case 1: FSDP and TP on different dims — [Shard(0), Shard(1)] ── - # Arises on row-parallel layers (o_proj, down_proj). No collision: - # - # 2×2 mesh, param (8, 8), rank (0, 0): - # initial → dim 0: [(0, 8)] dim 1: [(0, 8)] - # after Shard(0) → dim 0: [(0, 4)] dim 1: [(0, 8)] - # after Shard(1) → dim 0: [(0, 4)] dim 1: [(0, 4)] - # → source[0:4, 0:4] - # - # ── Case 2: FSDP and TP on the same dim — [_StridedShard(0, sf), - # Shard(0)] ──────────────────────────────────────────────────── - # Arises on column-parallel layers (q/k/v/qkv_proj, gate/up/gate_up - # _proj, lm_head). Both want dim 0 → collision; the stride resolves - # it so TP still sees contiguous chunks: - # - # 2×2 mesh, size-4 dim 0, sf=2: - # after _StridedShard(0, sf=2): split (0,4) into 2 groups - # (0,2) and (2,4); each FSDP rank keeps its half of each → - # FSDP 0 → [(0,1),(2,3)] FSDP 1 → [(1,2),(3,4)] - # after Shard(0): view the list flat, take TP rank's half → - # (0,0)→row 0 (0,1)→row 2 (1,0)→row 1 (1,1)→row 3 source_shape = list(source.shape) if isinstance(source, torch.Tensor) else source.get_shape() source_missing_leading_axis = self.param_ndim > len(source_shape) @@ -963,30 +931,15 @@ def shard_tensor( # Expert axis was the only sharding → keep the whole expert tensor. return source[...].to(device=device, dtype=dtype) - # How this maps to the code below - # ------------------------------- - # - intervals[d] (line: intervals = [[(0, size)] ...]) is the - # per-dim range list, initialized to the whole source axis. - # - Each loop iteration picks one helper to narrow - # intervals[source_dim]. The default rule is: - # Shard → _contiguous_intervals - # _StridedShard → _strided_intervals - # ...but there is one exception: if `source` has fewer dims - # than `param` (cases b or c), _StridedShard *also* falls back - # to _contiguous_intervals. That is what `is_interleaved` - # encodes in a single line. - - # All four combinations: - # - # placement type source ndim → branch case - # Shard same as param → _contiguous_intervals (a) - # _StridedShard same as param → _strided_intervals (a) - # Shard less than param → _contiguous_intervals (b)/(c) - # _StridedShard less than param → _contiguous_intervals (b)/(c) - + # Example - case (a) full weight + # input (source_shape=(8, 16), placements=[Shard(0)], world_size=2): + # intervals -> [[(0, 8)], [(0, 16)]] # whole range on every source dim + # output: + # rank 0: intervals -> [[(0, 4)], [(0, 16)]] # → source[0:4, 0:16] + # rank 1: intervals -> [[(4, 8)], [(0, 16)]] # → source[4:8, 0:16] intervals: list[list[tuple[int, int]]] = [[(0, size)] for size in source_shape] for mesh_dim, placement in placements: - source_dim = self._source_dim(placement.dim, source_shape) + source_dim = self._param_dim_to_source_dim(placement.dim, source_shape) sub_mesh = self._get_sub_mesh(mesh_dim) rank, world_size = sub_mesh.get_local_rank(), sub_mesh.size() @@ -997,7 +950,8 @@ def shard_tensor( ) else: intervals[source_dim] = self._contiguous_intervals(intervals[source_dim], rank, world_size) - # Finally, read the source with those intervals (concatenating along a multi-interval dim if any). + + # Read the source with those intervals (concatenating along a multi-interval dim if any). return self._slice_and_cat(source, intervals, device, dtype) def _strided_intervals( @@ -1033,12 +987,6 @@ def _contiguous_intervals( ) -> list[tuple[int, int]]: """Shard semantics: treat `intervals` as one flat logical sequence and take this rank's contiguous chunk of it. - For a single input interval this is a plain per-interval cut. For - multiple intervals (left by a prior _StridedShard on the same dim), - the chunk can cross range boundaries — walk the intervals and emit - the sub-ranges that fall inside [my_offset, my_offset + my_size) in - the flat view. - Examples: Single interval, plain split: intervals = [(0, 8)], world_size = 2 @@ -1051,15 +999,6 @@ def _contiguous_intervals( flat view = [row 0, row 2], rank 0 owns first half rank 0 → [(0, 1)] (row 0 only) rank 1 → [(2, 3)] (row 2 only) - - Multiple intervals, chunk crosses a boundary: - intervals = [(0, 2), (4, 6)], world_size = 2, rank = 0 - flat view = [0,1,2,3] mapped to rows [0,1,4,5]; - rank 0 owns flat positions 0..1 → rows 0..1 - → [(0, 2)] - Same with world_size = 4, rank = 1: - rank 1 owns flat position 1 → row 1 - → [(1, 2)] """ total = sum(end - start for start, end in intervals) my_size, my_offset = Shard.local_shard_size_and_offset(total, world_size, rank) @@ -1090,6 +1029,21 @@ def _slice_and_cat(self, source, intervals, device, dtype): owns on that source dim. At most one dim may have more than one piece (from _StridedShard); two such dims would require a 2D outer product of reads and are rejected. + + Examples: + 1) Single interval per dim — one contiguous read: + source shape = [8, 4] + intervals = [[(0, 4)], [(0, 4)]] + → source[0:4, 0:4] + + 2) Multi-interval on one dim — read each piece, concat on that dim: + source shape = [8, 4] + intervals = [[(0, 2), (4, 6)], [(0, 4)]] + → cat([source[0:2, 0:4], source[4:6, 0:4]], dim=0) + + 3) Multi-interval on two dims — rejected: + intervals = [[(0, 2), (4, 6)], [(0, 1), (2, 3)]] + → ValueError (would require a 2D outer product of reads) """ multi_interval_dim: int | None = None slices: list[slice] = [] @@ -1126,21 +1080,40 @@ def _owns_expert(self, expert_idx: int) -> bool: return first_owned_expert <= expert_idx < first_owned_expert + self.local_shape[0] def _get_sub_mesh(self, mesh_dim: int): - """1-D sub-mesh along a single mesh axis (e.g. just the TP axis).""" if self.device_mesh.ndim > 1: return self.device_mesh[self.device_mesh.mesh_dim_names[mesh_dim]] return self.device_mesh def _norm_dim(self, dim: int) -> int: - """Normalize a (possibly negative) param dim to a non-negative index.""" return dim if dim >= 0 else self.param_ndim + dim - def _source_dim(self, placement_dim: int, source_shape) -> int: + def _param_dim_to_source_dim(self, placement_dim: int, source_shape) -> int: """Map a placement's param dim onto the corresponding source-tensor dim. When the source has fewer dims than the param (expert or pre-pack source), leading axes are absent: any placement dim past the missing prefix shifts down by the difference. + + Examples: + 1) No missing dims (common case) — source and param align, no shift: + param shape = [out, in], source shape = [out, in] + placement_dim = 1 → source_dim = 1 + + 2) Stacked experts — source is one expert (missing `num_experts` axis): + param shape = [8, 4096, 2048], source shape = [4096, 2048] + missing_leading_dims = 1 + placement_dim = 1 (out) → source_dim = 0 + placement_dim = 2 (in) → source_dim = 1 + + 3) Packed QKV — source is one of Q/K/V (missing pack axis): + param shape = [3, 4096, 4096], source shape = [4096, 4096] + missing_leading_dims = 1 + placement_dim = 1 → source_dim = 0 + + 4) Two missing leading dims (stacked experts + packed QKV): + param shape = [8, 3, 4096, 4096], source shape = [4096, 4096] + missing_leading_dims = 2 + placement_dim = 3 → source_dim = 1 """ dim = self._norm_dim(placement_dim) missing_leading_dims = self.param_ndim - len(source_shape) diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index 045057975607..0c406041c037 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -282,11 +282,13 @@ def output_hook(mod, inputs, output): def _accumulate_local_param_grad(original_param: DTensor, local_grad: torch.Tensor) -> torch.Tensor: - """Copy a detached local grad into the original DTensor parameter. + """Stitch a local grad back onto the original DTensor parameter. - Packed ``_StridedShard`` parameters cannot rely on autograd through - ``DTensor.to_local()`` on older torch releases, so we materialize a local leaf - parameter for the forward and stitch its gradient back manually here. + During forward we replace the DTensor param with a detached plain-tensor + leaf (see ``_local_dtensor_params``) because ``grouped_mm`` / fused ops do + not accept DTensor inputs. That swap breaks the autograd link between the + local leaf's grad and the DTensor param's ``.grad``, so this tensor hook + runs on the leaf and copies/accumulates the grad onto the original DTensor. """ tensor_meta = original_param._spec.tensor_meta detached_grad = local_grad.detach() @@ -315,9 +317,10 @@ def _accumulate_local_param_grad(original_param: DTensor, local_grad: torch.Tens def _local_dtensor_params(module): """Temporarily swap DTensor params for local leaf params during one forward. - Needed because grouped_mm / fused ops on DTensors trigger broken autograd - paths for ``_StridedShard``. We run forward on a plain-tensor leaf param and - stitch its gradient back onto the original DTensor via a hook. Restores the + Needed because ``grouped_mm`` / fused ops do not accept DTensor inputs: we + forward through a detached plain-tensor leaf, then rely on + ``_accumulate_local_param_grad`` (registered as a tensor hook on the leaf) + to copy the backward grad onto the original DTensor param. Restores the DTensor params on exit (even on exception). """ shadows = {} From 8529d7c2e71278d4fada09f16c0281c94e806d6a Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 22 Apr 2026 15:59:34 +0000 Subject: [PATCH 040/116] fix tp olmo hybrid and exaone --- .../exaone_moe/configuration_exaone_moe.py | 9 +++++++ .../models/exaone_moe/modular_exaone_moe.py | 27 +++++++++++++++++++ .../olmo_hybrid/configuration_olmo_hybrid.py | 9 ++++--- .../models/olmo_hybrid/modular_olmo_hybrid.py | 16 +++++++++++ 4 files changed, 57 insertions(+), 4 deletions(-) diff --git a/src/transformers/models/exaone_moe/configuration_exaone_moe.py b/src/transformers/models/exaone_moe/configuration_exaone_moe.py index 874490da3422..8aa2605dbda5 100644 --- a/src/transformers/models/exaone_moe/configuration_exaone_moe.py +++ b/src/transformers/models/exaone_moe/configuration_exaone_moe.py @@ -75,6 +75,7 @@ class ExaoneMoeConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { "embed_tokens": TPStyle("vocab", "reduce_scatter"), "layers.*.input_layernorm": TPStyle("activation", "none"), @@ -146,5 +147,13 @@ def __post_init__(self, **kwargs): super().__post_init__(**kwargs) + # Dense layers can keep the Exaone4 MLP sharding, but sparse MoE blocks + # need to split their replicated output back to the sequence shard. + self.base_model_sp_plan = self.base_model_sp_plan.copy() + for layer_idx, mlp_layer_type in enumerate(self.mlp_layer_types): + self.base_model_sp_plan[f"layers.{layer_idx}.mlp"] = TPStyle( + "module", "allgather" if mlp_layer_type == "dense" else "allgather_split" + ) + __all__ = ["ExaoneMoeConfig"] diff --git a/src/transformers/models/exaone_moe/modular_exaone_moe.py b/src/transformers/models/exaone_moe/modular_exaone_moe.py index 75ec2b0bfd27..0b72cb4dc33a 100644 --- a/src/transformers/models/exaone_moe/modular_exaone_moe.py +++ b/src/transformers/models/exaone_moe/modular_exaone_moe.py @@ -20,6 +20,7 @@ from ... import initialization as init from ...cache_utils import Cache +from ...integrations.tensor_parallel import TPStyle from ...modeling_outputs import CausalLMOutputWithPast from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack @@ -80,6 +81,24 @@ class ExaoneMoeConfig(Exaone4Config): >>> configuration = model.config ```""" + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), + "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), + "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + } + vocab_size: int = 102400 hidden_size: int = 4096 intermediate_size: int = 16384 @@ -119,6 +138,14 @@ def __post_init__(self, **kwargs): super().__post_init__(**kwargs) + # Dense layers can keep the Exaone4 MLP sharding, but sparse MoE blocks + # need to split their replicated output back to the sequence shard. + self.base_model_sp_plan = self.base_model_sp_plan.copy() + for layer_idx, mlp_layer_type in enumerate(self.mlp_layer_types): + self.base_model_sp_plan[f"layers.{layer_idx}.mlp"] = TPStyle( + "module", "allgather" if mlp_layer_type == "dense" else "allgather_split" + ) + class ExaoneMoeAttention(Exaone4Attention): pass diff --git a/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py index c78c9bc548b7..2dee021756cb 100644 --- a/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py +++ b/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py @@ -94,10 +94,11 @@ class OlmoHybridConfig(PreTrainedConfig): "embed_tokens": TPStyle("vocab", "reduce_scatter"), "layers.*.input_layernorm": TPStyle("activation", "none"), "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "allgather"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "allgather"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "allgather"), + "layers.*.self_attn.o_proj": TPStyle("vocab", "reduce_scatter"), + "layers.*.linear_attn": TPStyle("module", "allgather_split", input_key="hidden_states"), "layers.*.post_attention_layernorm": TPStyle("activation", "none"), "layers.*.mlp": TPStyle("module", "allgather"), "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), diff --git a/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py index 2518a5bc247c..4b299d5d1b19 100644 --- a/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py +++ b/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py @@ -137,6 +137,22 @@ class OlmoHybridConfig(LlamaConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } + base_model_sp_plan = { + "embed_tokens": TPStyle("vocab", "reduce_scatter"), + "layers.*.input_layernorm": TPStyle("activation", "none"), + "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), + "layers.*.self_attn.q_proj": TPStyle("colwise", "allgather"), + "layers.*.self_attn.k_proj": TPStyle("colwise", "allgather"), + "layers.*.self_attn.v_proj": TPStyle("colwise", "allgather"), + "layers.*.self_attn.o_proj": TPStyle("vocab", "reduce_scatter"), + "layers.*.linear_attn": TPStyle("module", "allgather_split", input_key="hidden_states"), + "layers.*.post_attention_layernorm": TPStyle("activation", "none"), + "layers.*.mlp": TPStyle("module", "allgather"), + "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), + "layers.*.mlp.up_proj": TPStyle("colwise", "none"), + "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), + "norm": TPStyle("activation", "none"), + } vocab_size: int = 100352 hidden_size: int = 3840 From 43b792b9fb385559e12e78b60bd1827003a26afc Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 22 Apr 2026 16:28:11 +0000 Subject: [PATCH 041/116] Enhance tensor parallel weight tying logic to prevent clobbering of lm_head when embed_tokens is not in the plan. --- src/transformers/integrations/tensor_parallel.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index 0c406041c037..7ca762f2ec9f 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -648,6 +648,15 @@ def apply_tensor_parallel(model, tp_mesh, tp_plan): else: tp_plan = dict(model._tp_plan or {}) + # tie_weights() replaces lm_head.weight with embed_tokens.weight after TP is applied. + # If embed_tokens isn't in the plan, sharding lm_head as a DTensor causes tie to + # clobber it with a plain tensor (and forward then mixes DTensor/Tensor). Skip + # lm_head TP in that case so both ends stay plain and the tie is a real alias. + if getattr(model.config, "tie_word_embeddings", False): + tied_source_in_plan = any(k.endswith("embed_tokens") for k in tp_plan) + if not tied_source_in_plan: + tp_plan.pop("lm_head", None) + parallelize_plan = {} for name, _ in model.named_modules(): From 0dbef901eeaf0c6d2c379203d0864a90c23ea777 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 28 Apr 2026 05:32:14 +0000 Subject: [PATCH 042/116] fix fsdp mixin test due to missing args --- tests/test_fsdp_mixin.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py index f6f7f7e6e2ab..89ccc5f221ed 100644 --- a/tests/test_fsdp_mixin.py +++ b/tests/test_fsdp_mixin.py @@ -374,7 +374,7 @@ def train_fsdp2( ): # -- Phase 1: Pre-checkpoint run -- train only the first `checkpoint_step` steps, then save _set_determinism(SEED) - distributed_config = DistributedConfig(fsdp_plan=fsdp_plan) + distributed_config = DistributedConfig(fsdp_size=dist.get_world_size(), fsdp_plan=fsdp_plan) pre_ckpt_model = AutoModelForCausalLM.from_pretrained( init_model_dir, torch_dtype=dtype, @@ -460,7 +460,7 @@ def _test_fsdp2_save_load_impl(rank, config_class, config_dict): batches = _build_repeated_training_batches(config, device, 3) - distributed_config = DistributedConfig(fsdp_plan="auto") + distributed_config = DistributedConfig(fsdp_size=dist.get_world_size(), fsdp_plan="auto") init_tmpdir, init_tmpdir_obj = _save_init_pretrained(rank, config, torch.float32) try: From da83f324119a5ec8341a72326f1758ff067bffd2 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 28 Apr 2026 05:39:52 +0000 Subject: [PATCH 043/116] fix test non model --- src/transformers/distributed/configuration_utils.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/transformers/distributed/configuration_utils.py b/src/transformers/distributed/configuration_utils.py index 89281b0a9a39..9c0c0bd196ec 100644 --- a/src/transformers/distributed/configuration_utils.py +++ b/src/transformers/distributed/configuration_utils.py @@ -48,10 +48,11 @@ def __post_init__(self): if self.fsdp_size > 1 and self.fsdp_plan is None: self.fsdp_plan = "auto" - world_size = torch.distributed.get_world_size() - assert self.tp_size * self.fsdp_size == world_size, ( - f"tp_size ({self.tp_size}) * fsdp_size ({self.fsdp_size}) must be equal to world_size ({world_size})" - ) + if torch.distributed.is_available() and torch.distributed.is_initialized(): + world_size = torch.distributed.get_world_size() + assert self.tp_size * self.fsdp_size == world_size, ( + f"tp_size ({self.tp_size}) * fsdp_size ({self.fsdp_size}) must be equal to world_size ({world_size})" + ) @classmethod def from_dict(cls, config_dict: dict, **kwargs) -> "DistributedConfig": From 3903757c3427c51f38cb72b78fca53e61048167a Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 28 Apr 2026 05:50:01 +0000 Subject: [PATCH 044/116] skip sp plan for exaone and olmo hybrid --- .../exaone_moe/configuration_exaone_moe.py | 26 +----------------- .../models/exaone_moe/modular_exaone_moe.py | 27 +------------------ .../olmo_hybrid/configuration_olmo_hybrid.py | 17 +----------- .../models/olmo_hybrid/modular_olmo_hybrid.py | 17 +----------- 4 files changed, 4 insertions(+), 83 deletions(-) diff --git a/src/transformers/models/exaone_moe/configuration_exaone_moe.py b/src/transformers/models/exaone_moe/configuration_exaone_moe.py index 8aa2605dbda5..1dcb5399a49a 100644 --- a/src/transformers/models/exaone_moe/configuration_exaone_moe.py +++ b/src/transformers/models/exaone_moe/configuration_exaone_moe.py @@ -76,23 +76,7 @@ class ExaoneMoeConfig(PreTrainedConfig): "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } - base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), - } + base_model_sp_plan = None base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), @@ -147,13 +131,5 @@ def __post_init__(self, **kwargs): super().__post_init__(**kwargs) - # Dense layers can keep the Exaone4 MLP sharding, but sparse MoE blocks - # need to split their replicated output back to the sequence shard. - self.base_model_sp_plan = self.base_model_sp_plan.copy() - for layer_idx, mlp_layer_type in enumerate(self.mlp_layer_types): - self.base_model_sp_plan[f"layers.{layer_idx}.mlp"] = TPStyle( - "module", "allgather" if mlp_layer_type == "dense" else "allgather_split" - ) - __all__ = ["ExaoneMoeConfig"] diff --git a/src/transformers/models/exaone_moe/modular_exaone_moe.py b/src/transformers/models/exaone_moe/modular_exaone_moe.py index 0b72cb4dc33a..88b6a60e0e16 100644 --- a/src/transformers/models/exaone_moe/modular_exaone_moe.py +++ b/src/transformers/models/exaone_moe/modular_exaone_moe.py @@ -20,7 +20,6 @@ from ... import initialization as init from ...cache_utils import Cache -from ...integrations.tensor_parallel import TPStyle from ...modeling_outputs import CausalLMOutputWithPast from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack @@ -81,23 +80,7 @@ class ExaoneMoeConfig(Exaone4Config): >>> configuration = model.config ```""" - base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), - } + base_model_sp_plan = None vocab_size: int = 102400 hidden_size: int = 4096 @@ -138,14 +121,6 @@ def __post_init__(self, **kwargs): super().__post_init__(**kwargs) - # Dense layers can keep the Exaone4 MLP sharding, but sparse MoE blocks - # need to split their replicated output back to the sequence shard. - self.base_model_sp_plan = self.base_model_sp_plan.copy() - for layer_idx, mlp_layer_type in enumerate(self.mlp_layer_types): - self.base_model_sp_plan[f"layers.{layer_idx}.mlp"] = TPStyle( - "module", "allgather" if mlp_layer_type == "dense" else "allgather_split" - ) - class ExaoneMoeAttention(Exaone4Attention): pass diff --git a/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py index 2dee021756cb..5a5f2b1d8a33 100644 --- a/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py +++ b/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py @@ -90,22 +90,7 @@ class OlmoHybridConfig(PreTrainedConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } - base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "allgather"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "allgather"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "allgather"), - "layers.*.self_attn.o_proj": TPStyle("vocab", "reduce_scatter"), - "layers.*.linear_attn": TPStyle("module", "allgather_split", input_key="hidden_states"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), - } + base_model_sp_plan = None base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py index 4b299d5d1b19..c8b05d83532d 100644 --- a/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py +++ b/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py @@ -137,22 +137,7 @@ class OlmoHybridConfig(LlamaConfig): "layers.*.mlp.up_proj": TPStyle("colwise", "none"), "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), } - base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "allgather"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "allgather"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "allgather"), - "layers.*.self_attn.o_proj": TPStyle("vocab", "reduce_scatter"), - "layers.*.linear_attn": TPStyle("module", "allgather_split", input_key="hidden_states"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), - } + base_model_sp_plan = None vocab_size: int = 100352 hidden_size: int = 3840 From e51f663333f3d026c33a9eb5e60d6a187f452715 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 28 Apr 2026 05:56:13 +0000 Subject: [PATCH 045/116] linting --- src/transformers/core_model_loading.py | 8 +++----- src/transformers/distributed/configuration_utils.py | 6 +++--- tests/utils/test_core_model_loading.py | 4 ++-- 3 files changed, 8 insertions(+), 10 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 8ef35d46e7a9..7e0d14a21237 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -923,14 +923,14 @@ def shard_tensor( source_is_one_expert = tensor_idx is not None and self.param_ndim == len(source_shape) + 1 if source_is_one_expert and any(self._norm_dim(p.dim) == 0 for _, p in placements): - if not self._owns_expert(tensor_idx): # Case (b) -> Not owned + if not self._owns_expert(tensor_idx): # Case (b) -> Not owned return None # Case (b) -> Owned, drop expert axis and continue with the generic interval loop placements = [(mesh_dim, p) for mesh_dim, p in placements if self._norm_dim(p.dim) != 0] if not placements: # Expert axis was the only sharding → keep the whole expert tensor. return source[...].to(device=device, dtype=dtype) - + # Example - case (a) full weight # input (source_shape=(8, 16), placements=[Shard(0)], world_size=2): # intervals -> [[(0, 8)], [(0, 16)]] # whole range on every source dim @@ -1053,9 +1053,7 @@ def _slice_and_cat(self, source, intervals, device, dtype): 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." - ) + 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 diff --git a/src/transformers/distributed/configuration_utils.py b/src/transformers/distributed/configuration_utils.py index 9c0c0bd196ec..160283b5311b 100644 --- a/src/transformers/distributed/configuration_utils.py +++ b/src/transformers/distributed/configuration_utils.py @@ -48,9 +48,9 @@ def __post_init__(self): if self.fsdp_size > 1 and self.fsdp_plan is None: self.fsdp_plan = "auto" - if torch.distributed.is_available() and torch.distributed.is_initialized(): - world_size = torch.distributed.get_world_size() - assert self.tp_size * self.fsdp_size == world_size, ( + if torch.distributed.is_available() and torch.distributed.is_initialized(): + world_size = torch.distributed.get_world_size() + assert self.tp_size * self.fsdp_size == world_size, ( f"tp_size ({self.tp_size}) * fsdp_size ({self.fsdp_size}) must be equal to world_size ({world_size})" ) diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py index 6da2762f95bf..4f0e26b40b5e 100644 --- a/tests/utils/test_core_model_loading.py +++ b/tests/utils/test_core_model_loading.py @@ -1055,8 +1055,8 @@ def test_moe_expert_owned_with_inner_tp(self): expected = { 0: tensor[:2], # (FSDP=0, TP=0) — owns expert 1, inner rows 0-1 1: tensor[2:], # (FSDP=0, TP=1) — owns expert 1, inner rows 2-3 - 2: None, # (FSDP=1, TP=0) — does not own expert 1 - 3: None, # (FSDP=1, TP=1) — does not own expert 1 + 2: None, # (FSDP=1, TP=0) — does not own expert 1 + 3: None, # (FSDP=1, TP=1) — does not own expert 1 } for rank in range(4): mesh = FakeMesh(shape=(2, 2), rank=rank) From 96f3f29666611c7141b71d7135c07bbefc5fec7b Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 28 Apr 2026 07:44:50 +0000 Subject: [PATCH 046/116] fix import for ci --- .../integrations/tensor_parallel.py | 24 +++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index 7ca762f2ec9f..4b692938d6bf 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -18,18 +18,6 @@ from dataclasses import dataclass from typing import Literal -from torch.distributed.tensor import DTensor, Partial, Replicate, Shard, distribute_tensor -from torch.distributed.tensor.parallel import ( - ColwiseParallel, - PrepareModuleInput, - RowwiseParallel, - SequenceParallel, - parallelize_module, -) -from torch.distributed.tensor.parallel.style import ParallelStyle -from torch.distributed.tensor.placement_types import _StridedShard - -from ..distributed.patches import patch_dtensor_ops from ..utils import logging from ..utils.import_utils import is_torch_available @@ -37,6 +25,18 @@ if is_torch_available(): import torch import torch.distributed as dist + from torch.distributed.tensor import DTensor, Partial, Replicate, Shard, distribute_tensor + from torch.distributed.tensor.parallel import ( + ColwiseParallel, + PrepareModuleInput, + RowwiseParallel, + SequenceParallel, + parallelize_module, + ) + from torch.distributed.tensor.parallel.style import ParallelStyle + from torch.distributed.tensor.placement_types import _StridedShard + + from ..distributed.patches import patch_dtensor_ops # Cache this result has it's a C FFI call which can be pretty time-consuming _torch_distributed_available = torch.distributed.is_available() From dfb448e75792142f09c792e0f02fbf655beae6ef Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 28 Apr 2026 07:51:08 +0000 Subject: [PATCH 047/116] test distributed config --- .../distributed/configuration_utils.py | 16 +++++++++++----- tests/test_distributed_config.py | 8 +++++++- 2 files changed, 18 insertions(+), 6 deletions(-) diff --git a/src/transformers/distributed/configuration_utils.py b/src/transformers/distributed/configuration_utils.py index 160283b5311b..0f3abaccc352 100644 --- a/src/transformers/distributed/configuration_utils.py +++ b/src/transformers/distributed/configuration_utils.py @@ -35,17 +35,23 @@ class DistributedConfig: FSDP wrapping plan. Use `"auto"` to wrap each transformer layer + root. """ - tp_size: int = 1 + tp_size: int | None = None tp_plan: str | dict[str, str] | None = None enable_sequence_parallel: bool = False - fsdp_size: int = 1 + fsdp_size: int | None = None fsdp_plan: str | dict | None = None def __post_init__(self): - # If a size is set without a plan, default the plan to "auto" - if self.tp_size > 1 and self.tp_plan is None: + 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 self.tp_plan is None: self.tp_plan = "auto" - if self.fsdp_size > 1 and self.fsdp_plan is None: + if self.fsdp_plan is None: self.fsdp_plan = "auto" if torch.distributed.is_available() and torch.distributed.is_initialized(): diff --git a/tests/test_distributed_config.py b/tests/test_distributed_config.py index 53e9a063c123..6057d79d6dcc 100644 --- a/tests/test_distributed_config.py +++ b/tests/test_distributed_config.py @@ -50,7 +50,13 @@ def test_from_dict_kwargs_override(self): def test_to_dict(self): dc = DistributedConfig(tp_size=2, fsdp_size=4) d = dc.to_dict() - assert d == {"tp_size": 2, "tp_plan": "auto", "fsdp_size": 4, "fsdp_plan": "auto"} + assert d == { + "tp_size": 2, + "tp_plan": "auto", + "enable_sequence_parallel": False, + "fsdp_size": 4, + "fsdp_plan": "auto", + } def test_to_dict_is_a_copy(self): dc = DistributedConfig(tp_plan={"layer": "colwise"}) From 0a74b7d702627bb7f381816208f57ea89eca18c9 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 28 Apr 2026 08:26:34 +0000 Subject: [PATCH 048/116] attempt to fix guarding import ci --- .../integrations/tensor_parallel.py | 28 ++++++++++++++----- tests/utils/test_modeling_utils.py | 24 ++++++++-------- 2 files changed, 33 insertions(+), 19 deletions(-) diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index 4b692938d6bf..99fdc72bb18b 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -16,15 +16,19 @@ import contextlib import re from dataclasses import dataclass -from typing import Literal +from typing import ABC, Literal, abstractmethod from ..utils import logging -from ..utils.import_utils import is_torch_available +from ..utils.import_utils import is_torch_available, is_torch_greater_or_equal if is_torch_available(): import torch + import torch.nn as nn + +if is_torch_available() and is_torch_greater_or_equal("2.5"): import torch.distributed as dist + from torch.distributed.device_mesh import DeviceMesh from torch.distributed.tensor import DTensor, Partial, Replicate, Shard, distribute_tensor from torch.distributed.tensor.parallel import ( ColwiseParallel, @@ -73,11 +77,6 @@ def _get_parameter_tp_plan(parameter_name: str, tp_plan: dict[str, str], is_weig return None -# ============================================================================= -# Tensor Sharding Utilities -# ============================================================================= - - # ============================================================================= # High-Level API Functions # ============================================================================= @@ -249,6 +248,21 @@ def verify_tp_plan(expected_keys: list[str], tp_plan: dict[str, str | TPStyle] | logger.warning(f"The following layers were not sharded: {', '.join(unsharded_layers)}") +class ParallelStyle(ABC): + """ + Import from torch.distributed.tensor.parallel.style.ParallelStyle to avoid import guarding every class that inherits from it. + The parallel style contract defines how the module or submodule should be parallelized. + + It only defines the ``apply`` method for ``parallelize_module`` to use, this allows maximum + flexibility for different kind of style implementations. + """ + + src_data_rank: int | None = 0 + + @abstractmethod + def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module: ... + + class PrepareModuleInputOutput(ParallelStyle): """Allgather input (Shard(1) → Replicate) + local split output (Replicate → Shard(1)). diff --git a/tests/utils/test_modeling_utils.py b/tests/utils/test_modeling_utils.py index 59177fec5061..defc2add35bb 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 ( @@ -438,25 +439,24 @@ def test_model_from_pretrained_fsdp_distributes_before_loading(self): model.save_pretrained(tmp_dir) call_order = [] - def fake_distribute_model(model, tp_plan, distributed_config, device_mesh, tp_size, fsdp_plan=None): + fake_mesh = mock.Mock() + fake_mesh.mesh_dim_names = ("fsdp",) + fake_mesh.ndim = 1 + + def fake_apply_fsdp(model, fsdp_mesh, fsdp_plan): call_order.append("distribute") - self.assertEqual(fsdp_plan, {"mode": "auto"}) - model._tp_plan = {"model.layers.*.mlp.experts.gate_up_proj": "packed_colwise"} - model._is_fsdp_managed_module = True + self.assertIs(fsdp_mesh, fake_mesh) + self.assertEqual(fsdp_plan, "auto") return model def fake_load_pretrained_model(model, state_dict, checkpoint_files, load_config, expected_keys=None): call_order.append("load") - self.assertEqual(load_config.device_mesh, "fake-mesh") - self.assertEqual(load_config.device_map, {"": torch.device("cpu")}) - self.assertIsNone(load_config.tp_plan) + self.assertIs(load_config.device_mesh, fake_mesh) return mock.Mock(), None with ( - patch( - "transformers.modeling_utils.initialize_fsdp", return_value=(torch.device("cpu"), "fake-mesh", 2) - ), - patch("transformers.modeling_utils.distribute_model", side_effect=fake_distribute_model), + patch("transformers.modeling_utils.init_device_mesh", return_value=fake_mesh), + patch("transformers.modeling_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, @@ -464,7 +464,7 @@ def fake_load_pretrained_model(model, state_dict, checkpoint_files, load_config, side_effect=lambda model, load_config, loading_info: loading_info, ), ): - GPT2LMHeadModel.from_pretrained(tmp_dir, fsdp_plan={"mode": "auto"}) + GPT2LMHeadModel.from_pretrained(tmp_dir, distributed_config=DistributedConfig(fsdp_size=2)) self.assertEqual(call_order, ["distribute", "load"]) From c50e49cc20903b77ed0c3034b205f4ea9e0ec9a9 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 28 Apr 2026 08:43:36 +0000 Subject: [PATCH 049/116] fix ci check repro --- docs/source/en/model_doc/nomic_bert.md | 2 +- src/transformers/integrations/tensor_parallel.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/docs/source/en/model_doc/nomic_bert.md b/docs/source/en/model_doc/nomic_bert.md index 73b3adc8a35f..2017805fe42a 100644 --- a/docs/source/en/model_doc/nomic_bert.md +++ b/docs/source/en/model_doc/nomic_bert.md @@ -23,7 +23,7 @@ limitations under the License. ## Overview -NomicBERT was proposed in [Nomic Embed: Training a Reproducible Long Context Text Embedder](https://arxiv.org/abs/2402.01613) by +NomicBERT was proposed in [Nomic Embed: Training a Reproducible Long Context Text Embedder](https://huggingface.co/papers/2402.01613) by Zach Nussbaum, John X. Morris, Brandon Duderstadt, and Andriy Mulyar. It is BERT-inspired with the most notable extension applying [Rotary Position Embeddings](https://huggingface.co/papers/2104.09864.pdf) to an encoder model. diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index 99fdc72bb18b..ce8992b22085 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -15,8 +15,9 @@ import contextlib import re +from abc import ABC from dataclasses import dataclass -from typing import ABC, Literal, abstractmethod +from typing import Literal, abstractmethod from ..utils import logging from ..utils.import_utils import is_torch_available, is_torch_greater_or_equal From f9daf7b1f7433af73fd9c8d02fb68fc2cea7413f Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 28 Apr 2026 10:07:31 +0000 Subject: [PATCH 050/116] add ALL_PARALLEL_STYLES registry alongside TPStyle --- src/transformers/integrations/__init__.py | 2 + .../integrations/tensor_parallel.py | 89 +++++++++++++++---- 2 files changed, 74 insertions(+), 17 deletions(-) diff --git a/src/transformers/integrations/__init__.py b/src/transformers/integrations/__init__.py index d274b31837e2..bbe3461402c1 100755 --- a/src/transformers/integrations/__init__.py +++ b/src/transformers/integrations/__init__.py @@ -161,6 +161,7 @@ ] _import_structure["tensor_parallel"] = [ + "ALL_PARALLEL_STYLES", "TPStyle", "apply_tensor_parallel", "convert_strided_to_shard", @@ -299,6 +300,7 @@ from .sinq import SinqDeserialize, SinqQuantize from .spqr import replace_with_spqr_linear from .tensor_parallel import ( + ALL_PARALLEL_STYLES, TPStyle, apply_tensor_parallel, convert_strided_to_shard, diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index ce8992b22085..c7cbca1e4da1 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -20,6 +20,7 @@ from typing import Literal, abstractmethod from ..utils import logging +from ..utils.generic import GeneralInterface from ..utils.import_utils import is_torch_available, is_torch_greater_or_equal @@ -443,24 +444,26 @@ def __repr__(self) -> str: } -class _AllReduceBackward(torch.autograd.Function): - """Identity forward, allreduce-sum backward. +if is_torch_available() and is_torch_greater_or_equal("2.5"): - 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. - """ + class _AllReduceBackward(torch.autograd.Function): + """Identity forward, allreduce-sum backward. - @staticmethod - def forward(ctx, x, process_group): - ctx.process_group = process_group - return x + 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 backward(ctx, grad): - dist.all_reduce(grad, group=ctx.process_group) - return grad, None + @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(ParallelStyle): @@ -480,10 +483,10 @@ class MoEExpertsParallel(ParallelStyle): contributed); all-reduce to get the complete hidden state. """ - def __init__(self, output_layouts=None): + def __init__(self, output_layouts=None, shard_plan: dict[str, str] | None = None): super().__init__() self.output_layouts = output_layouts or Replicate() - self._moe_shard_plan: dict[str, str] = {} + self._moe_shard_plan: dict[str, str] = shard_plan or {} @staticmethod def _partition_fn(name, module, device_mesh, shard_plan): @@ -643,6 +646,58 @@ def __str__(self): return f"{self.kind}_{self.comm}" +class ParallelInterface(GeneralInterface): + """Registry of named TP styles. Configs and modeling files reference these by string name. + + Adding a new entry here is the supported way to introduce a new TP style. + Users can also override or extend at runtime via ``ALL_PARALLEL_STYLES["my_style"] = ...``. + + Naming convention: ``{kind}[_{comm}][_{extra}]``. The ``_{comm}`` suffix is dropped only when + comm is ``"none"`` (no collective). All entries are eager instances; the dict literal lives + behind a torch-availability guard so this module remains importable without torch. + """ + + _global_mapping = ( + { + # Column-parallel + "colwise": ColwiseParallel(input_layouts=Replicate(), output_layouts=Shard(-1)), + "colwise_allgather": ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()), + "colwise_loss_parallel": ColwiseParallel( + input_layouts=Shard(1), output_layouts=Shard(-1), use_local_output=False + ), + "packed_colwise": PackedColwiseParallel(input_layouts=Replicate()), + # Row-parallel + "rowwise_allreduce": RowwiseParallel(input_layouts=Shard(-1), output_layouts=Replicate()), + "rowwise_reduce_scatter": RowwiseParallel(input_layouts=Shard(-1), output_layouts=Shard(1)), + # Vocab / embedding (rowwise sharding on vocab dim) + "vocab_allreduce": RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()), + "vocab_reduce_scatter": RowwiseParallel(input_layouts=Replicate(), output_layouts=Shard(1)), + # Activation / norm (sequence-parallel passthrough) + "activation": SequenceParallel(), + "activation_seq_dim_2": SequenceParallel(sequence_dim=2), + # Module-level prepare-input + "module_allgather": PrepareModuleInput( + input_layouts=(Shard(1),), desired_input_layouts=(Replicate(),) + ), + "module_allgather_hidden_states": PrepareModuleInput( + input_kwarg_layouts={"hidden_states": Shard(1)}, + desired_input_kwarg_layouts={"hidden_states": Replicate()}, + ), + "module_allgather_split": PrepareModuleInputOutput(), + # MoE — canonical shard_plan baked in (only variant in use across configs) + "moe_experts_allreduce": MoEExpertsParallel( + output_layouts=Replicate(), + shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + ), + } + if is_torch_available() and is_torch_greater_or_equal("2.5") and _torch_distributed_available + else {} + ) + + +ALL_PARALLEL_STYLES: ParallelInterface = ParallelInterface() + + def apply_tensor_parallel(model, tp_mesh, tp_plan): """Apply tensor parallelism using PyTorch's parallelize_module. From 8a1a9e53407642d0776069f413de18c95829de58 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 28 Apr 2026 10:09:05 +0000 Subject: [PATCH 051/116] route apply_tensor_parallel through ALL_PARALLEL_STYLES --- .../integrations/tensor_parallel.py | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index c7cbca1e4da1..50d604e4a71a 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -734,7 +734,14 @@ def apply_tensor_parallel(model, tp_mesh, tp_plan): if style_value is None: continue - if isinstance(style_value, TPStyle): + if isinstance(style_value, str): + if style_value not in ALL_PARALLEL_STYLES: + raise ValueError( + f"Unknown TP style {style_value!r} for module {name!r}. " + f"Valid styles: {sorted(ALL_PARALLEL_STYLES)}" + ) + parallelize_plan[name] = ALL_PARALLEL_STYLES[style_value] + elif isinstance(style_value, TPStyle): dtensor_style = style_value.to_dtensor_style() parallelize_plan[name] = dtensor_style # For MoE modules, attach the per-parameter shard plan from TPStyle @@ -746,8 +753,8 @@ def apply_tensor_parallel(model, tp_mesh, tp_plan): else: raise TypeError( f"Unsupported plan value for '{name}': {style_value!r} (type {type(style_value).__name__}). " - f"TP plan values must be TPStyle instances or ParallelStyle instances, not strings. " - f"Migrate string plan values to TPStyle (e.g., 'colwise' -> TPStyle('colwise', 'none'))." + f"TP plan values must be strings (looked up in ALL_PARALLEL_STYLES), TPStyle, " + f"or ParallelStyle instances." ) parallelize_module(model, tp_mesh, parallelize_plan) @@ -778,7 +785,10 @@ def _inject_sp_metadata(mod, args, kwargs): # 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(isinstance(v, TPStyle) and v.comm == "loss_parallel" for v in tp_plan.values()) + has_loss_parallel = any( + v == "colwise_loss_parallel" or (isinstance(v, TPStyle) and v.comm == "loss_parallel") + for v in tp_plan.values() + ) if has_loss_parallel: from torch.distributed.tensor.parallel import loss_parallel From 7819783f4df51aa2ffa75e031c314b81e62ee5e0 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 28 Apr 2026 10:19:40 +0000 Subject: [PATCH 052/116] migrate modular files to string-based TP plans --- .../models/afmoe/modeling_afmoe.py | 5 +- .../models/afmoe/modular_afmoe.py | 5 +- .../models/apertus/configuration_apertus.py | 43 +++++---- .../models/apertus/modular_apertus.py | 43 +++++---- .../models/arcee/configuration_arcee.py | 37 ++++---- .../models/arcee/modular_arcee.py | 37 ++++---- .../models/aria/configuration_aria.py | 14 +-- src/transformers/models/aria/modular_aria.py | 15 ++-- .../models/cohere2/configuration_cohere2.py | 41 +++++---- .../models/cohere2/modular_cohere2.py | 41 +++++---- src/transformers/models/dbrx/modeling_dbrx.py | 5 +- src/transformers/models/dbrx/modular_dbrx.py | 5 +- .../deepseek_v2/configuration_deepseek_v2.py | 63 ++++++------- .../models/deepseek_v2/modular_deepseek_v2.py | 63 ++++++------- .../models/doge/configuration_doge.py | 23 +++-- src/transformers/models/doge/modular_doge.py | 23 +++-- .../models/dots1/configuration_dots1.py | 27 +++--- .../models/dots1/modular_dots1.py | 27 +++--- .../configuration_ernie4_5_vl_moe.py | 29 +++--- .../modular_ernie4_5_vl_moe.py | 29 +++--- .../models/esm/configuration_esm.py | 4 +- .../models/eurobert/modeling_eurobert.py | 5 +- .../models/eurobert/modular_eurobert.py | 5 +- .../models/exaone4/configuration_exaone4.py | 45 +++++----- .../models/exaone4/modular_exaone4.py | 45 +++++----- .../exaone_moe/configuration_exaone_moe.py | 15 ++-- .../flex_olmo/configuration_flex_olmo.py | 17 +--- .../models/flex_olmo/modular_flex_olmo.py | 17 +--- .../models/gemma/configuration_gemma.py | 41 +++++---- .../models/gemma/modular_gemma.py | 41 +++++---- .../models/gemma2/configuration_gemma2.py | 41 +++++---- .../models/gemma2/modular_gemma2.py | 41 +++++---- .../models/gemma3/configuration_gemma3.py | 45 +++++----- .../models/gemma3/modular_gemma3.py | 45 +++++----- .../models/gemma3n/configuration_gemma3n.py | 15 ++-- .../models/gemma3n/modular_gemma3n.py | 15 ++-- .../models/glm4_moe/configuration_glm4_moe.py | 27 +++--- .../models/glm4_moe/modular_glm4_moe.py | 27 +++--- .../configuration_glm4_moe_lite.py | 19 ++-- .../glm4_moe_lite/modular_glm4_moe_lite.py | 19 ++-- .../models/glm4v/configuration_glm4v.py | 15 ++-- .../models/glm4v/modular_glm4v.py | 15 ++-- .../glm4v_moe/configuration_glm4v_moe.py | 15 ++-- .../models/glm4v_moe/modular_glm4v_moe.py | 15 ++-- .../glm_image/configuration_glm_image.py | 15 ++-- .../glm_moe_dsa/configuration_glm_moe_dsa.py | 25 +++--- .../models/glm_moe_dsa/modular_glm_moe_dsa.py | 25 +++--- .../models/glm_ocr/configuration_glm_ocr.py | 15 ++-- .../models/gpt_neox/modeling_gpt_neox.py | 3 +- .../models/gpt_neox/modular_gpt_neox.py | 3 +- .../models/jais2/configuration_jais2.py | 37 ++++---- .../models/jais2/modular_jais2.py | 37 ++++---- .../models/minimax/configuration_minimax.py | 15 ++-- .../models/minimax/modular_minimax.py | 15 ++-- .../minimax_m2/configuration_minimax_m2.py | 41 ++++----- .../models/minimax_m2/modular_minimax_m2.py | 41 ++++----- .../models/nanochat/modeling_nanochat.py | 5 +- .../models/nanochat/modular_nanochat.py | 5 +- .../models/olmo2/configuration_olmo2.py | 53 +++++------ .../models/olmo2/modular_olmo2.py | 53 +++++------ .../models/olmo3/configuration_olmo3.py | 53 +++++------ .../models/olmo3/modular_olmo3.py | 53 +++++------ .../olmo_hybrid/configuration_olmo_hybrid.py | 23 ++--- .../models/olmo_hybrid/modular_olmo_hybrid.py | 23 ++--- src/transformers/models/pi0/modeling_pi0.py | 3 +- src/transformers/models/pi0/modular_pi0.py | 3 +- .../configuration_qwen2_5_omni.py | 15 ++-- .../qwen2_5_omni/modular_qwen2_5_omni.py | 15 ++-- .../models/qwen2_moe/modeling_qwen2_moe.py | 5 +- .../models/qwen2_moe/modular_qwen2_moe.py | 5 +- .../models/qwen3_5/configuration_qwen3_5.py | 15 ++-- .../models/qwen3_5/modular_qwen3_5.py | 15 ++-- .../qwen3_5_moe/configuration_qwen3_5_moe.py | 21 ++--- .../qwen3_5_moe/modeling_qwen3_5_moe.py | 4 +- .../models/qwen3_5_moe/modular_qwen3_5_moe.py | 25 +++--- .../configuration_qwen3_omni_moe.py | 14 +-- .../qwen3_omni_moe/modeling_qwen3_omni_moe.py | 2 +- .../qwen3_omni_moe/modular_qwen3_omni_moe.py | 17 ++-- .../configuration_qwen3_vl_moe.py | 14 +-- .../qwen3_vl_moe/modular_qwen3_vl_moe.py | 15 ++-- .../models/smollm3/configuration_smollm3.py | 41 +++++---- .../models/smollm3/modular_smollm3.py | 41 +++++---- .../solar_open/configuration_solar_open.py | 41 ++++----- .../models/solar_open/modular_solar_open.py | 41 ++++----- .../models/t5gemma/configuration_t5gemma.py | 41 +++++---- .../models/t5gemma/modeling_t5gemma.py | 3 +- .../models/t5gemma/modular_t5gemma.py | 3 +- .../models/t5gemma2/configuration_t5gemma2.py | 89 +++++++++---------- .../models/t5gemma2/modeling_t5gemma2.py | 3 +- .../models/t5gemma2/modular_t5gemma2.py | 3 +- .../vaultgemma/configuration_vaultgemma.py | 41 +++++---- .../models/youtu/configuration_youtu.py | 7 +- .../models/youtu/modular_youtu.py | 7 +- 93 files changed, 999 insertions(+), 1254 deletions(-) diff --git a/src/transformers/models/afmoe/modeling_afmoe.py b/src/transformers/models/afmoe/modeling_afmoe.py index bf4e39eb5327..a60ce02c5f19 100644 --- a/src/transformers/models/afmoe/modeling_afmoe.py +++ b/src/transformers/models/afmoe/modeling_afmoe.py @@ -34,7 +34,6 @@ use_kernel_func_from_hub, use_kernelized_func, ) -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast @@ -616,8 +615,8 @@ def forward( @auto_docstring class AfmoeForCausalLM(AfmoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/afmoe/modular_afmoe.py b/src/transformers/models/afmoe/modular_afmoe.py index 30cc67787772..2200f13ce4ee 100644 --- a/src/transformers/models/afmoe/modular_afmoe.py +++ b/src/transformers/models/afmoe/modular_afmoe.py @@ -21,7 +21,6 @@ from ... import initialization as init from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast @@ -396,8 +395,8 @@ def forward( class AfmoeForCausalLM(LlamaForCausalLM, AfmoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/apertus/configuration_apertus.py b/src/transformers/models/apertus/configuration_apertus.py index f116bf8324a7..33251e5f4b75 100644 --- a/src/transformers/models/apertus/configuration_apertus.py +++ b/src/transformers/models/apertus/configuration_apertus.py @@ -21,7 +21,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -47,29 +46,29 @@ class ApertusConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] default_theta = 12000000.0 base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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/apertus/modular_apertus.py b/src/transformers/models/apertus/modular_apertus.py index 2da727db820f..1850fc0cb43e 100644 --- a/src/transformers/models/apertus/modular_apertus.py +++ b/src/transformers/models/apertus/modular_apertus.py @@ -21,7 +21,6 @@ from ...activations import ACT2CLS from ...cache_utils import Cache from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...modeling_utils import ALL_ATTENTION_FUNCTIONS from ...processing_utils import Unpack @@ -65,29 +64,29 @@ class ApertusConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] default_theta = 12000000.0 base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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/arcee/configuration_arcee.py b/src/transformers/models/arcee/configuration_arcee.py index 6c2a75ee2da2..b23d249f435d 100644 --- a/src/transformers/models/arcee/configuration_arcee.py +++ b/src/transformers/models/arcee/configuration_arcee.py @@ -23,7 +23,6 @@ from transformers.utils import auto_docstring from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters @@ -47,26 +46,26 @@ class ArceeConfig(PreTrainedConfig): model_type = "arcee" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), diff --git a/src/transformers/models/arcee/modular_arcee.py b/src/transformers/models/arcee/modular_arcee.py index 5703ad4e29dd..a382b2e7191b 100644 --- a/src/transformers/models/arcee/modular_arcee.py +++ b/src/transformers/models/arcee/modular_arcee.py @@ -17,7 +17,6 @@ from transformers.utils import auto_docstring, logging -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ..llama.configuration_llama import LlamaConfig from ..llama.modeling_llama import ( @@ -51,26 +50,26 @@ class ArceeConfig(LlamaConfig): model_type = "arcee" base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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 bb5a1abb3c62..b8d9e834e37b 100644 --- a/src/transformers/models/aria/configuration_aria.py +++ b/src/transformers/models/aria/configuration_aria.py @@ -42,13 +42,13 @@ class AriaTextConfig(PreTrainedConfig): model_type = "aria_text" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), + "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.shared_experts.gate_proj": "colwise", + "layers.*.mlp.shared_experts.up_proj": "colwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", } base_model_sp_plan = { "embed_tokens": TPStyle("vocab", "reduce_scatter"), diff --git a/src/transformers/models/aria/modular_aria.py b/src/transformers/models/aria/modular_aria.py index d3484849b3a8..1e4c712f63d0 100644 --- a/src/transformers/models/aria/modular_aria.py +++ b/src/transformers/models/aria/modular_aria.py @@ -30,7 +30,6 @@ SizeDict, get_image_size, ) -from ...integrations.tensor_parallel import TPStyle from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_outputs import BaseModelOutputWithPooling from ...modeling_utils import PreTrainedModel @@ -111,13 +110,13 @@ class AriaTextConfig(LlamaConfig): model_type = "aria_text" base_config_key = "text_config" base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), + "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.shared_experts.gate_proj": "colwise", + "layers.*.mlp.shared_experts.up_proj": "colwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", } intermediate_size: int = 4096 diff --git a/src/transformers/models/cohere2/configuration_cohere2.py b/src/transformers/models/cohere2/configuration_cohere2.py index 9a0aca93cc20..f596e3cee86f 100644 --- a/src/transformers/models/cohere2/configuration_cohere2.py +++ b/src/transformers/models/cohere2/configuration_cohere2.py @@ -21,7 +21,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -50,28 +49,28 @@ class Cohere2Config(PreTrainedConfig): model_type = "cohere2" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), diff --git a/src/transformers/models/cohere2/modular_cohere2.py b/src/transformers/models/cohere2/modular_cohere2.py index e6272e1e8e16..574004647d0d 100644 --- a/src/transformers/models/cohere2/modular_cohere2.py +++ b/src/transformers/models/cohere2/modular_cohere2.py @@ -20,7 +20,6 @@ from ...cache_utils import Cache, DynamicCache from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_outputs import BaseModelOutputWithPast from ...modeling_rope_utils import ( @@ -71,28 +70,28 @@ class Cohere2Config(PreTrainedConfig): model_type = "cohere2" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), diff --git a/src/transformers/models/dbrx/modeling_dbrx.py b/src/transformers/models/dbrx/modeling_dbrx.py index db1d86c24a6b..f009aa23d2b9 100644 --- a/src/transformers/models/dbrx/modeling_dbrx.py +++ b/src/transformers/models/dbrx/modeling_dbrx.py @@ -29,7 +29,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_func_from_hub -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast @@ -643,8 +642,8 @@ def load_balancing_loss_func( class DbrxForCausalLM(DbrxPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "transformer.wte.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: DbrxConfig): diff --git a/src/transformers/models/dbrx/modular_dbrx.py b/src/transformers/models/dbrx/modular_dbrx.py index e3eabe9fbe08..500d8acc5915 100644 --- a/src/transformers/models/dbrx/modular_dbrx.py +++ b/src/transformers/models/dbrx/modular_dbrx.py @@ -23,7 +23,6 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import ( GradientCheckpointingLayer, @@ -431,8 +430,8 @@ def forward( class DbrxForCausalLM(DbrxPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "transformer.wte.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: DbrxConfig): diff --git a/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py b/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py index b0214bd3ed28..5ab83e16ad87 100644 --- a/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py +++ b/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py @@ -21,7 +21,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -53,49 +52,41 @@ class DeepseekV2Config(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.q_b_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.kv_b_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), - "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.q_b_proj": "colwise", + "layers.*.self_attn.kv_b_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", + "layers.*.mlp.shared_experts.gate_proj": "colwise", + "layers.*.mlp.shared_experts.up_proj": "colwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), + "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": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_b_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.kv_b_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather_split"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_b_proj": "colwise", + "layers.*.self_attn.kv_b_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts": "moe_experts_allreduce", # Shared experts output must stay Replicate to match experts output # (they're summed inside the MoE block, before the outer allgather_split # handles the SP boundary). - "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), diff --git a/src/transformers/models/deepseek_v2/modular_deepseek_v2.py b/src/transformers/models/deepseek_v2/modular_deepseek_v2.py index 7e3fc5d7fa25..a6c6494c0415 100644 --- a/src/transformers/models/deepseek_v2/modular_deepseek_v2.py +++ b/src/transformers/models/deepseek_v2/modular_deepseek_v2.py @@ -21,7 +21,6 @@ from ... import initialization as init from ...cache_utils import Cache -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters, dynamic_rope_update from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...utils import auto_docstring, logging @@ -68,49 +67,41 @@ class DeepseekV2Config(LlamaConfig): """ base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.q_b_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.kv_b_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), - "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.q_b_proj": "colwise", + "layers.*.self_attn.kv_b_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", + "layers.*.mlp.shared_experts.gate_proj": "colwise", + "layers.*.mlp.shared_experts.up_proj": "colwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), + "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": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_b_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.kv_b_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather_split"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_b_proj": "colwise", + "layers.*.self_attn.kv_b_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts": "moe_experts_allreduce", # Shared experts output must stay Replicate to match experts output # (they're summed inside the MoE block, before the outer allgather_split # handles the SP boundary). - "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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/doge/configuration_doge.py b/src/transformers/models/doge/configuration_doge.py index e4ccda76fafc..4fbbc393fb2a 100644 --- a/src/transformers/models/doge/configuration_doge.py +++ b/src/transformers/models/doge/configuration_doge.py @@ -22,7 +22,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -53,17 +52,17 @@ class DogeConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `DogeModel` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.dt_proj": TPStyle("rowwise", "allreduce"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.router_gate": TPStyle("colwise", "allgather"), - "layers.*.mlp.down_embed": TPStyle("vocab", "allreduce"), - "layers.*.mlp.up_embed": TPStyle("vocab", "allreduce"), + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "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_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"]), diff --git a/src/transformers/models/doge/modular_doge.py b/src/transformers/models/doge/modular_doge.py index cf92f343eb26..afe2469fc9b2 100644 --- a/src/transformers/models/doge/modular_doge.py +++ b/src/transformers/models/doge/modular_doge.py @@ -29,7 +29,6 @@ from ...cache_utils import Cache from ...configuration_utils import PreTrainedConfig from ...integrations.flex_attention import compile_friendly_flex_attention -from ...integrations.tensor_parallel import TPStyle from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast from ...modeling_rope_utils import RopeParameters @@ -82,17 +81,17 @@ class DogeConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `DogeModel` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.dt_proj": TPStyle("rowwise", "allreduce"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.router_gate": TPStyle("colwise", "allgather"), - "layers.*.mlp.down_embed": TPStyle("vocab", "allreduce"), - "layers.*.mlp.up_embed": TPStyle("vocab", "allreduce"), + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "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_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"]), diff --git a/src/transformers/models/dots1/configuration_dots1.py b/src/transformers/models/dots1/configuration_dots1.py index e452acf1802c..9bf2bce8fd4a 100644 --- a/src/transformers/models/dots1/configuration_dots1.py +++ b/src/transformers/models/dots1/configuration_dots1.py @@ -20,7 +20,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -49,21 +48,17 @@ class Dots1Config(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), - "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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": "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 = { diff --git a/src/transformers/models/dots1/modular_dots1.py b/src/transformers/models/dots1/modular_dots1.py index 86a9093a49b9..459f651e7b22 100644 --- a/src/transformers/models/dots1/modular_dots1.py +++ b/src/transformers/models/dots1/modular_dots1.py @@ -15,7 +15,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_outputs import CausalLMOutputWithPast from ...modeling_rope_utils import RopeParameters from ...processing_utils import Unpack @@ -64,21 +63,17 @@ class Dots1Config(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), - "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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": "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 = { 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 610b9647ee75..cd48b6adae67 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 @@ -20,7 +20,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring, logging @@ -50,10 +49,10 @@ class Ernie4_5_VLMoeVisionConfig(PreTrainedConfig): initializer_range: float = 0.02 base_model_tp_plan = { - "blocks.*.attn.qkv": TPStyle("colwise", "none"), - "blocks.*.attn.proj": TPStyle("rowwise", "allreduce"), - "blocks.*.mlp.fc1": TPStyle("colwise", "none"), - "blocks.*.mlp.fc2": TPStyle("rowwise", "allreduce"), + "blocks.*.attn.qkv": "colwise", + "blocks.*.attn.proj": "rowwise_allreduce", + "blocks.*.mlp.fc1": "colwise", + "blocks.*.mlp.fc2": "rowwise_allreduce", } intermediate_size: int = 4 * 1280 temporal_merge_size: int = 2 @@ -84,16 +83,16 @@ class Ernie4_5_VLMoeTextConfig(PreTrainedConfig): default_theta = 500000.0 base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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.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"]), 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 43e5d780f8c9..bf213108bab4 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 @@ -37,7 +37,6 @@ PILImageResampling, SizeDict, ) -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -94,10 +93,10 @@ class Ernie4_5_VLMoeVisionConfig(Qwen2VLVisionConfig): model_type = "ernie4_5_vl_moe_vision" base_model_tp_plan = { - "blocks.*.attn.qkv": TPStyle("colwise", "none"), - "blocks.*.attn.proj": TPStyle("rowwise", "allreduce"), - "blocks.*.mlp.fc1": TPStyle("colwise", "none"), - "blocks.*.mlp.fc2": TPStyle("rowwise", "allreduce"), + "blocks.*.attn.qkv": "colwise", + "blocks.*.attn.proj": "rowwise_allreduce", + "blocks.*.mlp.fc1": "colwise", + "blocks.*.mlp.fc2": "rowwise_allreduce", } hidden_size: int = 1280 @@ -132,16 +131,16 @@ class Ernie4_5_VLMoeTextConfig(Ernie4_5_MoeConfig): base_config_key = "text_config" base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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.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", } ignore_keys_at_rope_validation = {"mrope_section"} diff --git a/src/transformers/models/esm/configuration_esm.py b/src/transformers/models/esm/configuration_esm.py index a00dcf8b39e3..7875d88ecee8 100644 --- a/src/transformers/models/esm/configuration_esm.py +++ b/src/transformers/models/esm/configuration_esm.py @@ -159,12 +159,12 @@ class EsmConfig(PreTrainedConfig): mask_token_id (`int`, *optional*): The index of the mask token in the vocabulary. This must be included in the config because of the "mask-dropout" scaling trick, which will scale the inputs depending on the number of masked tokens. + rope_theta (`float`, defaults to 10000.0): + The base period of the RoPE embeddings. Only used when `position_embedding_type` is set to `"rotary"`. position_embedding_type (`str`, *optional*, defaults to `"absolute"`): Type of position embedding. Choose either `"absolute"` or "rotary"`. emb_layer_norm_before (`bool`, *optional*): Whether to apply layer normalization after embeddings but before the main stem of the network. - rope_theta (`float`, defaults to 10000.0): - The base period of the RoPE embeddings. Only used when `position_embedding_type` is set to `"rotary"`. token_dropout (`bool`, defaults to `False`): When this is enabled, masked tokens are treated as if they had been dropped out by input dropout. is_folding_model (`bool`, defaults to `False`): diff --git a/src/transformers/models/eurobert/modeling_eurobert.py b/src/transformers/models/eurobert/modeling_eurobert.py index 3d32c7f4bae8..149f0ef247d1 100644 --- a/src/transformers/models/eurobert/modeling_eurobert.py +++ b/src/transformers/models/eurobert/modeling_eurobert.py @@ -29,7 +29,6 @@ from ...activations import ACT2FN from ...cache_utils import Cache from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_bidirectional_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, MaskedLMOutput, SequenceClassifierOutput, TokenClassifierOutput @@ -409,8 +408,8 @@ def forward( @auto_docstring class EuroBertForMaskedLM(EuroBertPreTrainedModel): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: EuroBertConfig): diff --git a/src/transformers/models/eurobert/modular_eurobert.py b/src/transformers/models/eurobert/modular_eurobert.py index 1ae4cc7292af..fcfce231cd53 100644 --- a/src/transformers/models/eurobert/modular_eurobert.py +++ b/src/transformers/models/eurobert/modular_eurobert.py @@ -18,7 +18,6 @@ from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss from ...configuration_utils import strict -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_bidirectional_mask from ...modeling_outputs import BaseModelOutput, MaskedLMOutput, SequenceClassifierOutput, TokenClassifierOutput from ...modeling_rope_utils import RopeParameters @@ -142,8 +141,8 @@ def forward( @auto_docstring class EuroBertForMaskedLM(EuroBertPreTrainedModel): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: EuroBertConfig): diff --git a/src/transformers/models/exaone4/configuration_exaone4.py b/src/transformers/models/exaone4/configuration_exaone4.py index 652939a9c5e0..f52652f69b4f 100644 --- a/src/transformers/models/exaone4/configuration_exaone4.py +++ b/src/transformers/models/exaone4/configuration_exaone4.py @@ -21,7 +21,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -61,30 +60,30 @@ class Exaone4Config(PreTrainedConfig): 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": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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/exaone4/modular_exaone4.py b/src/transformers/models/exaone4/modular_exaone4.py index 5d4926ed33bf..ccbd0926290c 100644 --- a/src/transformers/models/exaone4/modular_exaone4.py +++ b/src/transformers/models/exaone4/modular_exaone4.py @@ -22,7 +22,6 @@ from ...cache_utils import Cache, DynamicCache from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_outputs import ( BaseModelOutputWithPast, @@ -90,30 +89,30 @@ class Exaone4Config(PreTrainedConfig): 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": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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/exaone_moe/configuration_exaone_moe.py b/src/transformers/models/exaone_moe/configuration_exaone_moe.py index 1dcb5399a49a..f76ed1c97baa 100644 --- a/src/transformers/models/exaone_moe/configuration_exaone_moe.py +++ b/src/transformers/models/exaone_moe/configuration_exaone_moe.py @@ -21,7 +21,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...utils import auto_docstring @@ -67,13 +66,13 @@ class ExaoneMoeConfig(PreTrainedConfig): 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": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_sp_plan = None diff --git a/src/transformers/models/flex_olmo/configuration_flex_olmo.py b/src/transformers/models/flex_olmo/configuration_flex_olmo.py index fb45082fd430..c80b792944a0 100644 --- a/src/transformers/models/flex_olmo/configuration_flex_olmo.py +++ b/src/transformers/models/flex_olmo/configuration_flex_olmo.py @@ -22,7 +22,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -51,18 +50,10 @@ class FlexOlmoConfig(PreTrainedConfig): attribute_map = {"num_local_experts": "num_experts"} default_theta = 500000.0 base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle( - "colwise", "allgather" - ), # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.k_proj": TPStyle( - "colwise", "allgather" - ), # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.v_proj": TPStyle( - "colwise", "allgather" - ), # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.o_proj": TPStyle( - "vocab", "allreduce" - ), # 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 } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/flex_olmo/modular_flex_olmo.py b/src/transformers/models/flex_olmo/modular_flex_olmo.py index 68e21e53fa30..e3dbd850f3ce 100644 --- a/src/transformers/models/flex_olmo/modular_flex_olmo.py +++ b/src/transformers/models/flex_olmo/modular_flex_olmo.py @@ -18,7 +18,6 @@ from ...cache_utils import Cache, DynamicCache from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_outputs import MoeModelOutputWithPast from ...modeling_rope_utils import RopeParameters @@ -61,18 +60,10 @@ class FlexOlmoConfig(PreTrainedConfig): attribute_map = {"num_local_experts": "num_experts"} default_theta = 500000.0 base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle( - "colwise", "allgather" - ), # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.k_proj": TPStyle( - "colwise", "allgather" - ), # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.v_proj": TPStyle( - "colwise", "allgather" - ), # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.o_proj": TPStyle( - "vocab", "allreduce" - ), # 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 } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/gemma/configuration_gemma.py b/src/transformers/models/gemma/configuration_gemma.py index 0bcde94f890f..77d5d3203277 100644 --- a/src/transformers/models/gemma/configuration_gemma.py +++ b/src/transformers/models/gemma/configuration_gemma.py @@ -23,7 +23,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -48,28 +47,28 @@ class GemmaConfig(PreTrainedConfig): model_type = "gemma" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), diff --git a/src/transformers/models/gemma/modular_gemma.py b/src/transformers/models/gemma/modular_gemma.py index e975c4c96ffd..06168ff81cbb 100644 --- a/src/transformers/models/gemma/modular_gemma.py +++ b/src/transformers/models/gemma/modular_gemma.py @@ -21,7 +21,6 @@ from ... import initialization as init from ...cache_utils import Cache, DynamicCache from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_outputs import BaseModelOutputWithPast from ...modeling_rope_utils import RopeParameters @@ -67,28 +66,28 @@ class GemmaConfig(PreTrainedConfig): model_type = "gemma" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), diff --git a/src/transformers/models/gemma2/configuration_gemma2.py b/src/transformers/models/gemma2/configuration_gemma2.py index 4b840136836d..673041502934 100644 --- a/src/transformers/models/gemma2/configuration_gemma2.py +++ b/src/transformers/models/gemma2/configuration_gemma2.py @@ -21,7 +21,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -52,28 +51,28 @@ class Gemma2Config(PreTrainedConfig): model_type = "gemma2" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), diff --git a/src/transformers/models/gemma2/modular_gemma2.py b/src/transformers/models/gemma2/modular_gemma2.py index 54c62724e526..6c52bff8d9dd 100644 --- a/src/transformers/models/gemma2/modular_gemma2.py +++ b/src/transformers/models/gemma2/modular_gemma2.py @@ -21,7 +21,6 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -79,28 +78,28 @@ class Gemma2Config(PreTrainedConfig): model_type = "gemma2" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), diff --git a/src/transformers/models/gemma3/configuration_gemma3.py b/src/transformers/models/gemma3/configuration_gemma3.py index 9b9a06d4ae10..0969e23f8f5e 100644 --- a/src/transformers/models/gemma3/configuration_gemma3.py +++ b/src/transformers/models/gemma3/configuration_gemma3.py @@ -23,7 +23,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...utils import auto_docstring, logging from ..siglip import SiglipVisionConfig @@ -59,30 +58,30 @@ class Gemma3TextConfig(PreTrainedConfig): model_type = "gemma3_text" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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/gemma3/modular_gemma3.py b/src/transformers/models/gemma3/modular_gemma3.py index b034bcfe87a5..aabc605fcf5c 100644 --- a/src/transformers/models/gemma3/modular_gemma3.py +++ b/src/transformers/models/gemma3/modular_gemma3.py @@ -22,7 +22,6 @@ from ... import initialization as init from ...cache_utils import Cache, DynamicCache from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_masks_for_generate, create_sliding_window_causal_mask from ...modeling_layers import GenericForSequenceClassification, GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, SequenceClassifierOutputWithPast @@ -87,30 +86,30 @@ class Gemma3TextConfig(Gemma2Config, PreTrainedConfig): model_type = "gemma3_text" base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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", } default_theta = {"global": 1_000_000.0, "local": 10_000.0} diff --git a/src/transformers/models/gemma3n/configuration_gemma3n.py b/src/transformers/models/gemma3n/configuration_gemma3n.py index 7d299c44ee92..48f7ec482372 100644 --- a/src/transformers/models/gemma3n/configuration_gemma3n.py +++ b/src/transformers/models/gemma3n/configuration_gemma3n.py @@ -24,7 +24,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...utils import auto_docstring, is_timm_available, logging, requires_backends @@ -80,13 +79,13 @@ class Gemma3nTextConfig(PreTrainedConfig): model_type = "gemma3n_text" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "allgather"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "allgather"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "allgather"), - "layers.*.self_attn.o_proj": TPStyle("vocab", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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_allreduce", } base_model_sp_plan = None base_model_pp_plan = { diff --git a/src/transformers/models/gemma3n/modular_gemma3n.py b/src/transformers/models/gemma3n/modular_gemma3n.py index 09aaf0e9fbd5..05105bcd4b1b 100644 --- a/src/transformers/models/gemma3n/modular_gemma3n.py +++ b/src/transformers/models/gemma3n/modular_gemma3n.py @@ -26,7 +26,6 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS @@ -107,13 +106,13 @@ class Gemma3nTextConfig(Gemma3TextConfig): model_type = "gemma3n_text" base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "allgather"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "allgather"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "allgather"), - "layers.*.self_attn.o_proj": TPStyle("vocab", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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_allreduce", } base_model_sp_plan = None default_theta = {"global": 1_000_000.0, "local": 10_000.0} diff --git a/src/transformers/models/glm4_moe/configuration_glm4_moe.py b/src/transformers/models/glm4_moe/configuration_glm4_moe.py index b09c9d61eef9..b08ba1db6d6f 100644 --- a/src/transformers/models/glm4_moe/configuration_glm4_moe.py +++ b/src/transformers/models/glm4_moe/configuration_glm4_moe.py @@ -20,7 +20,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -55,21 +54,17 @@ class Glm4MoeConfig(PreTrainedConfig): # Default tensor parallel plan for base model `Glm4Moe` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), - "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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": "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"]), diff --git a/src/transformers/models/glm4_moe/modular_glm4_moe.py b/src/transformers/models/glm4_moe/modular_glm4_moe.py index 0a0f6d9da610..f1caa381ec3b 100644 --- a/src/transformers/models/glm4_moe/modular_glm4_moe.py +++ b/src/transformers/models/glm4_moe/modular_glm4_moe.py @@ -18,7 +18,6 @@ from torch import nn from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring, logging from ..cohere.modeling_cohere import CohereAttention @@ -68,21 +67,17 @@ class Glm4MoeConfig(PreTrainedConfig): # Default tensor parallel plan for base model `Glm4Moe` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), - "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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": "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"]), 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 012d1d06526b..3885caa01e6e 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 @@ -22,7 +22,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -53,17 +52,13 @@ class Glm4MoeLiteConfig(PreTrainedConfig): model_type = "glm4_moe_lite" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_b_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.kv_b_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "layers.*.self_attn.q_b_proj": "colwise", + "layers.*.self_attn.kv_b_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), 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 06741613822e..41bf6b0aff0d 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 @@ -17,7 +17,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring from ..deepseek_v3.modeling_deepseek_v3 import DeepseekV3Attention @@ -61,17 +60,13 @@ class Glm4MoeLiteConfig(PreTrainedConfig): model_type = "glm4_moe_lite" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_b_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.kv_b_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "layers.*.self_attn.q_b_proj": "colwise", + "layers.*.self_attn.kv_b_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/glm4v/configuration_glm4v.py b/src/transformers/models/glm4v/configuration_glm4v.py index e887cd44f6d2..98200d3e3495 100644 --- a/src/transformers/models/glm4v/configuration_glm4v.py +++ b/src/transformers/models/glm4v/configuration_glm4v.py @@ -20,7 +20,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -91,14 +90,12 @@ class Glm4vTextConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `Glm4v` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_up_proj": TPStyle( - "colwise", "allgather" - ), # we need to replicate here due to the `chunk` operation - "layers.*.mlp.down_proj": TPStyle("vocab", "allreduce"), # input is replicated due to the `chunk` operation + "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_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"]), diff --git a/src/transformers/models/glm4v/modular_glm4v.py b/src/transformers/models/glm4v/modular_glm4v.py index ff53d2e287c6..cf2263263241 100644 --- a/src/transformers/models/glm4v/modular_glm4v.py +++ b/src/transformers/models/glm4v/modular_glm4v.py @@ -27,7 +27,6 @@ from ...configuration_utils import PreTrainedConfig from ...feature_extraction_utils import BatchFeature from ...image_utils import ImageInput -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -135,14 +134,12 @@ class Glm4vTextConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `Glm4v` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_up_proj": TPStyle( - "colwise", "allgather" - ), # we need to replicate here due to the `chunk` operation - "layers.*.mlp.down_proj": TPStyle("vocab", "allreduce"), # input is replicated due to the `chunk` operation + "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_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"]), diff --git a/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py b/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py index c95729daa67e..e5cbbdf2872c 100644 --- a/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py +++ b/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py @@ -20,7 +20,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -54,13 +53,13 @@ class Glm4vMoeTextConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `Glm4vMoe` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/glm4v_moe/modular_glm4v_moe.py b/src/transformers/models/glm4v_moe/modular_glm4v_moe.py index 05c48b004623..8116285ea881 100644 --- a/src/transformers/models/glm4v_moe/modular_glm4v_moe.py +++ b/src/transformers/models/glm4v_moe/modular_glm4v_moe.py @@ -19,7 +19,6 @@ from ... import initialization as init from ...cache_utils import Cache, DynamicCache -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_outputs import MoeModelOutputWithPast @@ -86,13 +85,13 @@ class Glm4vMoeTextConfig(Glm4MoeConfig): keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `Glm4vMoe` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/glm_image/configuration_glm_image.py b/src/transformers/models/glm_image/configuration_glm_image.py index de3b70617151..94f7c61dc4fa 100644 --- a/src/transformers/models/glm_image/configuration_glm_image.py +++ b/src/transformers/models/glm_image/configuration_glm_image.py @@ -21,7 +21,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -104,14 +103,12 @@ class GlmImageTextConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `GlmImage` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_up_proj": TPStyle( - "colwise", "allgather" - ), # we need to replicate here due to the `chunk` operation - "layers.*.mlp.down_proj": TPStyle("vocab", "allreduce"), # input is replicated due to the `chunk` operation + "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_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"]), 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 88de8f2f2e3e..4a1187f2389c 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 @@ -21,7 +21,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -58,20 +57,16 @@ class GlmMoeDsaConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_b_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.kv_b_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), - "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "layers.*.self_attn.q_b_proj": "colwise", + "layers.*.self_attn.kv_b_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", + "layers.*.mlp.shared_experts.gate_proj": "colwise", + "layers.*.mlp.shared_experts.up_proj": "colwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_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"]), 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 7ca32fdd2143..0271dbb6e238 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 @@ -21,7 +21,6 @@ from ...cache_utils import Cache from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_utils import ALL_ATTENTION_FUNCTIONS from ...models.llama.modeling_llama import rotate_half @@ -104,20 +103,16 @@ class GlmMoeDsaConfig(Glm4MoeLiteConfig): ```""" base_model_tp_plan = { - "layers.*.self_attn.q_b_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.kv_b_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), - "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "layers.*.self_attn.q_b_proj": "colwise", + "layers.*.self_attn.kv_b_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", + "layers.*.mlp.shared_experts.gate_proj": "colwise", + "layers.*.mlp.shared_experts.up_proj": "colwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } hidden_size: int = 6144 diff --git a/src/transformers/models/glm_ocr/configuration_glm_ocr.py b/src/transformers/models/glm_ocr/configuration_glm_ocr.py index fdad0285f965..dcd86079241d 100644 --- a/src/transformers/models/glm_ocr/configuration_glm_ocr.py +++ b/src/transformers/models/glm_ocr/configuration_glm_ocr.py @@ -21,7 +21,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -92,14 +91,12 @@ class GlmOcrTextConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `GlmOcr` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_up_proj": TPStyle( - "colwise", "allgather" - ), # we need to replicate here due to the `chunk` operation - "layers.*.mlp.down_proj": TPStyle("vocab", "allreduce"), # input is replicated due to the `chunk` operation + "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_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"]), diff --git a/src/transformers/models/gpt_neox/modeling_gpt_neox.py b/src/transformers/models/gpt_neox/modeling_gpt_neox.py index 7e4426f258c1..cfb81a96227d 100755 --- a/src/transformers/models/gpt_neox/modeling_gpt_neox.py +++ b/src/transformers/models/gpt_neox/modeling_gpt_neox.py @@ -13,7 +13,6 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -395,7 +394,7 @@ def set_input_embeddings(self, value): ) class GPTNeoXForCausalLM(GPTNeoXPreTrainedModel, GenerationMixin): _tied_weights_keys = {"embed_out.weight": "gpt_neox.embed_in.weight"} - _tp_plan = {"embed_out": TPStyle("colwise", "allgather")} + _tp_plan = {"embed_out": "colwise_allgather"} _pp_plan = {"embed_out": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/gpt_neox/modular_gpt_neox.py b/src/transformers/models/gpt_neox/modular_gpt_neox.py index 5833150550e2..cabda14021f0 100644 --- a/src/transformers/models/gpt_neox/modular_gpt_neox.py +++ b/src/transformers/models/gpt_neox/modular_gpt_neox.py @@ -7,7 +7,6 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -342,7 +341,7 @@ def forward( ) class GPTNeoXForCausalLM(GPTNeoXPreTrainedModel, GenerationMixin): _tied_weights_keys = {"embed_out.weight": "gpt_neox.embed_in.weight"} - _tp_plan = {"embed_out": TPStyle("colwise", "allgather")} + _tp_plan = {"embed_out": "colwise_allgather"} _pp_plan = {"embed_out": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/jais2/configuration_jais2.py b/src/transformers/models/jais2/configuration_jais2.py index b5f03d44cc99..5e0fa934f7a9 100644 --- a/src/transformers/models/jais2/configuration_jais2.py +++ b/src/transformers/models/jais2/configuration_jais2.py @@ -22,7 +22,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring from ...utils.type_validators import interval @@ -48,26 +47,26 @@ class Jais2Config(PreTrainedConfig): model_type = "jais2" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), diff --git a/src/transformers/models/jais2/modular_jais2.py b/src/transformers/models/jais2/modular_jais2.py index a5dc6b0302ba..c760e70c4b15 100644 --- a/src/transformers/models/jais2/modular_jais2.py +++ b/src/transformers/models/jais2/modular_jais2.py @@ -16,7 +16,6 @@ import torch.nn as nn from huggingface_hub.dataclasses import strict -from ...integrations.tensor_parallel import TPStyle from ...utils import auto_docstring, can_return_tuple from ..llama.configuration_llama import LlamaConfig from ..llama.modeling_llama import ( @@ -32,26 +31,26 @@ @strict class Jais2Config(LlamaConfig): base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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/minimax/configuration_minimax.py b/src/transformers/models/minimax/configuration_minimax.py index 699a768bc311..e53ecd12617f 100644 --- a/src/transformers/models/minimax/configuration_minimax.py +++ b/src/transformers/models/minimax/configuration_minimax.py @@ -21,7 +21,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -63,15 +62,11 @@ class MiniMaxConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] default_theta = 1000000.0 base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), + "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": "moe_experts_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/minimax/modular_minimax.py b/src/transformers/models/minimax/modular_minimax.py index 1dd83f5d148a..485417e274dd 100644 --- a/src/transformers/models/minimax/modular_minimax.py +++ b/src/transformers/models/minimax/modular_minimax.py @@ -23,7 +23,6 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -89,15 +88,11 @@ class MiniMaxConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] default_theta = 1000000.0 base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), + "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": "moe_experts_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/minimax_m2/configuration_minimax_m2.py b/src/transformers/models/minimax_m2/configuration_minimax_m2.py index 64cb731bc2c9..2f360c509c28 100644 --- a/src/transformers/models/minimax_m2/configuration_minimax_m2.py +++ b/src/transformers/models/minimax_m2/configuration_minimax_m2.py @@ -22,7 +22,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -49,32 +48,24 @@ class MiniMaxM2Config(PreTrainedConfig): model_type = "minimax_m2" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "allgather"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "allgather"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "allgather"), - "layers.*.self_attn.o_proj": TPStyle("vocab", "allreduce"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_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.experts": "moe_experts_allreduce", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "allgather"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "allgather"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "allgather"), - "layers.*.self_attn.o_proj": TPStyle("vocab", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather_split"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), - "norm": TPStyle("activation", "none"), + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise_allgather", + "layers.*.self_attn.k_proj": "colwise_allgather", + "layers.*.self_attn.v_proj": "colwise_allgather", + "layers.*.self_attn.o_proj": "vocab_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts": "moe_experts_allreduce", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/minimax_m2/modular_minimax_m2.py b/src/transformers/models/minimax_m2/modular_minimax_m2.py index cf62aae16d7d..127638fb38ab 100644 --- a/src/transformers/models/minimax_m2/modular_minimax_m2.py +++ b/src/transformers/models/minimax_m2/modular_minimax_m2.py @@ -21,7 +21,6 @@ from ... import initialization as init from ...cache_utils import Cache, DynamicCache from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_outputs import MoeModelOutputWithPast from ...modeling_rope_utils import RopeParameters @@ -68,32 +67,24 @@ class MiniMaxM2Config(PreTrainedConfig): model_type = "minimax_m2" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "allgather"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "allgather"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "allgather"), - "layers.*.self_attn.o_proj": TPStyle("vocab", "allreduce"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_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.experts": "moe_experts_allreduce", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "allgather"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "allgather"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "allgather"), - "layers.*.self_attn.o_proj": TPStyle("vocab", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather_split"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), - "norm": TPStyle("activation", "none"), + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise_allgather", + "layers.*.self_attn.k_proj": "colwise_allgather", + "layers.*.self_attn.v_proj": "colwise_allgather", + "layers.*.self_attn.o_proj": "vocab_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts": "moe_experts_allreduce", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/nanochat/modeling_nanochat.py b/src/transformers/models/nanochat/modeling_nanochat.py index 5431b5492da9..4f5e1b7fe3c7 100644 --- a/src/transformers/models/nanochat/modeling_nanochat.py +++ b/src/transformers/models/nanochat/modeling_nanochat.py @@ -30,7 +30,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_func_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -433,8 +432,8 @@ def forward( @auto_docstring class NanoChatForCausalLM(NanoChatPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/nanochat/modular_nanochat.py b/src/transformers/models/nanochat/modular_nanochat.py index b70ab836b62d..6bc3ddba6457 100644 --- a/src/transformers/models/nanochat/modular_nanochat.py +++ b/src/transformers/models/nanochat/modular_nanochat.py @@ -20,7 +20,6 @@ from ... import initialization as init from ...cache_utils import Cache, DynamicCache -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel @@ -199,8 +198,8 @@ def forward( @auto_docstring class NanoChatForCausalLM(Gemma2ForCausalLM): - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} def forward(self, **super_kwargs) -> CausalLMOutputWithPast: r""" diff --git a/src/transformers/models/olmo2/configuration_olmo2.py b/src/transformers/models/olmo2/configuration_olmo2.py index 8b6b745812b3..408e15b4bffd 100644 --- a/src/transformers/models/olmo2/configuration_olmo2.py +++ b/src/transformers/models/olmo2/configuration_olmo2.py @@ -26,7 +26,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -54,38 +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": TPStyle( - "colwise", "allgather" - ), # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.k_proj": TPStyle( - "colwise", "allgather" - ), # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.v_proj": TPStyle( - "colwise", "allgather" - ), # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.o_proj": TPStyle( - "vocab", "allreduce" - ), # input is replicated due to the added norm on q and k - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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_allreduce", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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/olmo2/modular_olmo2.py b/src/transformers/models/olmo2/modular_olmo2.py index a4aac76c44d6..d560317f19a2 100644 --- a/src/transformers/models/olmo2/modular_olmo2.py +++ b/src/transformers/models/olmo2/modular_olmo2.py @@ -26,7 +26,6 @@ from transformers.utils.generic import TransformersKwargs from ...cache_utils import Cache -from ...integrations.tensor_parallel import TPStyle from ...modeling_utils import ALL_ATTENTION_FUNCTIONS from ...processing_utils import Unpack from ...utils import auto_docstring, logging @@ -67,38 +66,30 @@ class Olmo2Config(OlmoConfig): model_type = "olmo2" base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle( - "colwise", "allgather" - ), # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.k_proj": TPStyle( - "colwise", "allgather" - ), # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.v_proj": TPStyle( - "colwise", "allgather" - ), # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.o_proj": TPStyle( - "vocab", "allreduce" - ), # input is replicated due to the added norm on q and k - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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_allreduce", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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 a317fcdb80c9..c4597535850c 100644 --- a/src/transformers/models/olmo3/configuration_olmo3.py +++ b/src/transformers/models/olmo3/configuration_olmo3.py @@ -21,7 +21,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -49,38 +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": TPStyle( - "colwise", "allgather" - ), # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.k_proj": TPStyle( - "colwise", "allgather" - ), # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.v_proj": TPStyle( - "colwise", "allgather" - ), # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.o_proj": TPStyle( - "vocab", "allreduce" - ), # input is replicated due to the added norm on q and k - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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_allreduce", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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/modular_olmo3.py b/src/transformers/models/olmo3/modular_olmo3.py index aa70895ab944..bc68746a7c52 100644 --- a/src/transformers/models/olmo3/modular_olmo3.py +++ b/src/transformers/models/olmo3/modular_olmo3.py @@ -19,7 +19,6 @@ from huggingface_hub.dataclasses import strict from ...cache_utils import Cache, DynamicCache -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_outputs import BaseModelOutputWithPast from ...modeling_utils import ALL_ATTENTION_FUNCTIONS @@ -63,38 +62,30 @@ class Olmo3Config(Olmo2Config): model_type = "olmo3" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle( - "colwise", "allgather" - ), # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.k_proj": TPStyle( - "colwise", "allgather" - ), # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.v_proj": TPStyle( - "colwise", "allgather" - ), # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.o_proj": TPStyle( - "vocab", "allreduce" - ), # input is replicated due to the added norm on q and k - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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_allreduce", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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 5a5f2b1d8a33..a6eabbb7b99b 100644 --- a/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py +++ b/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py @@ -21,7 +21,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring from ...utils.type_validators import interval @@ -74,21 +73,13 @@ class OlmoHybridConfig(PreTrainedConfig): model_type = "olmo_hybrid" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle( - "colwise", "allgather" - ), # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.k_proj": TPStyle( - "colwise", "allgather" - ), # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.v_proj": TPStyle( - "colwise", "allgather" - ), # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.o_proj": TPStyle( - "vocab", "allreduce" - ), # input is replicated due to the added norm on q and k - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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_allreduce", } base_model_sp_plan = None base_model_pp_plan = { diff --git a/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py index c8b05d83532d..7c883ec5f823 100644 --- a/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py +++ b/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py @@ -27,7 +27,6 @@ from ...activations import ACT2FN from ...cache_utils import Cache from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_outputs import BaseModelOutputWithPast from ...modeling_rope_utils import dynamic_rope_update @@ -121,21 +120,13 @@ class OlmoHybridConfig(LlamaConfig): model_type = "olmo_hybrid" base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle( - "colwise", "allgather" - ), # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.k_proj": TPStyle( - "colwise", "allgather" - ), # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.v_proj": TPStyle( - "colwise", "allgather" - ), # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.o_proj": TPStyle( - "vocab", "allreduce" - ), # input is replicated due to the added norm on q and k - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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_allreduce", } base_model_sp_plan = None diff --git a/src/transformers/models/pi0/modeling_pi0.py b/src/transformers/models/pi0/modeling_pi0.py index a816d633bd04..cfd11f20a3a4 100644 --- a/src/transformers/models/pi0/modeling_pi0.py +++ b/src/transformers/models/pi0/modeling_pi0.py @@ -27,7 +27,6 @@ from ... import initialization as init from ...cache_utils import Cache -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_bidirectional_mask from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from ...modeling_utils import PreTrainedModel @@ -227,7 +226,7 @@ def forward( class PI0ForConditionalGeneration(PI0PreTrainedModel): """PI0 model with action projection heads and flow matching.""" - _tp_plan = {"action_out_proj": TPStyle("colwise", "allgather")} + _tp_plan = {"action_out_proj": "colwise_allgather"} def __init__(self, config: PI0Config): super().__init__(config) diff --git a/src/transformers/models/pi0/modular_pi0.py b/src/transformers/models/pi0/modular_pi0.py index dfa25e9fc121..c34309bebba2 100644 --- a/src/transformers/models/pi0/modular_pi0.py +++ b/src/transformers/models/pi0/modular_pi0.py @@ -27,7 +27,6 @@ from ...configuration_utils import PreTrainedConfig from ...feature_extraction_utils import BatchFeature from ...image_utils import ImageInput, make_nested_list_of_images -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_bidirectional_mask from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from ...modeling_utils import PreTrainedModel @@ -477,7 +476,7 @@ def forward( class PI0ForConditionalGeneration(PI0PreTrainedModel): """PI0 model with action projection heads and flow matching.""" - _tp_plan = {"action_out_proj": TPStyle("colwise", "allgather")} + _tp_plan = {"action_out_proj": "colwise_allgather"} def __init__(self, config: PI0Config): 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 2f8b406315ae..c6b0311a039e 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 @@ -21,7 +21,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring, logging @@ -150,13 +149,13 @@ class Qwen2_5OmniTextConfig(PreTrainedConfig): # Default tensor parallel plan for base model `Qwen25OmniText` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), 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 75ba364d0656..6906856787af 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 @@ -30,7 +30,6 @@ from ...cache_utils import Cache from ...configuration_utils import PreTrainedConfig from ...generation import GenerationMixin -from ...integrations.tensor_parallel import TPStyle from ...modeling_outputs import BaseModelOutputWithPooling, ModelOutput from ...modeling_rope_utils import RopeParameters from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel @@ -162,13 +161,13 @@ class Qwen2_5OmniTextConfig(PreTrainedConfig): # Default tensor parallel plan for base model `Qwen25OmniText` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py b/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py index 7d3a1d2c5209..1afe8ac3c2d2 100644 --- a/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py +++ b/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py @@ -40,7 +40,6 @@ use_kernel_func_from_hub, use_kernelized_func, ) -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_layers import ( GenericForQuestionAnswering, @@ -618,8 +617,8 @@ 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": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/qwen2_moe/modular_qwen2_moe.py b/src/transformers/models/qwen2_moe/modular_qwen2_moe.py index d351dfa764e6..8a5e9fb7751b 100644 --- a/src/transformers/models/qwen2_moe/modular_qwen2_moe.py +++ b/src/transformers/models/qwen2_moe/modular_qwen2_moe.py @@ -25,7 +25,6 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_layers import ( GenericForQuestionAnswering, @@ -231,8 +230,8 @@ def forward( class Qwen2MoeForCausalLM(MixtralForCausalLM, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/qwen3_5/configuration_qwen3_5.py b/src/transformers/models/qwen3_5/configuration_qwen3_5.py index 2f2f6ecca1b9..a955c38f3e0b 100644 --- a/src/transformers/models/qwen3_5/configuration_qwen3_5.py +++ b/src/transformers/models/qwen3_5/configuration_qwen3_5.py @@ -20,7 +20,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -58,13 +57,13 @@ class Qwen3_5TextConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/qwen3_5/modular_qwen3_5.py b/src/transformers/models/qwen3_5/modular_qwen3_5.py index f01e3392c237..ec06f5f2daa4 100644 --- a/src/transformers/models/qwen3_5/modular_qwen3_5.py +++ b/src/transformers/models/qwen3_5/modular_qwen3_5.py @@ -22,7 +22,6 @@ from ... import initialization as init from ...cache_utils import Cache, DynamicCache -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GenericForSequenceClassification, GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling @@ -89,13 +88,13 @@ class Qwen3_5TextConfig(Qwen3NextConfig): base_config_key = "text_config" base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } 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 1f7455c5be9e..fda2ef77dd3b 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 @@ -20,7 +20,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -58,18 +57,14 @@ class Qwen3_5MoeTextConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), - "layers.*.mlp.shared_expert.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_expert.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_expert.down_proj": TPStyle("rowwise", "allreduce"), + "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": "moe_experts_allreduce", + "layers.*.mlp.shared_expert.gate_proj": "colwise", + "layers.*.mlp.shared_expert.up_proj": "colwise", + "layers.*.mlp.shared_expert.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), 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 f3afa5bdc169..d2c9ff7df97f 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 @@ -2003,8 +2003,8 @@ class Qwen3_5MoeForConditionalGeneration(Qwen3_5MoePreTrainedModel, GenerationMi # Reference: fix gemma3 grad acc #37208 accepts_loss_kwargs = False config: Qwen3_5MoeConfig - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/qwen3_5_moe/modular_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/modular_qwen3_5_moe.py index 428cd54248d9..2c4ad2095b74 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 @@ -17,7 +17,6 @@ from huggingface_hub.dataclasses import strict from ... import initialization as init -from ...integrations.tensor_parallel import TPStyle from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPooling from ...modeling_utils import PreTrainedModel @@ -87,18 +86,14 @@ class Qwen3_5MoeTextConfig(Qwen3NextConfig): base_config_key = "text_config" base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), - "layers.*.mlp.shared_expert.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_expert.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_expert.down_proj": TPStyle("rowwise", "allreduce"), + "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": "moe_experts_allreduce", + "layers.*.mlp.shared_expert.gate_proj": "colwise", + "layers.*.mlp.shared_expert.up_proj": "colwise", + "layers.*.mlp.shared_expert.down_proj": "rowwise_allreduce", } ignore_keys_at_rope_validation = {"mrope_section", "mrope_interleaved"} @@ -249,8 +244,8 @@ def __init__(self, config): class Qwen3_5MoeForConditionalGeneration(Qwen3VLMoeForConditionalGeneration): - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} def forward(self, **super_kwargs): r""" 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 bc6d86ae7edb..f511c9114478 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 @@ -131,13 +131,13 @@ class Qwen3OmniMoeTextConfig(PreTrainedConfig): # Default tensor parallel plan for base model `Qwen3OmniMoeText` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), 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 e1a53909071d..1308bcf1fcfa 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 @@ -3006,7 +3006,7 @@ def get_input_embeddings(self): @auto_docstring class Qwen3OmniMoeTalkerForConditionalGeneration(Qwen3OmniMoeThinkerTextPreTrainedModel, GenerationMixin): _tied_weights_keys = {"codec_head": "model.codec_embedding.weight"} - _tp_plan = {"codec_head": TPStyle("colwise", "allgather")} + _tp_plan = {"codec_head": "colwise_allgather"} _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} _pp_plan = {"codec_head": (["hidden_states"], ["logits"])} config_class = Qwen3OmniMoeTalkerConfig 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 8b8cbdc8218e..55315ec5d6b0 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 @@ -32,7 +32,6 @@ from ...feature_extraction_utils import BatchFeature from ...generation import GenerationMixin from ...image_utils import ImageInput -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( @@ -183,13 +182,13 @@ class Qwen3OmniMoeTextConfig(PreTrainedConfig): # Default tensor parallel plan for base model `Qwen3OmniMoeText` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -1577,7 +1576,7 @@ def get_input_embeddings(self): class Qwen3OmniMoeTalkerForConditionalGeneration(Qwen3MoeForCausalLM): _tied_weights_keys = {"codec_head": "model.codec_embedding.weight"} - _tp_plan = {"codec_head": TPStyle("colwise", "allgather")} + _tp_plan = {"codec_head": "colwise_allgather"} _pp_plan = {"codec_head": (["hidden_states"], ["logits"])} config_class = Qwen3OmniMoeTalkerConfig base_model_prefix = "talker" diff --git a/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py b/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py index 68fec77811af..997dbe290a71 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 @@ -57,13 +57,13 @@ class Qwen3VLMoeTextConfig(PreTrainedConfig): } # Default tensor parallel plan for base model `Qwen3VLMoe` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_sp_plan = { "embed_tokens": TPStyle("vocab", "reduce_scatter"), 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 1d6b4b3fcfd5..0d18a89ba7ab 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 @@ -20,7 +20,6 @@ from ... import initialization as init from ...cache_utils import Cache, DynamicCache -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_outputs import MoeModelOutputWithPast @@ -84,13 +83,13 @@ class Qwen3VLMoeTextConfig(Qwen3MoeConfig): default_theta = 500000.0 # Default tensor parallel plan for base model `Qwen3VLMoe` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/smollm3/configuration_smollm3.py b/src/transformers/models/smollm3/configuration_smollm3.py index 9b5a113b7c2f..5de4e2c1cedf 100644 --- a/src/transformers/models/smollm3/configuration_smollm3.py +++ b/src/transformers/models/smollm3/configuration_smollm3.py @@ -21,7 +21,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -56,28 +55,28 @@ class SmolLM3Config(PreTrainedConfig): default_theta = 2000000.0 base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), diff --git a/src/transformers/models/smollm3/modular_smollm3.py b/src/transformers/models/smollm3/modular_smollm3.py index 89f0c813f3f9..2fb76822c447 100644 --- a/src/transformers/models/smollm3/modular_smollm3.py +++ b/src/transformers/models/smollm3/modular_smollm3.py @@ -19,7 +19,6 @@ from ...cache_utils import Cache from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_rope_utils import RopeParameters from ...modeling_utils import ALL_ATTENTION_FUNCTIONS @@ -72,28 +71,28 @@ class SmolLM3Config(PreTrainedConfig): default_theta = 2000000.0 base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), diff --git a/src/transformers/models/solar_open/configuration_solar_open.py b/src/transformers/models/solar_open/configuration_solar_open.py index e0e316dbd604..685ea1cbe4a5 100644 --- a/src/transformers/models/solar_open/configuration_solar_open.py +++ b/src/transformers/models/solar_open/configuration_solar_open.py @@ -21,7 +21,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -39,15 +38,11 @@ class SolarOpenConfig(PreTrainedConfig): # Default tensor parallel plan for base model `SolarOpenModel` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), + "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": "moe_experts_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -85,21 +80,17 @@ class SolarOpenConfig(PreTrainedConfig): pad_token_id: int | None = None default_theta = 1_000_000.0 base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather_split"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), - "norm": TPStyle("activation", "none"), + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts": "moe_experts_allreduce", + "norm": "activation", } head_dim: int = 128 diff --git a/src/transformers/models/solar_open/modular_solar_open.py b/src/transformers/models/solar_open/modular_solar_open.py index dcf766bbb7aa..92b212fcfbdd 100644 --- a/src/transformers/models/solar_open/modular_solar_open.py +++ b/src/transformers/models/solar_open/modular_solar_open.py @@ -16,7 +16,6 @@ from huggingface_hub.dataclasses import strict from torch import nn -from ...integrations.tensor_parallel import TPStyle from ...utils import auto_docstring, logging from ..glm4_moe.configuration_glm4_moe import Glm4MoeConfig from ..glm4_moe.modeling_glm4_moe import ( @@ -45,32 +44,24 @@ class SolarOpenConfig(Glm4MoeConfig): # Default tensor parallel plan for base model `SolarOpenModel` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), + "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": "moe_experts_allreduce", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather_split"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), - "norm": TPStyle("activation", "none"), + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts": "moe_experts_allreduce", + "norm": "activation", } vocab_size: int = 196608 diff --git a/src/transformers/models/t5gemma/configuration_t5gemma.py b/src/transformers/models/t5gemma/configuration_t5gemma.py index ba4e2282cb70..d5e97d034d73 100644 --- a/src/transformers/models/t5gemma/configuration_t5gemma.py +++ b/src/transformers/models/t5gemma/configuration_t5gemma.py @@ -23,7 +23,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -52,28 +51,28 @@ class T5GemmaModuleConfig(PreTrainedConfig): model_type = "t5_gemma_module" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), diff --git a/src/transformers/models/t5gemma/modeling_t5gemma.py b/src/transformers/models/t5gemma/modeling_t5gemma.py index 65f4cb52d846..6b7a6e358283 100644 --- a/src/transformers/models/t5gemma/modeling_t5gemma.py +++ b/src/transformers/models/t5gemma/modeling_t5gemma.py @@ -29,7 +29,6 @@ from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache from ...generation import GenerationMixin from ...integrations import use_kernel_func_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import ( create_bidirectional_mask, create_bidirectional_sliding_window_mask, @@ -946,7 +945,7 @@ def forward( class T5GemmaForConditionalGeneration(T5GemmaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.out_proj.weight": "model.decoder.embed_tokens.weight"} - _tp_plan = {"lm_head.out_proj": TPStyle("colwise", "allgather")} + _tp_plan = {"lm_head.out_proj": "colwise_allgather"} _pp_plan = {"lm_head.out_proj": (["hidden_states"], ["logits"])} def __init__(self, config: T5GemmaConfig): diff --git a/src/transformers/models/t5gemma/modular_t5gemma.py b/src/transformers/models/t5gemma/modular_t5gemma.py index 363e3ca76430..58dc0c0b41bb 100644 --- a/src/transformers/models/t5gemma/modular_t5gemma.py +++ b/src/transformers/models/t5gemma/modular_t5gemma.py @@ -23,7 +23,6 @@ from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache from ...configuration_utils import PreTrainedConfig from ...generation import GenerationMixin -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import ( create_bidirectional_mask, create_bidirectional_sliding_window_mask, @@ -785,7 +784,7 @@ def forward( class T5GemmaForConditionalGeneration(T5GemmaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.out_proj.weight": "model.decoder.embed_tokens.weight"} - _tp_plan = {"lm_head.out_proj": TPStyle("colwise", "allgather")} + _tp_plan = {"lm_head.out_proj": "colwise_allgather"} _pp_plan = {"lm_head.out_proj": (["hidden_states"], ["logits"])} def __init__(self, config: T5GemmaConfig): diff --git a/src/transformers/models/t5gemma2/configuration_t5gemma2.py b/src/transformers/models/t5gemma2/configuration_t5gemma2.py index 38f88e669360..b90174f4f9ec 100644 --- a/src/transformers/models/t5gemma2/configuration_t5gemma2.py +++ b/src/transformers/models/t5gemma2/configuration_t5gemma2.py @@ -23,7 +23,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...utils import auto_docstring, logging from ..siglip import SiglipVisionConfig @@ -46,30 +45,30 @@ class T5Gemma2TextConfig(PreTrainedConfig): model_type = "t5gemma2_text" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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,30 +235,30 @@ class T5Gemma2DecoderConfig(PreTrainedConfig): model_type = "t5gemma2_decoder" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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/t5gemma2/modeling_t5gemma2.py b/src/transformers/models/t5gemma2/modeling_t5gemma2.py index 3de40aeaec64..811d561f971a 100644 --- a/src/transformers/models/t5gemma2/modeling_t5gemma2.py +++ b/src/transformers/models/t5gemma2/modeling_t5gemma2.py @@ -30,7 +30,6 @@ from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache, StaticCache from ...generation import GenerationConfig, GenerationMixin, GenerationMode from ...integrations import use_kernel_func_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_bidirectional_mask, create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -1186,7 +1185,7 @@ class T5Gemma2ForConditionalGeneration(T5Gemma2PreTrainedModel, GenerationMixin) _tied_weights_keys = { "lm_head.out_proj.weight": "model.encoder.text_model.embed_tokens.weight", } - _tp_plan = {"lm_head.out_proj": TPStyle("colwise", "allgather")} + _tp_plan = {"lm_head.out_proj": "colwise_allgather"} _pp_plan = {"lm_head.out_proj": (["hidden_states"], ["logits"])} def __init__(self, config: T5Gemma2Config): diff --git a/src/transformers/models/t5gemma2/modular_t5gemma2.py b/src/transformers/models/t5gemma2/modular_t5gemma2.py index db8b82ec6b88..c022eb55791b 100644 --- a/src/transformers/models/t5gemma2/modular_t5gemma2.py +++ b/src/transformers/models/t5gemma2/modular_t5gemma2.py @@ -24,7 +24,6 @@ from ...cache_utils import DynamicCache, EncoderDecoderCache, StaticCache from ...configuration_utils import PreTrainedConfig from ...generation import GenerationConfig, GenerationMixin, GenerationMode -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_bidirectional_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_outputs import ( @@ -974,7 +973,7 @@ class T5Gemma2ForConditionalGeneration(T5Gemma2PreTrainedModel, GenerationMixin) _tied_weights_keys = { "lm_head.out_proj.weight": "model.encoder.text_model.embed_tokens.weight", } - _tp_plan = {"lm_head.out_proj": TPStyle("colwise", "allgather")} + _tp_plan = {"lm_head.out_proj": "colwise_allgather"} _pp_plan = {"lm_head.out_proj": (["hidden_states"], ["logits"])} def __init__(self, config: T5Gemma2Config): diff --git a/src/transformers/models/vaultgemma/configuration_vaultgemma.py b/src/transformers/models/vaultgemma/configuration_vaultgemma.py index 3b535d51bfa9..be5737835253 100644 --- a/src/transformers/models/vaultgemma/configuration_vaultgemma.py +++ b/src/transformers/models/vaultgemma/configuration_vaultgemma.py @@ -22,7 +22,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -51,28 +50,28 @@ class VaultGemmaConfig(PreTrainedConfig): model_type = "vaultgemma" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), diff --git a/src/transformers/models/youtu/configuration_youtu.py b/src/transformers/models/youtu/configuration_youtu.py index 9a3592c871e9..6da28284a577 100644 --- a/src/transformers/models/youtu/configuration_youtu.py +++ b/src/transformers/models/youtu/configuration_youtu.py @@ -27,7 +27,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -52,9 +51,9 @@ class YoutuConfig(PreTrainedConfig): model_type = "youtu" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("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"]), diff --git a/src/transformers/models/youtu/modular_youtu.py b/src/transformers/models/youtu/modular_youtu.py index 09ccb4e8063d..f3218061670c 100644 --- a/src/transformers/models/youtu/modular_youtu.py +++ b/src/transformers/models/youtu/modular_youtu.py @@ -23,7 +23,6 @@ from torch import nn from ... import initialization as init -from ...integrations.tensor_parallel import TPStyle from ...modeling_utils import PreTrainedModel from ...utils import auto_docstring, logging from ..deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config @@ -61,9 +60,9 @@ class YoutuConfig(DeepseekV3Config): model_type = "youtu" base_model_tp_plan = { - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_sp_plan = None attribute_map = {} From e70ac3785c1bdd962b8f3e85b132f24ea32e1b3c Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 28 Apr 2026 10:23:51 +0000 Subject: [PATCH 053/116] migrate standalone configs and modelings to string-based TP plans --- .../integrations/tensor_parallel.py | 4 +- .../models/apertus/modeling_apertus.py | 5 +- .../models/arcee/modeling_arcee.py | 5 +- .../models/aria/configuration_aria.py | 27 +++-- src/transformers/models/aria/modeling_aria.py | 5 +- .../models/bamba/modeling_bamba.py | 5 +- .../models/cohere/configuration_cohere.py | 45 ++++---- .../models/cohere/modeling_cohere.py | 5 +- .../models/cohere2/modeling_cohere2.py | 5 +- .../models/cwm/configuration_cwm.py | 41 ++++--- src/transformers/models/cwm/modeling_cwm.py | 5 +- .../deepseek_v2/modeling_deepseek_v2.py | 5 +- .../deepseek_v3/configuration_deepseek_v3.py | 22 ++-- .../deepseek_v3/modeling_deepseek_v3.py | 5 +- .../models/diffllama/modeling_diffllama.py | 5 +- src/transformers/models/doge/modeling_doge.py | 5 +- .../models/dots1/modeling_dots1.py | 5 +- src/transformers/models/emu3/modeling_emu3.py | 5 +- .../models/ernie4_5/configuration_ernie4_5.py | 41 ++++--- .../models/ernie4_5/modeling_ernie4_5.py | 5 +- .../configuration_ernie4_5_moe.py | 27 ++--- .../ernie4_5_moe/modeling_ernie4_5_moe.py | 5 +- .../models/exaone4/modeling_exaone4.py | 5 +- .../models/exaone_moe/modeling_exaone_moe.py | 5 +- .../models/falcon_h1/modeling_falcon_h1.py | 5 +- .../models/flex_olmo/modeling_flex_olmo.py | 5 +- .../models/gemma/modeling_gemma.py | 5 +- .../models/gemma2/modeling_gemma2.py | 5 +- .../models/gemma3/modeling_gemma3.py | 5 +- .../models/gemma3n/modeling_gemma3n.py | 5 +- .../models/gemma4/configuration_gemma4.py | 29 +++-- .../models/gemma4/modeling_gemma4.py | 5 +- .../models/glm/configuration_glm.py | 37 ++++--- src/transformers/models/glm/modeling_glm.py | 5 +- .../models/glm4/configuration_glm4.py | 39 ++++--- src/transformers/models/glm4/modeling_glm4.py | 5 +- .../models/glm4_moe/modeling_glm4_moe.py | 5 +- .../glm4_moe_lite/modeling_glm4_moe_lite.py | 5 +- .../glm_moe_dsa/modeling_glm_moe_dsa.py | 5 +- .../models/gpt_neox/configuration_gpt_neox.py | 23 ++-- .../models/gpt_oss/modeling_gpt_oss.py | 5 +- .../models/granite/configuration_granite.py | 41 ++++--- .../models/granite/modeling_granite.py | 5 +- .../models/granitemoe/modeling_granitemoe.py | 5 +- .../modeling_granitemoehybrid.py | 5 +- .../modeling_granitemoeshared.py | 5 +- .../models/helium/configuration_helium.py | 41 ++++--- .../models/helium/modeling_helium.py | 5 +- .../configuration_higgs_audio_v2.py | 41 ++++--- .../modeling_hunyuan_v1_dense.py | 5 +- .../hunyuan_v1_moe/modeling_hunyuan_v1_moe.py | 5 +- .../models/jais2/modeling_jais2.py | 5 +- .../models/jamba/modeling_jamba.py | 5 +- .../modeling_kyutai_speech_to_text.py | 5 +- src/transformers/models/lfm2/modeling_lfm2.py | 5 +- .../models/lfm2_moe/modeling_lfm2_moe.py | 5 +- .../models/llama/configuration_llama.py | 41 ++++--- .../models/llama/modeling_llama.py | 5 +- .../models/llama4/configuration_llama4.py | 35 +++--- .../models/llama4/modeling_llama4.py | 5 +- .../configuration_longcat_flash.py | 13 ++- .../longcat_flash/modeling_longcat_flash.py | 5 +- .../models/minimax/modeling_minimax.py | 5 +- .../models/minimax_m2/modeling_minimax_m2.py | 5 +- .../ministral/configuration_ministral.py | 41 ++++--- .../models/ministral/modeling_ministral.py | 5 +- .../ministral3/configuration_ministral3.py | 41 ++++--- .../models/ministral3/modeling_ministral3.py | 5 +- .../models/mistral/configuration_mistral.py | 41 ++++--- .../models/mistral/modeling_mistral.py | 5 +- .../models/mistral4/configuration_mistral4.py | 19 ++-- .../models/mistral4/modeling_mistral4.py | 5 +- .../models/mixtral/configuration_mixtral.py | 47 +++----- .../models/mixtral/modeling_mixtral.py | 5 +- .../models/nanochat/configuration_nanochat.py | 13 ++- .../models/olmo/configuration_olmo.py | 41 ++++--- src/transformers/models/olmo/modeling_olmo.py | 5 +- .../models/olmo2/modeling_olmo2.py | 5 +- .../models/olmo3/modeling_olmo3.py | 5 +- .../olmo_hybrid/modeling_olmo_hybrid.py | 5 +- .../models/olmoe/configuration_olmoe.py | 41 +++---- .../models/olmoe/modeling_olmoe.py | 5 +- .../configuration_paddleocr_vl.py | 41 ++++--- .../models/phi/configuration_phi.py | 37 ++++--- src/transformers/models/phi/modeling_phi.py | 5 +- .../models/phi3/configuration_phi3.py | 33 +++--- src/transformers/models/phi3/modeling_phi3.py | 5 +- .../configuration_phi4_multimodal.py | 33 +++--- .../modeling_phi4_multimodal.py | 5 +- .../models/phimoe/modeling_phimoe.py | 5 +- .../models/qwen2/configuration_qwen2.py | 41 ++++--- .../models/qwen2/modeling_qwen2.py | 5 +- .../qwen2_5_vl/configuration_qwen2_5_vl.py | 15 ++- .../qwen2_moe/configuration_qwen2_moe.py | 15 ++- .../models/qwen2_vl/configuration_qwen2_vl.py | 15 ++- .../models/qwen3/configuration_qwen3.py | 45 ++++---- .../models/qwen3/modeling_qwen3.py | 5 +- .../models/qwen3_5/modeling_qwen3_5.py | 5 +- .../qwen3_5_moe/modeling_qwen3_5_moe.py | 5 +- .../qwen3_moe/configuration_qwen3_moe.py | 57 +++++----- .../models/qwen3_moe/modeling_qwen3_moe.py | 5 +- .../qwen3_next/configuration_qwen3_next.py | 27 ++--- .../models/qwen3_next/modeling_qwen3_next.py | 5 +- .../configuration_qwen3_omni_moe.py | 101 ++++++++---------- .../qwen3_omni_moe/modeling_qwen3_omni_moe.py | 7 +- .../configuration_qwen3_vl_moe.py | 37 +++---- .../models/seed_oss/configuration_seed_oss.py | 41 ++++--- .../models/seed_oss/modeling_seed_oss.py | 5 +- .../models/smollm3/modeling_smollm3.py | 5 +- .../models/solar_open/modeling_solar_open.py | 5 +- .../starcoder2/configuration_starcoder2.py | 37 ++++--- .../models/starcoder2/modeling_starcoder2.py | 5 +- .../models/vaultgemma/modeling_vaultgemma.py | 5 +- .../modeling_voxtral_realtime.py | 5 +- .../models/youtu/modeling_youtu.py | 5 +- 115 files changed, 798 insertions(+), 985 deletions(-) diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index 50d604e4a71a..065e7ff314cb 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -676,9 +676,7 @@ class ParallelInterface(GeneralInterface): "activation": SequenceParallel(), "activation_seq_dim_2": SequenceParallel(sequence_dim=2), # Module-level prepare-input - "module_allgather": PrepareModuleInput( - input_layouts=(Shard(1),), desired_input_layouts=(Replicate(),) - ), + "module_allgather": PrepareModuleInput(input_layouts=(Shard(1),), desired_input_layouts=(Replicate(),)), "module_allgather_hidden_states": PrepareModuleInput( input_kwarg_layouts={"hidden_states": Shard(1)}, desired_input_kwarg_layouts={"hidden_states": Replicate()}, diff --git a/src/transformers/models/apertus/modeling_apertus.py b/src/transformers/models/apertus/modeling_apertus.py index 96c588160a6d..de88a5e0023b 100644 --- a/src/transformers/models/apertus/modeling_apertus.py +++ b/src/transformers/models/apertus/modeling_apertus.py @@ -28,7 +28,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GenericForTokenClassification, GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -423,8 +422,8 @@ def forward( @auto_docstring class ApertusForCausalLM(ApertusPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/arcee/modeling_arcee.py b/src/transformers/models/arcee/modeling_arcee.py index 625b9154e336..a2681c4681a4 100644 --- a/src/transformers/models/arcee/modeling_arcee.py +++ b/src/transformers/models/arcee/modeling_arcee.py @@ -30,7 +30,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import ( GenericForQuestionAnswering, @@ -425,8 +424,8 @@ 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": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/aria/configuration_aria.py b/src/transformers/models/aria/configuration_aria.py index b8d9e834e37b..42380c73387a 100644 --- a/src/transformers/models/aria/configuration_aria.py +++ b/src/transformers/models/aria/configuration_aria.py @@ -20,7 +20,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring from ...utils.type_validators import interval @@ -51,19 +50,19 @@ class AriaTextConfig(PreTrainedConfig): "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), diff --git a/src/transformers/models/aria/modeling_aria.py b/src/transformers/models/aria/modeling_aria.py index d8112acaa870..05e4a40f29fb 100644 --- a/src/transformers/models/aria/modeling_aria.py +++ b/src/transformers/models/aria/modeling_aria.py @@ -29,7 +29,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -759,8 +758,8 @@ def forward( @auto_docstring class AriaTextForCausalLM(AriaTextPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: AriaTextConfig): diff --git a/src/transformers/models/bamba/modeling_bamba.py b/src/transformers/models/bamba/modeling_bamba.py index fe9b2e95c942..0fb69a490b80 100644 --- a/src/transformers/models/bamba/modeling_bamba.py +++ b/src/transformers/models/bamba/modeling_bamba.py @@ -35,7 +35,6 @@ from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernelized_func from ...integrations.hub_kernels import lazy_load_kernel -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -1072,8 +1071,8 @@ 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": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/cohere/configuration_cohere.py b/src/transformers/models/cohere/configuration_cohere.py index b2605904e84d..5b12131d72ba 100644 --- a/src/transformers/models/cohere/configuration_cohere.py +++ b/src/transformers/models/cohere/configuration_cohere.py @@ -21,7 +21,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -51,30 +50,30 @@ class CohereConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] default_theta = 500000.0 base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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/cohere/modeling_cohere.py b/src/transformers/models/cohere/modeling_cohere.py index 0315bf78ab3a..ed715930260d 100644 --- a/src/transformers/models/cohere/modeling_cohere.py +++ b/src/transformers/models/cohere/modeling_cohere.py @@ -36,7 +36,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -455,8 +454,8 @@ def forward( @auto_docstring class CohereForCausalLM(CoherePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/cohere2/modeling_cohere2.py b/src/transformers/models/cohere2/modeling_cohere2.py index 948958759e41..8b35fd4ab25a 100644 --- a/src/transformers/models/cohere2/modeling_cohere2.py +++ b/src/transformers/models/cohere2/modeling_cohere2.py @@ -28,7 +28,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -434,8 +433,8 @@ def forward( @auto_docstring class Cohere2ForCausalLM(Cohere2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/cwm/configuration_cwm.py b/src/transformers/models/cwm/configuration_cwm.py index cee303f7b618..1e086edaaa43 100644 --- a/src/transformers/models/cwm/configuration_cwm.py +++ b/src/transformers/models/cwm/configuration_cwm.py @@ -22,7 +22,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...utils import auto_docstring @@ -47,28 +46,28 @@ class CwmConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `CwmModel` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), diff --git a/src/transformers/models/cwm/modeling_cwm.py b/src/transformers/models/cwm/modeling_cwm.py index f3583a78f8ec..862f1cfcebe7 100644 --- a/src/transformers/models/cwm/modeling_cwm.py +++ b/src/transformers/models/cwm/modeling_cwm.py @@ -29,7 +29,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -427,8 +426,8 @@ def forward( @auto_docstring class CwmForCausalLM(CwmPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/deepseek_v2/modeling_deepseek_v2.py b/src/transformers/models/deepseek_v2/modeling_deepseek_v2.py index 0c2eacdb5d01..0eb276fb30f8 100644 --- a/src/transformers/models/deepseek_v2/modeling_deepseek_v2.py +++ b/src/transformers/models/deepseek_v2/modeling_deepseek_v2.py @@ -30,7 +30,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_experts_implementation, use_kernel_forward_from_hub -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GenericForSequenceClassification, GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -542,8 +541,8 @@ def forward( @auto_docstring class DeepseekV2ForCausalLM(DeepseekV2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py index 651a2fc2af5f..5dcad86422f6 100644 --- a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py @@ -18,7 +18,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -50,20 +49,13 @@ class DeepseekV3Config(PreTrainedConfig): model_type = "deepseek_v3" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={ - "gate_up_proj": TPStyle("packed_colwise", "none"), - "down_proj": TPStyle("rowwise", "allreduce"), - }, - ), - "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.experts": "moe_experts_allreduce", + "layers.*.mlp.shared_experts.gate_proj": "colwise", + "layers.*.mlp.shared_experts.up_proj": "colwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_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"]), diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index 8c7be7e1e019..c6f3d035051b 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -17,7 +17,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernel_func_from_hub -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -635,8 +634,8 @@ def forward( @auto_docstring class DeepseekV3ForCausalLM(DeepseekV3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/diffllama/modeling_diffllama.py b/src/transformers/models/diffllama/modeling_diffllama.py index 21e01341050d..1603ba8022a0 100644 --- a/src/transformers/models/diffllama/modeling_diffllama.py +++ b/src/transformers/models/diffllama/modeling_diffllama.py @@ -32,7 +32,6 @@ from ...cache_utils import Cache, DynamicCache, StaticCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import _flash_attention_forward, flash_attn_supports_top_left_mask from ...modeling_layers import ( @@ -661,8 +660,8 @@ def forward( @auto_docstring class DiffLlamaForCausalLM(DiffLlamaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/doge/modeling_doge.py b/src/transformers/models/doge/modeling_doge.py index 531a2758548e..7d2f91f444cf 100644 --- a/src/transformers/models/doge/modeling_doge.py +++ b/src/transformers/models/doge/modeling_doge.py @@ -34,7 +34,6 @@ from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub from ...integrations.flex_attention import compile_friendly_flex_attention -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_layers import GenericForSequenceClassification, GradientCheckpointingLayer from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast @@ -717,8 +716,8 @@ 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": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/dots1/modeling_dots1.py b/src/transformers/models/dots1/modeling_dots1.py index b937fe12c014..c6caf4e6fd91 100644 --- a/src/transformers/models/dots1/modeling_dots1.py +++ b/src/transformers/models/dots1/modeling_dots1.py @@ -29,7 +29,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -546,8 +545,8 @@ def forward( @auto_docstring class Dots1ForCausalLM(Dots1PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/emu3/modeling_emu3.py b/src/transformers/models/emu3/modeling_emu3.py index b77e735f092f..0691845d6bba 100644 --- a/src/transformers/models/emu3/modeling_emu3.py +++ b/src/transformers/models/emu3/modeling_emu3.py @@ -34,7 +34,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, CausalLMOutputWithPast @@ -1276,8 +1275,8 @@ def forward( @auto_docstring class Emu3ForCausalLM(Emu3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} config: Emu3TextConfig diff --git a/src/transformers/models/ernie4_5/configuration_ernie4_5.py b/src/transformers/models/ernie4_5/configuration_ernie4_5.py index facfc55d07f0..ac2d6fead084 100644 --- a/src/transformers/models/ernie4_5/configuration_ernie4_5.py +++ b/src/transformers/models/ernie4_5/configuration_ernie4_5.py @@ -16,7 +16,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -48,28 +47,28 @@ class Ernie4_5Config(PreTrainedConfig): default_theta = 500000.0 # Default tensor parallel plan for base model `Ernie4_5Model` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), diff --git a/src/transformers/models/ernie4_5/modeling_ernie4_5.py b/src/transformers/models/ernie4_5/modeling_ernie4_5.py index 86512c389707..367ce8c88323 100644 --- a/src/transformers/models/ernie4_5/modeling_ernie4_5.py +++ b/src/transformers/models/ernie4_5/modeling_ernie4_5.py @@ -28,7 +28,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -423,8 +422,8 @@ def forward( @auto_docstring class Ernie4_5ForCausalLM(Ernie4_5PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, 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 22b38d1d5d33..5967808f5607 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 @@ -16,7 +16,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -64,21 +63,17 @@ class Ernie4_5_MoeConfig(PreTrainedConfig): # Default tensor parallel plan for base model `Ernie4_5_MoE` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), - "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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": "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"]), 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 e7b3f38181c1..a594cd4e306e 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 @@ -30,7 +30,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast @@ -655,8 +654,8 @@ 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": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/exaone4/modeling_exaone4.py b/src/transformers/models/exaone4/modeling_exaone4.py index 888e3c091a8b..570b5e6fe160 100644 --- a/src/transformers/models/exaone4/modeling_exaone4.py +++ b/src/transformers/models/exaone4/modeling_exaone4.py @@ -29,7 +29,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_layers import ( GenericForQuestionAnswering, @@ -441,8 +440,8 @@ def forward( @auto_docstring class Exaone4ForCausalLM(Exaone4PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/exaone_moe/modeling_exaone_moe.py b/src/transformers/models/exaone_moe/modeling_exaone_moe.py index d5d0ff38fa60..8c3a4fd53ab9 100644 --- a/src/transformers/models/exaone_moe/modeling_exaone_moe.py +++ b/src/transformers/models/exaone_moe/modeling_exaone_moe.py @@ -31,7 +31,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernel_func_from_hub -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -564,8 +563,8 @@ def forward( @auto_docstring class ExaoneMoeForCausalLM(ExaoneMoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/falcon_h1/modeling_falcon_h1.py b/src/transformers/models/falcon_h1/modeling_falcon_h1.py index 41520e97518f..60c18ac0264c 100644 --- a/src/transformers/models/falcon_h1/modeling_falcon_h1.py +++ b/src/transformers/models/falcon_h1/modeling_falcon_h1.py @@ -36,7 +36,6 @@ from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func from ...integrations.hub_kernels import lazy_load_kernel -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -1167,8 +1166,8 @@ 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": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/flex_olmo/modeling_flex_olmo.py b/src/transformers/models/flex_olmo/modeling_flex_olmo.py index fa6b44143646..fa1811906643 100644 --- a/src/transformers/models/flex_olmo/modeling_flex_olmo.py +++ b/src/transformers/models/flex_olmo/modeling_flex_olmo.py @@ -31,7 +31,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast @@ -598,8 +597,8 @@ 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": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/gemma/modeling_gemma.py b/src/transformers/models/gemma/modeling_gemma.py index e2b6b7667667..79a61722c36a 100644 --- a/src/transformers/models/gemma/modeling_gemma.py +++ b/src/transformers/models/gemma/modeling_gemma.py @@ -31,7 +31,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_func_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import ( GenericForSequenceClassification, @@ -451,8 +450,8 @@ def forward( @auto_docstring class GemmaForCausalLM(GemmaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/gemma2/modeling_gemma2.py b/src/transformers/models/gemma2/modeling_gemma2.py index 1685a7f8d3bf..7dc49c2f9e28 100644 --- a/src/transformers/models/gemma2/modeling_gemma2.py +++ b/src/transformers/models/gemma2/modeling_gemma2.py @@ -29,7 +29,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_func_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -477,8 +476,8 @@ def forward( @auto_docstring class Gemma2ForCausalLM(Gemma2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/gemma3/modeling_gemma3.py b/src/transformers/models/gemma3/modeling_gemma3.py index d3f8e5029f7e..2c77e7dcc0a4 100644 --- a/src/transformers/models/gemma3/modeling_gemma3.py +++ b/src/transformers/models/gemma3/modeling_gemma3.py @@ -31,7 +31,6 @@ from ...configuration_utils import PreTrainedConfig from ...generation import GenerationMixin from ...integrations import use_kernel_func_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_masks_for_generate, create_sliding_window_causal_mask from ...modeling_layers import GenericForSequenceClassification, GradientCheckpointingLayer from ...modeling_outputs import ( @@ -592,8 +591,8 @@ def forward( @auto_docstring class Gemma3ForCausalLM(Gemma3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} config: Gemma3TextConfig diff --git a/src/transformers/models/gemma3n/modeling_gemma3n.py b/src/transformers/models/gemma3n/modeling_gemma3n.py index ca38d342fedb..343fbeb35e6e 100644 --- a/src/transformers/models/gemma3n/modeling_gemma3n.py +++ b/src/transformers/models/gemma3n/modeling_gemma3n.py @@ -32,7 +32,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, CausalLMOutputWithPast @@ -1770,8 +1769,8 @@ 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": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} config: Gemma3nTextConfig diff --git a/src/transformers/models/gemma4/configuration_gemma4.py b/src/transformers/models/gemma4/configuration_gemma4.py index 4bb3b142eb77..cc2dab52f747 100644 --- a/src/transformers/models/gemma4/configuration_gemma4.py +++ b/src/transformers/models/gemma4/configuration_gemma4.py @@ -17,7 +17,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...utils import auto_docstring, logging from ...utils.type_validators import interval @@ -126,13 +125,13 @@ class Gemma4TextConfig(PreTrainedConfig): base_model_tp_plan = { # 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": TPStyle("colwise", "allgather"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "allgather"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "allgather"), - "layers.*.self_attn.o_proj": TPStyle("vocab", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -224,13 +223,13 @@ class Gemma4VisionConfig(PreTrainedConfig): model_type = "gemma4_vision" base_model_tp_plan = { - "encoder.layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "encoder.layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "encoder.layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "encoder.layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "encoder.layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "encoder.layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "encoder.layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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.o_proj": "rowwise_allreduce", + "encoder.layers.*.mlp.gate_proj": "colwise", + "encoder.layers.*.mlp.up_proj": "colwise", + "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 91cfe090afeb..fc14687ffba3 100644 --- a/src/transformers/models/gemma4/modeling_gemma4.py +++ b/src/transformers/models/gemma4/modeling_gemma4.py @@ -34,7 +34,6 @@ from ...configuration_utils import PreTrainedConfig from ...generation import GenerationMixin from ...integrations import use_experts_implementation, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import ( create_bidirectional_mask, create_causal_mask, @@ -1700,8 +1699,8 @@ 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": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} config: Gemma4TextConfig base_model_prefix = "model" diff --git a/src/transformers/models/glm/configuration_glm.py b/src/transformers/models/glm/configuration_glm.py index 222d1a8dadfa..52197a8d596f 100644 --- a/src/transformers/models/glm/configuration_glm.py +++ b/src/transformers/models/glm/configuration_glm.py @@ -17,7 +17,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -41,26 +40,26 @@ class GlmConfig(PreTrainedConfig): model_type = "glm" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_up_proj": TPStyle("packed_colwise", "none"), # fused gate/up shards stay local for chunk - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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_up_proj": "packed_colwise", # fused gate/up shards stay local for chunk + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_up_proj": TPStyle("packed_colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), diff --git a/src/transformers/models/glm/modeling_glm.py b/src/transformers/models/glm/modeling_glm.py index 084f8b55a12b..e87e10841bab 100644 --- a/src/transformers/models/glm/modeling_glm.py +++ b/src/transformers/models/glm/modeling_glm.py @@ -28,7 +28,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import ( GenericForSequenceClassification, @@ -440,8 +439,8 @@ def forward( @auto_docstring class GlmForCausalLM(GlmPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/glm4/configuration_glm4.py b/src/transformers/models/glm4/configuration_glm4.py index e6f51c996b66..d4ab4053f93d 100644 --- a/src/transformers/models/glm4/configuration_glm4.py +++ b/src/transformers/models/glm4/configuration_glm4.py @@ -17,7 +17,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -41,28 +40,26 @@ class Glm4Config(PreTrainedConfig): model_type = "glm4" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_up_proj": TPStyle( - "colwise", "allgather" - ), # we need to replicate here due to the `chunk` operation - "layers.*.mlp.down_proj": TPStyle("vocab", "allreduce"), # input is replicated due to the `chunk` operation + "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_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": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_up_proj": TPStyle("colwise", "allgather"), # fused gate/up needs full tensor for chunk - "layers.*.mlp.down_proj": TPStyle("vocab", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), diff --git a/src/transformers/models/glm4/modeling_glm4.py b/src/transformers/models/glm4/modeling_glm4.py index 0edde697243d..5c4eb666da85 100644 --- a/src/transformers/models/glm4/modeling_glm4.py +++ b/src/transformers/models/glm4/modeling_glm4.py @@ -29,7 +29,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -445,8 +444,8 @@ def forward( @auto_docstring class Glm4ForCausalLM(Glm4PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/glm4_moe/modeling_glm4_moe.py b/src/transformers/models/glm4_moe/modeling_glm4_moe.py index bed0c0153c59..6a3ceb451cc1 100644 --- a/src/transformers/models/glm4_moe/modeling_glm4_moe.py +++ b/src/transformers/models/glm4_moe/modeling_glm4_moe.py @@ -30,7 +30,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -578,8 +577,8 @@ def forward( @auto_docstring class Glm4MoeForCausalLM(Glm4MoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): 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 1175490d801b..c00450336647 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 @@ -32,7 +32,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernel_func_from_hub -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -652,8 +651,8 @@ def forward( @auto_docstring class Glm4MoeLiteForCausalLM(Glm4MoeLitePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): 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 037a1b4a6b41..96b9c512e00c 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 @@ -30,7 +30,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_experts_implementation, use_kernel_forward_from_hub -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -806,8 +805,8 @@ def forward( @auto_docstring class GlmMoeDsaForCausalLM(GlmMoeDsaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/gpt_neox/configuration_gpt_neox.py b/src/transformers/models/gpt_neox/configuration_gpt_neox.py index 83df236f3741..5b4b7908e213 100644 --- a/src/transformers/models/gpt_neox/configuration_gpt_neox.py +++ b/src/transformers/models/gpt_neox/configuration_gpt_neox.py @@ -16,7 +16,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -47,19 +46,19 @@ class GPTNeoXConfig(PreTrainedConfig): model_type = "gpt_neox" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.attention.query_key_value": TPStyle("colwise", "none"), - "layers.*.attention.dense": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.dense_h_to_4h": TPStyle("colwise", "none"), - "layers.*.mlp.dense_4h_to_h": TPStyle("rowwise", "allreduce"), + "layers.*.attention.query_key_value": "colwise", + "layers.*.attention.dense": "rowwise_allreduce", + "layers.*.mlp.dense_h_to_4h": "colwise", + "layers.*.mlp.dense_4h_to_h": "rowwise_allreduce", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.dense_h_to_4h": TPStyle("colwise", "none"), - "layers.*.mlp.dense_4h_to_h": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), diff --git a/src/transformers/models/gpt_oss/modeling_gpt_oss.py b/src/transformers/models/gpt_oss/modeling_gpt_oss.py index 4f39413d4560..629d8f56b632 100644 --- a/src/transformers/models/gpt_oss/modeling_gpt_oss.py +++ b/src/transformers/models/gpt_oss/modeling_gpt_oss.py @@ -28,7 +28,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_layers import ( GenericForSequenceClassification, @@ -589,8 +588,8 @@ 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": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/granite/configuration_granite.py b/src/transformers/models/granite/configuration_granite.py index fbaaf32ab37f..063452fbe08d 100644 --- a/src/transformers/models/granite/configuration_granite.py +++ b/src/transformers/models/granite/configuration_granite.py @@ -21,7 +21,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -48,28 +47,28 @@ class GraniteConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `GraniteModel` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), diff --git a/src/transformers/models/granite/modeling_granite.py b/src/transformers/models/granite/modeling_granite.py index 1621b5c86922..68c8877446d4 100644 --- a/src/transformers/models/granite/modeling_granite.py +++ b/src/transformers/models/granite/modeling_granite.py @@ -29,7 +29,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -446,8 +445,8 @@ def forward( @auto_docstring class GraniteForCausalLM(GranitePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/granitemoe/modeling_granitemoe.py b/src/transformers/models/granitemoe/modeling_granitemoe.py index 7a8ab944bdc5..6731c00e178a 100644 --- a/src/transformers/models/granitemoe/modeling_granitemoe.py +++ b/src/transformers/models/granitemoe/modeling_granitemoe.py @@ -31,7 +31,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast @@ -627,8 +626,8 @@ 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": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: GraniteMoeConfig): diff --git a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py index df81d7eea1cd..01a8ba6b3243 100644 --- a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py +++ b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py @@ -31,7 +31,6 @@ from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func from ...integrations.hub_kernels import lazy_load_kernel -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, MoeCausalLMOutputWithPast, MoeModelOutputWithPast @@ -1307,8 +1306,8 @@ 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": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: GraniteMoeHybridConfig): diff --git a/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py b/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py index 33e9d0bd144f..0c0a4e92ed7f 100644 --- a/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py +++ b/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py @@ -30,7 +30,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast @@ -696,8 +695,8 @@ 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": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: GraniteMoeSharedConfig): diff --git a/src/transformers/models/helium/configuration_helium.py b/src/transformers/models/helium/configuration_helium.py index 8cec52902ceb..15cd1b77db0f 100644 --- a/src/transformers/models/helium/configuration_helium.py +++ b/src/transformers/models/helium/configuration_helium.py @@ -17,7 +17,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -42,28 +41,28 @@ class HeliumConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] default_theta = 100000.0 base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), diff --git a/src/transformers/models/helium/modeling_helium.py b/src/transformers/models/helium/modeling_helium.py index c008cc067365..4463cc02980a 100644 --- a/src/transformers/models/helium/modeling_helium.py +++ b/src/transformers/models/helium/modeling_helium.py @@ -29,7 +29,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import ( GenericForSequenceClassification, @@ -424,8 +423,8 @@ def forward( @auto_docstring class HeliumForCausalLM(HeliumPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, 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 bc33ecd6e8c7..38e57e716c9a 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 @@ -22,7 +22,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring from ...utils.type_validators import interval @@ -60,28 +59,28 @@ class HiggsAudioV2Config(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `HiggsAudioV2Model` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), 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 09727a8304b1..79b21c72f890 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 @@ -31,7 +31,6 @@ from ...cache_utils import DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GenericForSequenceClassification, GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -462,8 +461,8 @@ def forward( @auto_docstring class HunYuanDenseV1ForCausalLM(HunYuanDenseV1PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, 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 73392ca1fdeb..9bb57e90bca1 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 @@ -35,7 +35,6 @@ use_kernel_func_from_hub, use_kernelized_func, ) -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GenericForSequenceClassification, GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -551,8 +550,8 @@ def forward( @auto_docstring class HunYuanMoEV1ForCausalLM(HunYuanMoEV1PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/jais2/modeling_jais2.py b/src/transformers/models/jais2/modeling_jais2.py index 93316dca79aa..0aec15924b03 100644 --- a/src/transformers/models/jais2/modeling_jais2.py +++ b/src/transformers/models/jais2/modeling_jais2.py @@ -29,7 +29,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_func_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -398,8 +397,8 @@ def forward( @auto_docstring class Jais2ForCausalLM(Jais2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/jamba/modeling_jamba.py b/src/transformers/models/jamba/modeling_jamba.py index 439fdf2f8cfb..b93ef6197f4b 100755 --- a/src/transformers/models/jamba/modeling_jamba.py +++ b/src/transformers/models/jamba/modeling_jamba.py @@ -38,7 +38,6 @@ use_kernel_func_from_hub, use_kernelized_func, ) -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GenericForSequenceClassification, GradientCheckpointingLayer from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast @@ -844,8 +843,8 @@ 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": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: JambaConfig): 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 6e408304fa28..80f5f756b98b 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 @@ -30,7 +30,6 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache, StaticCache from ...generation import GenerationConfig, GenerationMixin -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import flash_attn_supports_top_left_mask, is_flash_attn_available from ...modeling_layers import GradientCheckpointingLayer @@ -875,8 +874,8 @@ def forward( @auto_docstring class KyutaiSpeechToTextForConditionalGeneration(KyutaiSpeechToTextPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _keep_in_fp32_modules_strict = ["codec_model"] output_modalities = ("audio", "text") diff --git a/src/transformers/models/lfm2/modeling_lfm2.py b/src/transformers/models/lfm2/modeling_lfm2.py index 6279867bf54d..c8d42b555903 100644 --- a/src/transformers/models/lfm2/modeling_lfm2.py +++ b/src/transformers/models/lfm2/modeling_lfm2.py @@ -27,7 +27,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -541,8 +540,8 @@ def forward( @auto_docstring class Lfm2ForCausalLM(Lfm2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py b/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py index 122816e4c1e3..ed6df4a3e34b 100644 --- a/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py +++ b/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py @@ -34,7 +34,6 @@ use_kernel_func_from_hub, use_kernelized_func, ) -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, MoeModelOutputWithPast @@ -631,8 +630,8 @@ def forward( @auto_docstring class Lfm2MoeForCausalLM(Lfm2MoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/llama/configuration_llama.py b/src/transformers/models/llama/configuration_llama.py index 1f1d378f8e71..73b898d619bd 100644 --- a/src/transformers/models/llama/configuration_llama.py +++ b/src/transformers/models/llama/configuration_llama.py @@ -21,7 +21,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring from ...utils.type_validators import interval @@ -48,28 +47,28 @@ class LlamaConfig(PreTrainedConfig): 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": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), diff --git a/src/transformers/models/llama/modeling_llama.py b/src/transformers/models/llama/modeling_llama.py index 270ef1b0d228..c12f5e1966b1 100644 --- a/src/transformers/models/llama/modeling_llama.py +++ b/src/transformers/models/llama/modeling_llama.py @@ -26,7 +26,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import ( GenericForQuestionAnswering, @@ -429,8 +428,8 @@ def forward( @auto_docstring class LlamaForCausalLM(LlamaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/llama4/configuration_llama4.py b/src/transformers/models/llama4/configuration_llama4.py index 1d8c350d4b20..bf809aeeb6c7 100644 --- a/src/transformers/models/llama4/configuration_llama4.py +++ b/src/transformers/models/llama4/configuration_llama4.py @@ -17,7 +17,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring, logging @@ -43,13 +42,13 @@ class Llama4VisionConfig(PreTrainedConfig): """ base_model_tp_plan = { - "model.layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "model.layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "model.layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "model.layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "vision_adapter.mlp.fc1": TPStyle("colwise", "none"), - "vision_adapter.mlp.fc2": TPStyle("rowwise", "allreduce"), - "patch_embedding.linear": TPStyle("colwise", "allgather"), + "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_allreduce", + "vision_adapter.mlp.fc1": "colwise", + "vision_adapter.mlp.fc2": "rowwise_allreduce", + "patch_embedding.linear": "colwise_allgather", } model_type = "llama4_vision_model" base_config_key = "vision_config" @@ -111,16 +110,16 @@ class Llama4TextConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] default_theta = 500000.0 base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.feed_forward.shared_expert.gate_proj": TPStyle("colwise", "none"), - "layers.*.feed_forward.shared_expert.up_proj": TPStyle("colwise", "none"), - "layers.*.feed_forward.shared_expert.down_proj": TPStyle("rowwise", "allreduce"), - "layers.*.feed_forward.gate_proj": TPStyle("colwise", "none"), - "layers.*.feed_forward.up_proj": TPStyle("colwise", "none"), - "layers.*.feed_forward.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_ep_plan = { "layers.*.self_attn.q_proj": "colwise", diff --git a/src/transformers/models/llama4/modeling_llama4.py b/src/transformers/models/llama4/modeling_llama4.py index 6e10d0085057..a3e5f4dad8dd 100644 --- a/src/transformers/models/llama4/modeling_llama4.py +++ b/src/transformers/models/llama4/modeling_llama4.py @@ -28,7 +28,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_chunked_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -591,8 +590,8 @@ 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": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} config: Llama4TextConfig def __init__(self, config: Llama4TextConfig): diff --git a/src/transformers/models/longcat_flash/configuration_longcat_flash.py b/src/transformers/models/longcat_flash/configuration_longcat_flash.py index 3ad063190541..fcdcbb8cae28 100644 --- a/src/transformers/models/longcat_flash/configuration_longcat_flash.py +++ b/src/transformers/models/longcat_flash/configuration_longcat_flash.py @@ -17,7 +17,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -55,12 +54,12 @@ class LongcatFlashConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] default_theta = 10000000.0 base_model_tp_plan = { - "layers.*.self_attn.*.q_b_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.*.kv_b_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.*.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlps.*.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlps.*.up_proj": TPStyle("colwise", "none"), - "layers.*.mlps.*.down_proj": TPStyle("rowwise", "allreduce"), + "layers.*.self_attn.*.q_b_proj": "colwise", + "layers.*.self_attn.*.kv_b_proj": "colwise", + "layers.*.self_attn.*.o_proj": "rowwise_allreduce", + "layers.*.mlps.*.gate_proj": "colwise", + "layers.*.mlps.*.up_proj": "colwise", + "layers.*.mlps.*.down_proj": "rowwise_allreduce", } base_model_pp_plan = { diff --git a/src/transformers/models/longcat_flash/modeling_longcat_flash.py b/src/transformers/models/longcat_flash/modeling_longcat_flash.py index 8fb78077d9cc..443a860dcaae 100644 --- a/src/transformers/models/longcat_flash/modeling_longcat_flash.py +++ b/src/transformers/models/longcat_flash/modeling_longcat_flash.py @@ -31,7 +31,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -651,8 +650,8 @@ def forward( @auto_docstring class LongcatFlashForCausalLM(LongcatFlashPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _keys_to_ignore_on_load_unexpected = [r"model\.mtp.*"] diff --git a/src/transformers/models/minimax/modeling_minimax.py b/src/transformers/models/minimax/modeling_minimax.py index 53a0cd3e6238..7540f1ce5329 100644 --- a/src/transformers/models/minimax/modeling_minimax.py +++ b/src/transformers/models/minimax/modeling_minimax.py @@ -36,7 +36,6 @@ use_kernel_func_from_hub, use_kernelized_func, ) -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -790,8 +789,8 @@ 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": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/minimax_m2/modeling_minimax_m2.py b/src/transformers/models/minimax_m2/modeling_minimax_m2.py index 360f50b80bab..82389018bbbf 100644 --- a/src/transformers/models/minimax_m2/modeling_minimax_m2.py +++ b/src/transformers/models/minimax_m2/modeling_minimax_m2.py @@ -31,7 +31,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast @@ -589,8 +588,8 @@ 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": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/ministral/configuration_ministral.py b/src/transformers/models/ministral/configuration_ministral.py index f6002d51ef1a..84d69e92d7f3 100644 --- a/src/transformers/models/ministral/configuration_ministral.py +++ b/src/transformers/models/ministral/configuration_ministral.py @@ -22,7 +22,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -50,28 +49,28 @@ class MinistralConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `MinistralModel` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), diff --git a/src/transformers/models/ministral/modeling_ministral.py b/src/transformers/models/ministral/modeling_ministral.py index 3f76f68bdf3a..510364445c3c 100644 --- a/src/transformers/models/ministral/modeling_ministral.py +++ b/src/transformers/models/ministral/modeling_ministral.py @@ -29,7 +29,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -430,8 +429,8 @@ def forward( @auto_docstring class MinistralForCausalLM(MinistralPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/ministral3/configuration_ministral3.py b/src/transformers/models/ministral3/configuration_ministral3.py index 83dd9d21e75f..21a9a6515fba 100644 --- a/src/transformers/models/ministral3/configuration_ministral3.py +++ b/src/transformers/models/ministral3/configuration_ministral3.py @@ -16,7 +16,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring, logging @@ -56,28 +55,28 @@ class Ministral3Config(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `MistralModel` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), diff --git a/src/transformers/models/ministral3/modeling_ministral3.py b/src/transformers/models/ministral3/modeling_ministral3.py index c9c321d2b82a..bfc4d35df5f9 100644 --- a/src/transformers/models/ministral3/modeling_ministral3.py +++ b/src/transformers/models/ministral3/modeling_ministral3.py @@ -14,7 +14,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -413,8 +412,8 @@ def forward( @auto_docstring class Ministral3ForCausalLM(Ministral3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/mistral/configuration_mistral.py b/src/transformers/models/mistral/configuration_mistral.py index 913e6e3c96dc..27d125446329 100644 --- a/src/transformers/models/mistral/configuration_mistral.py +++ b/src/transformers/models/mistral/configuration_mistral.py @@ -16,7 +16,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring, logging @@ -47,28 +46,28 @@ class MistralConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `MistralModel` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), diff --git a/src/transformers/models/mistral/modeling_mistral.py b/src/transformers/models/mistral/modeling_mistral.py index 6faeb29c8a8e..4e4b24132f66 100644 --- a/src/transformers/models/mistral/modeling_mistral.py +++ b/src/transformers/models/mistral/modeling_mistral.py @@ -14,7 +14,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -402,8 +401,8 @@ def forward( @auto_docstring class MistralForCausalLM(MistralPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/mistral4/configuration_mistral4.py b/src/transformers/models/mistral4/configuration_mistral4.py index 74fed8abab2a..774b61016f38 100644 --- a/src/transformers/models/mistral4/configuration_mistral4.py +++ b/src/transformers/models/mistral4/configuration_mistral4.py @@ -16,7 +16,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -48,17 +47,13 @@ class Mistral4Config(PreTrainedConfig): model_type = "mistral4" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), - "layers.*.mlp.shared_experts.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_experts.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_experts.down_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "layers.*.mlp.experts": "moe_experts_allreduce", + "layers.*.mlp.shared_experts.gate_proj": "colwise", + "layers.*.mlp.shared_experts.up_proj": "colwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_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"]), diff --git a/src/transformers/models/mistral4/modeling_mistral4.py b/src/transformers/models/mistral4/modeling_mistral4.py index dbc9adc77fa7..a31527ce4b4d 100644 --- a/src/transformers/models/mistral4/modeling_mistral4.py +++ b/src/transformers/models/mistral4/modeling_mistral4.py @@ -29,7 +29,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernel_func_from_hub -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -641,8 +640,8 @@ def forward( @auto_docstring class Mistral4ForCausalLM(Mistral4PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/mixtral/configuration_mixtral.py b/src/transformers/models/mixtral/configuration_mixtral.py index e5e5294a0116..242664d356b5 100644 --- a/src/transformers/models/mixtral/configuration_mixtral.py +++ b/src/transformers/models/mixtral/configuration_mixtral.py @@ -16,7 +16,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -45,40 +44,26 @@ class MixtralConfig(PreTrainedConfig): default_theta = 1000000.0 # TP plan (for inference/generation). base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={ - "gate_up_proj": TPStyle("packed_colwise", "none"), - "down_proj": TPStyle("rowwise", "allreduce"), - }, - ), + "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": "moe_experts_allreduce", } # TP + Sequence Parallelism plan (for training). base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather_split"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={ - "gate_up_proj": TPStyle("packed_colwise", "none"), - "down_proj": TPStyle("rowwise", "allreduce"), - }, - ), - "norm": TPStyle("activation", "none"), + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts": "moe_experts_allreduce", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/mixtral/modeling_mixtral.py b/src/transformers/models/mixtral/modeling_mixtral.py index c755959639a5..f76680ace4c4 100644 --- a/src/transformers/models/mixtral/modeling_mixtral.py +++ b/src/transformers/models/mixtral/modeling_mixtral.py @@ -40,7 +40,6 @@ use_kernel_func_from_hub, use_kernelized_func, ) -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -581,8 +580,8 @@ 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": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/nanochat/configuration_nanochat.py b/src/transformers/models/nanochat/configuration_nanochat.py index 2f419697e952..d8af91fccce7 100644 --- a/src/transformers/models/nanochat/configuration_nanochat.py +++ b/src/transformers/models/nanochat/configuration_nanochat.py @@ -16,7 +16,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PretrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -44,12 +43,12 @@ class NanoChatConfig(PretrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.fc1": TPStyle("colwise", "none"), - "layers.*.mlp.fc2": TPStyle("rowwise", "allreduce"), + "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.fc1": "colwise", + "layers.*.mlp.fc2": "rowwise_allreduce", } base_model_sp_plan = None diff --git a/src/transformers/models/olmo/configuration_olmo.py b/src/transformers/models/olmo/configuration_olmo.py index f67cc37f8bd5..7f8fcfa2c168 100644 --- a/src/transformers/models/olmo/configuration_olmo.py +++ b/src/transformers/models/olmo/configuration_olmo.py @@ -21,7 +21,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -51,28 +50,28 @@ class OlmoConfig(PreTrainedConfig): model_type = "olmo" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), diff --git a/src/transformers/models/olmo/modeling_olmo.py b/src/transformers/models/olmo/modeling_olmo.py index cab9bb3b3e1a..b3d335b351d8 100644 --- a/src/transformers/models/olmo/modeling_olmo.py +++ b/src/transformers/models/olmo/modeling_olmo.py @@ -34,7 +34,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -426,8 +425,8 @@ def forward( @auto_docstring class OlmoForCausalLM(OlmoPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/olmo2/modeling_olmo2.py b/src/transformers/models/olmo2/modeling_olmo2.py index c319c4a3478a..2811ac1a6c1d 100644 --- a/src/transformers/models/olmo2/modeling_olmo2.py +++ b/src/transformers/models/olmo2/modeling_olmo2.py @@ -35,7 +35,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -430,8 +429,8 @@ def forward( @auto_docstring class Olmo2ForCausalLM(Olmo2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/olmo3/modeling_olmo3.py b/src/transformers/models/olmo3/modeling_olmo3.py index daf8db42d9d1..e8e8d93ba729 100644 --- a/src/transformers/models/olmo3/modeling_olmo3.py +++ b/src/transformers/models/olmo3/modeling_olmo3.py @@ -28,7 +28,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -434,8 +433,8 @@ def forward( @auto_docstring class Olmo3ForCausalLM(Olmo3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py index 788c73eca593..9cc922a0c6f3 100644 --- a/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py +++ b/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py @@ -32,7 +32,6 @@ from ...cache_utils import Cache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -1033,8 +1032,8 @@ 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": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/olmoe/configuration_olmoe.py b/src/transformers/models/olmoe/configuration_olmoe.py index 4bef0fdec27b..33258ec2c05f 100644 --- a/src/transformers/models/olmoe/configuration_olmoe.py +++ b/src/transformers/models/olmoe/configuration_olmoe.py @@ -14,7 +14,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -47,32 +46,24 @@ class OlmoeConfig(PreTrainedConfig): # Default tensor parallel plan for base model `Olmoe` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "allgather"), # due to the norm, we have to gather - "layers.*.self_attn.k_proj": TPStyle("colwise", "allgather"), # due to the norm, we have to gather - "layers.*.self_attn.v_proj": TPStyle("colwise", "allgather"), # due to the norm, we have to gather - "layers.*.self_attn.o_proj": TPStyle("vocab", "allreduce"), # due to the norm, we have to gather - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), + "layers.*.self_attn.q_proj": "colwise_allgather", # due to the norm, we have to gather + "layers.*.self_attn.k_proj": "colwise_allgather", # due to the norm, we have to gather + "layers.*.self_attn.v_proj": "colwise_allgather", # due to the norm, we have to gather + "layers.*.self_attn.o_proj": "vocab_allreduce", # due to the norm, we have to gather + "layers.*.mlp.experts": "moe_experts_allreduce", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "allgather"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "allgather"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "allgather"), - "layers.*.self_attn.o_proj": TPStyle("vocab", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather_split"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), - "norm": TPStyle("activation", "none"), + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise_allgather", + "layers.*.self_attn.k_proj": "colwise_allgather", + "layers.*.self_attn.v_proj": "colwise_allgather", + "layers.*.self_attn.o_proj": "vocab_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts": "moe_experts_allreduce", + "norm": "activation", } vocab_size: int = 50304 diff --git a/src/transformers/models/olmoe/modeling_olmoe.py b/src/transformers/models/olmoe/modeling_olmoe.py index 2b9ce11116a2..f69926761833 100644 --- a/src/transformers/models/olmoe/modeling_olmoe.py +++ b/src/transformers/models/olmoe/modeling_olmoe.py @@ -33,7 +33,6 @@ use_kernel_func_from_hub, use_kernelized_func, ) -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast @@ -605,8 +604,8 @@ 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": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py b/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py index cc2257fabef2..83991ac90f1c 100644 --- a/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py +++ b/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py @@ -28,7 +28,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -96,28 +95,28 @@ class PaddleOCRTextConfig(PreTrainedConfig): default_theta = 500000.0 # Default tensor parallel plan for base model `PaddleOCRTextModel` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), diff --git a/src/transformers/models/phi/configuration_phi.py b/src/transformers/models/phi/configuration_phi.py index 99432092ed4f..d10e30019bfd 100644 --- a/src/transformers/models/phi/configuration_phi.py +++ b/src/transformers/models/phi/configuration_phi.py @@ -17,7 +17,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -47,26 +46,26 @@ class PhiConfig(PreTrainedConfig): model_type = "phi" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.dense": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.fc1": TPStyle("colwise", "none"), - "layers.*.mlp.fc2": TPStyle("rowwise", "allreduce"), + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.dense": "rowwise_allreduce", + "layers.*.mlp.fc1": "colwise", + "layers.*.mlp.fc2": "rowwise_allreduce", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.dense": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.fc1": TPStyle("colwise", "none"), - "layers.*.mlp.fc2": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), diff --git a/src/transformers/models/phi/modeling_phi.py b/src/transformers/models/phi/modeling_phi.py index ed33a66731ed..70a0bb6d6bba 100644 --- a/src/transformers/models/phi/modeling_phi.py +++ b/src/transformers/models/phi/modeling_phi.py @@ -14,7 +14,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_func_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import ( GenericForSequenceClassification, @@ -407,8 +406,8 @@ def forward( @auto_docstring class PhiForCausalLM(PhiPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/phi3/configuration_phi3.py b/src/transformers/models/phi3/configuration_phi3.py index 38691f62356a..e2d4d8934b05 100644 --- a/src/transformers/models/phi3/configuration_phi3.py +++ b/src/transformers/models/phi3/configuration_phi3.py @@ -17,7 +17,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -48,26 +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": TPStyle( - "colwise", "allgather" - ), # we need to replicate here due to the slicing of qkv - "layers.*.self_attn.o_proj": TPStyle("vocab", "allreduce"), # input is replicated due to the slicing of qkv - "layers.*.mlp.gate_up_proj": TPStyle( - "colwise", "allgather" - ), # we need to replicate here due to the `chunk` operation - "layers.*.mlp.down_proj": TPStyle("vocab", "allreduce"), # 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": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.qkv_proj": TPStyle("colwise", "allgather"), # fused qkv needs full tensor for slicing - "layers.*.self_attn.o_proj": TPStyle("vocab", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_up_proj": TPStyle("colwise", "allgather"), # fused gate/up needs full tensor for chunk - "layers.*.mlp.down_proj": TPStyle("vocab", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), diff --git a/src/transformers/models/phi3/modeling_phi3.py b/src/transformers/models/phi3/modeling_phi3.py index a549f703b647..41d6634095b7 100644 --- a/src/transformers/models/phi3/modeling_phi3.py +++ b/src/transformers/models/phi3/modeling_phi3.py @@ -29,7 +29,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -433,8 +432,8 @@ def forward( @auto_docstring class Phi3ForCausalLM(Phi3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py b/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py index 03a06bce7094..db4cbd0eab6b 100644 --- a/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py +++ b/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py @@ -23,7 +23,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -182,26 +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": TPStyle( - "colwise", "allgather" - ), # we need to replicate here due to the slicing of qkv - "layers.*.self_attn.o_proj": TPStyle("vocab", "allreduce"), # input is replicated due to the slicing of qkv - "layers.*.mlp.gate_up_proj": TPStyle( - "colwise", "allgather" - ), # we need to replicate here due to the `chunk` operation - "layers.*.mlp.down_proj": TPStyle("vocab", "allreduce"), # 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": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.qkv_proj": TPStyle("colwise", "allgather"), # fused qkv needs full tensor for slicing - "layers.*.self_attn.o_proj": TPStyle("vocab", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_up_proj": TPStyle("colwise", "allgather"), # fused gate/up needs full tensor for chunk - "layers.*.mlp.down_proj": TPStyle("vocab", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), diff --git a/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py b/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py index abd1e9b09e7e..5c723f3d4365 100644 --- a/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py +++ b/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py @@ -32,7 +32,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_bidirectional_mask, create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -1597,8 +1596,8 @@ def forward( @auto_docstring class Phi4MultimodalForCausalLM(Phi4MultimodalPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/phimoe/modeling_phimoe.py b/src/transformers/models/phimoe/modeling_phimoe.py index cc76274c379a..97a095278065 100644 --- a/src/transformers/models/phimoe/modeling_phimoe.py +++ b/src/transformers/models/phimoe/modeling_phimoe.py @@ -30,7 +30,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_experts_implementation, use_kernel_func_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_layers import GenericForSequenceClassification, GradientCheckpointingLayer from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast @@ -773,8 +772,8 @@ 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": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/qwen2/configuration_qwen2.py b/src/transformers/models/qwen2/configuration_qwen2.py index dc16b48b56f6..599511b903a7 100644 --- a/src/transformers/models/qwen2/configuration_qwen2.py +++ b/src/transformers/models/qwen2/configuration_qwen2.py @@ -16,7 +16,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -45,28 +44,28 @@ class Qwen2Config(PreTrainedConfig): # Default tensor parallel plan for base model `Qwen2` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), diff --git a/src/transformers/models/qwen2/modeling_qwen2.py b/src/transformers/models/qwen2/modeling_qwen2.py index b8bcf9d9cdac..e22875b8c8a0 100644 --- a/src/transformers/models/qwen2/modeling_qwen2.py +++ b/src/transformers/models/qwen2/modeling_qwen2.py @@ -14,7 +14,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -417,8 +416,8 @@ def forward( @auto_docstring class Qwen2ForCausalLM(Qwen2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/qwen2_5_vl/configuration_qwen2_5_vl.py b/src/transformers/models/qwen2_5_vl/configuration_qwen2_5_vl.py index 4b21dda3712a..17e7d92f7d8f 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 @@ -27,7 +27,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -89,13 +88,13 @@ class Qwen2_5_VLTextConfig(PreTrainedConfig): default_theta = 1000000.0 # Default tensor parallel plan for base model `Qwen2_5_VL` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py b/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py index c0f5a7452fe8..5585a7454314 100644 --- a/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py +++ b/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py @@ -16,7 +16,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -55,13 +54,13 @@ class Qwen2MoeConfig(PreTrainedConfig): # Default tensor parallel plan for base model `Qwen2Moe` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py b/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py index 574f43b541df..272f8d4cfaf4 100644 --- a/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py +++ b/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py @@ -18,7 +18,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -66,13 +65,13 @@ class Qwen2VLTextConfig(PreTrainedConfig): default_theta = 1000000.0 # Default tensor parallel plan for base model `Qwen2VL` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/qwen3/configuration_qwen3.py b/src/transformers/models/qwen3/configuration_qwen3.py index 5aa1288e66a8..b3dc0e89a6bd 100644 --- a/src/transformers/models/qwen3/configuration_qwen3.py +++ b/src/transformers/models/qwen3/configuration_qwen3.py @@ -16,7 +16,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -46,13 +45,13 @@ class Qwen3Config(PreTrainedConfig): # 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": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } # TP + Sequence Parallelism plan (for training). @@ -62,21 +61,21 @@ class Qwen3Config(PreTrainedConfig): # 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": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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/qwen3/modeling_qwen3.py b/src/transformers/models/qwen3/modeling_qwen3.py index d8b2bf9973d1..2cf93f6f8ea2 100644 --- a/src/transformers/models/qwen3/modeling_qwen3.py +++ b/src/transformers/models/qwen3/modeling_qwen3.py @@ -28,7 +28,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -423,8 +422,8 @@ def forward( @auto_docstring class Qwen3ForCausalLM(Qwen3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/qwen3_5/modeling_qwen3_5.py b/src/transformers/models/qwen3_5/modeling_qwen3_5.py index 06a13e099c21..ade2351437ee 100644 --- a/src/transformers/models/qwen3_5/modeling_qwen3_5.py +++ b/src/transformers/models/qwen3_5/modeling_qwen3_5.py @@ -32,7 +32,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GenericForSequenceClassification, GradientCheckpointingLayer @@ -1688,8 +1687,8 @@ def forward( @auto_docstring class Qwen3_5ForCausalLM(Qwen3_5PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} config: Qwen3_5TextConfig _keys_to_ignore_on_load_unexpected = [r"^mtp.*", r"^model.visual.*"] 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 d2c9ff7df97f..a559d9cfb238 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 @@ -32,7 +32,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_experts_implementation, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -1895,8 +1894,8 @@ 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": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} config: Qwen3_5MoeTextConfig _keys_to_ignore_on_load_unexpected = [r"^mtp.*", r"^model.visual.*"] diff --git a/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py b/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py index f8f057b00ea8..66d3a22eec51 100644 --- a/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py +++ b/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py @@ -16,7 +16,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -55,40 +54,32 @@ class Qwen3MoeConfig(PreTrainedConfig): # Default tensor parallel plan for base model `Qwen3Moe` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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": "moe_experts_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather_split"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py b/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py index e101d4f78273..54b7ffe1d167 100644 --- a/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py +++ b/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py @@ -30,7 +30,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -586,8 +585,8 @@ 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": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/qwen3_next/configuration_qwen3_next.py b/src/transformers/models/qwen3_next/configuration_qwen3_next.py index c8bcaf895d93..1d1471d1db8c 100644 --- a/src/transformers/models/qwen3_next/configuration_qwen3_next.py +++ b/src/transformers/models/qwen3_next/configuration_qwen3_next.py @@ -16,7 +16,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -60,21 +59,17 @@ class Qwen3NextConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.shared_expert.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_expert.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.shared_expert.down_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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.shared_expert.gate_proj": "colwise", + "layers.*.mlp.shared_expert.up_proj": "colwise", + "layers.*.mlp.shared_expert.down_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/qwen3_next/modeling_qwen3_next.py b/src/transformers/models/qwen3_next/modeling_qwen3_next.py index c6ae090dfa3e..facb86e52521 100644 --- a/src/transformers/models/qwen3_next/modeling_qwen3_next.py +++ b/src/transformers/models/qwen3_next/modeling_qwen3_next.py @@ -30,7 +30,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_experts_implementation, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -1073,8 +1072,8 @@ 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": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, 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 f511c9114478..9c878256958c 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 @@ -21,7 +21,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring, logging @@ -261,13 +260,13 @@ class Qwen3OmniMoeTalkerCodePredictorConfig(PreTrainedConfig): # 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": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } # TP + Sequence Parallelism plan (for training). @@ -277,21 +276,21 @@ class Qwen3OmniMoeTalkerCodePredictorConfig(PreTrainedConfig): # 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": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), @@ -372,40 +371,32 @@ class Qwen3OmniMoeTalkerTextConfig(PreTrainedConfig): # Default tensor parallel plan for base model `Qwen3OmniMoeTalkerText` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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": "moe_experts_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather_split"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), 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 1308bcf1fcfa..b8998bd94709 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 @@ -35,7 +35,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -2621,8 +2620,8 @@ 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": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} config_class = Qwen3OmniMoeTalkerCodePredictorConfig base_model_prefix = "talker.code_predictor" @@ -3007,7 +3006,7 @@ def get_input_embeddings(self): class Qwen3OmniMoeTalkerForConditionalGeneration(Qwen3OmniMoeThinkerTextPreTrainedModel, GenerationMixin): _tied_weights_keys = {"codec_head": "model.codec_embedding.weight"} _tp_plan = {"codec_head": "colwise_allgather"} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"codec_head": (["hidden_states"], ["logits"])} config_class = Qwen3OmniMoeTalkerConfig base_model_prefix = "talker" diff --git a/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py b/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py index 997dbe290a71..7eba25bfc8b3 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 @@ -20,7 +20,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -66,26 +65,22 @@ class Qwen3VLMoeTextConfig(PreTrainedConfig): "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.self_attn.q_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.self_attn.k_norm": TPStyle("activation", "none", sequence_dim=2), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather_split"), - "layers.*.mlp.experts": TPStyle( - "moe_experts", - "allreduce", - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, - ), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/seed_oss/configuration_seed_oss.py b/src/transformers/models/seed_oss/configuration_seed_oss.py index 5902cbac4ae7..d60348c19149 100644 --- a/src/transformers/models/seed_oss/configuration_seed_oss.py +++ b/src/transformers/models/seed_oss/configuration_seed_oss.py @@ -16,7 +16,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -46,28 +45,28 @@ class SeedOssConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `SeedOssModel` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "allreduce"), + "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", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.gate_proj": TPStyle("colwise", "none"), - "layers.*.mlp.up_proj": TPStyle("colwise", "none"), - "layers.*.mlp.down_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), diff --git a/src/transformers/models/seed_oss/modeling_seed_oss.py b/src/transformers/models/seed_oss/modeling_seed_oss.py index b37084599b94..fa524a9de3b2 100644 --- a/src/transformers/models/seed_oss/modeling_seed_oss.py +++ b/src/transformers/models/seed_oss/modeling_seed_oss.py @@ -28,7 +28,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import ( GenericForQuestionAnswering, @@ -430,8 +429,8 @@ def forward( @auto_docstring class SeedOssForCausalLM(SeedOssPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/smollm3/modeling_smollm3.py b/src/transformers/models/smollm3/modeling_smollm3.py index 9623ef48a659..d9bb9d7473b5 100644 --- a/src/transformers/models/smollm3/modeling_smollm3.py +++ b/src/transformers/models/smollm3/modeling_smollm3.py @@ -28,7 +28,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -446,8 +445,8 @@ def forward( @auto_docstring class SmolLM3ForCausalLM(SmolLM3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/solar_open/modeling_solar_open.py b/src/transformers/models/solar_open/modeling_solar_open.py index ac114dbbb82b..1c3b11a19dc7 100644 --- a/src/transformers/models/solar_open/modeling_solar_open.py +++ b/src/transformers/models/solar_open/modeling_solar_open.py @@ -34,7 +34,6 @@ use_kernel_func_from_hub, use_kernelized_func, ) -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -553,8 +552,8 @@ def forward( @auto_docstring class SolarOpenForCausalLM(SolarOpenPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/starcoder2/configuration_starcoder2.py b/src/transformers/models/starcoder2/configuration_starcoder2.py index 45c349840963..f508dcce3cbd 100644 --- a/src/transformers/models/starcoder2/configuration_starcoder2.py +++ b/src/transformers/models/starcoder2/configuration_starcoder2.py @@ -16,7 +16,6 @@ from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig -from ...integrations.tensor_parallel import TPStyle from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -46,26 +45,26 @@ class Starcoder2Config(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `Starcoder2` base_model_tp_plan = { - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "allreduce"), - "layers.*.mlp.c_fc": TPStyle("colwise", "none"), - "layers.*.mlp.c_proj": TPStyle("rowwise", "allreduce"), + "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.c_fc": "colwise", + "layers.*.mlp.c_proj": "rowwise_allreduce", } base_model_sp_plan = { - "embed_tokens": TPStyle("vocab", "reduce_scatter"), - "layers.*.input_layernorm": TPStyle("activation", "none"), - "layers.*.self_attn": TPStyle("module", "allgather", input_key="hidden_states"), - "layers.*.self_attn.q_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.k_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.v_proj": TPStyle("colwise", "none"), - "layers.*.self_attn.o_proj": TPStyle("rowwise", "reduce_scatter"), - "layers.*.post_attention_layernorm": TPStyle("activation", "none"), - "layers.*.mlp": TPStyle("module", "allgather"), - "layers.*.mlp.c_fc": TPStyle("colwise", "none"), - "layers.*.mlp.c_proj": TPStyle("rowwise", "reduce_scatter"), - "norm": TPStyle("activation", "none"), + "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"]), diff --git a/src/transformers/models/starcoder2/modeling_starcoder2.py b/src/transformers/models/starcoder2/modeling_starcoder2.py index cbfb88c2db48..fff6c00cdd7d 100644 --- a/src/transformers/models/starcoder2/modeling_starcoder2.py +++ b/src/transformers/models/starcoder2/modeling_starcoder2.py @@ -33,7 +33,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_func_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -410,8 +409,8 @@ def forward( @auto_docstring class Starcoder2ForCausalLM(Starcoder2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/vaultgemma/modeling_vaultgemma.py b/src/transformers/models/vaultgemma/modeling_vaultgemma.py index 901d585660e8..e392829c8fad 100644 --- a/src/transformers/models/vaultgemma/modeling_vaultgemma.py +++ b/src/transformers/models/vaultgemma/modeling_vaultgemma.py @@ -30,7 +30,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_func_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -468,8 +467,8 @@ def forward( @auto_docstring class VaultGemmaForCausalLM(VaultGemmaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py b/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py index 633aa22edb87..d0af6cdcfd20 100644 --- a/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py +++ b/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py @@ -32,7 +32,6 @@ from ...cache_utils import Cache, DynamicCache, StaticCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -824,8 +823,8 @@ def forward( @auto_docstring class VoxtralRealtimeTextForCausalLM(VoxtralRealtimeTextPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/youtu/modeling_youtu.py b/src/transformers/models/youtu/modeling_youtu.py index c9ea544decf2..e76855d5f518 100644 --- a/src/transformers/models/youtu/modeling_youtu.py +++ b/src/transformers/models/youtu/modeling_youtu.py @@ -37,7 +37,6 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub -from ...integrations.tensor_parallel import TPStyle from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -534,8 +533,8 @@ def forward( @auto_docstring class YoutuForCausalLM(YoutuPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": TPStyle("colwise", "allgather")} - _sp_plan = {"lm_head": TPStyle("colwise", "loss_parallel")} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): From 061d4e69d63c0e343534e21cacb65cda7a43dc63 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 28 Apr 2026 10:29:47 +0000 Subject: [PATCH 054/116] delete TPStyle dataclass --- src/transformers/integrations/__init__.py | 2 - .../integrations/tensor_parallel.py | 147 +++--------------- 2 files changed, 18 insertions(+), 131 deletions(-) diff --git a/src/transformers/integrations/__init__.py b/src/transformers/integrations/__init__.py index bbe3461402c1..4be557be670a 100755 --- a/src/transformers/integrations/__init__.py +++ b/src/transformers/integrations/__init__.py @@ -162,7 +162,6 @@ _import_structure["tensor_parallel"] = [ "ALL_PARALLEL_STYLES", - "TPStyle", "apply_tensor_parallel", "convert_strided_to_shard", "gather_full_state_dict", @@ -301,7 +300,6 @@ from .spqr import replace_with_spqr_linear from .tensor_parallel import ( ALL_PARALLEL_STYLES, - TPStyle, apply_tensor_parallel, convert_strided_to_shard, gather_full_state_dict, diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index 065e7ff314cb..c6b2d43f7fc9 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -15,9 +15,7 @@ import contextlib import re -from abc import ABC -from dataclasses import dataclass -from typing import Literal, abstractmethod +from abc import ABC, abstractmethod from ..utils import logging from ..utils.generic import GeneralInterface @@ -213,7 +211,7 @@ def _resolve(d, dotted_key): container[leaf_key] = _replicate_dtensor(container[leaf_key]).redistribute(placements=original_placements) -def verify_tp_plan(expected_keys: list[str], tp_plan: dict[str, str | TPStyle] | None): +def verify_tp_plan(expected_keys: list[str], tp_plan: dict[str, str] | None): """ Verify the TP plan of the model, log a warning if the layers that were not sharded and the rules that were not applied. @@ -225,9 +223,12 @@ def verify_tp_plan(expected_keys: list[str], tp_plan: dict[str, str | TPStyle] | if tp_plan is None: return - # Filter out module-level comm hooks — they don't shard weights - _NON_WEIGHT_KINDS = {"activation", "module"} - weight_plan = {k: v for k, v in tp_plan.items() if not isinstance(v, TPStyle) or v.kind not in _NON_WEIGHT_KINDS} + # Filter out module-level comm hooks — they don't shard weights. + # Plan values are registry names; entries beginning with "activation" or "module" + # configure communication hooks rather than parameter sharding. + weight_plan = { + k: v for k, v in tp_plan.items() if not (v == "activation" or v.startswith(("activation_", "module_"))) + } generic_keys = {replace_layer_number_by_wildcard(key) for key in expected_keys} unsharded_layers = set(generic_keys) @@ -548,104 +549,6 @@ def tp_forward(hidden_states, top_k_index, top_k_weights): return module -@dataclass(frozen=True) -class TPStyle: - kind: Literal["colwise", "packed_colwise", "rowwise", "vocab", "activation", "module", "moe_experts"] - comm: Literal["none", "allreduce", "reduce_scatter", "allgather", "allgather_split", "loss_parallel"] - sequence_dim: int = 1 - use_local_output: bool = True - input_key: str | None = None - shard_plan: dict[str, str] | None = None - - def to_dtensor_style(self) -> ParallelStyle: - """Convert to the corresponding PyTorch DTensor ParallelStyle.""" - if self.kind == "colwise": - match self.comm: - case "none": - return ColwiseParallel( - input_layouts=Replicate(), output_layouts=Shard(-1), use_local_output=self.use_local_output - ) - case "allgather": - return ColwiseParallel( - input_layouts=Replicate(), - output_layouts=Replicate(), - use_local_output=self.use_local_output, - ) - case "loss_parallel": - return ColwiseParallel(input_layouts=Shard(1), output_layouts=Shard(-1), use_local_output=False) - elif self.kind == "packed_colwise": - match self.comm: - case "none": - return PackedColwiseParallel(input_layouts=Replicate(), use_local_output=self.use_local_output) - elif self.kind == "rowwise": - match self.comm: - case "allreduce": - return RowwiseParallel( - input_layouts=Shard(-1), - output_layouts=Replicate(), - use_local_output=self.use_local_output, - ) - case "reduce_scatter": - return RowwiseParallel( - input_layouts=Shard(-1), - output_layouts=Shard(1), - use_local_output=self.use_local_output, - ) - elif self.kind == "vocab": - match self.comm: - case "allreduce": - return RowwiseParallel( - input_layouts=Replicate(), - output_layouts=Replicate(), - use_local_output=self.use_local_output, - ) - case "reduce_scatter": - return RowwiseParallel( - input_layouts=Replicate(), output_layouts=Shard(1), use_local_output=self.use_local_output - ) - elif self.kind == "activation": - match self.comm: - case "none": - return SequenceParallel(sequence_dim=self.sequence_dim, use_local_output=self.use_local_output) - elif self.kind == "module": - match self.comm: - case "allgather": - if self.input_key is not None: - return PrepareModuleInput( - input_kwarg_layouts={self.input_key: Shard(1)}, - desired_input_kwarg_layouts={self.input_key: Replicate()}, - use_local_output=self.use_local_output, - ) - return PrepareModuleInput( - input_layouts=(Shard(1),), - desired_input_layouts=(Replicate(),), - use_local_output=self.use_local_output, - ) - case "allgather_split": - return PrepareModuleInputOutput(use_local_output=self.use_local_output) - elif self.kind == "moe_experts": - match self.comm: - case "allreduce": - return MoEExpertsParallel(output_layouts=Replicate()) - case "reduce_scatter": - return MoEExpertsParallel(output_layouts=Shard(1)) - raise ValueError( - f"Invalid TPStyle({self.kind!r}, {self.comm!r}). Valid combinations:\n" - f" colwise: none, allgather, loss_parallel\n" - f" packed_colwise: none\n" - f" rowwise: allreduce, reduce_scatter\n" - f" vocab: allreduce, reduce_scatter\n" - f" activation: none\n" - f" module: allgather, allgather_split\n" - f" moe_experts: allreduce, reduce_scatter" - ) - - def __str__(self): - if self.comm == "none": - return self.kind - return f"{self.kind}_{self.comm}" - - class ParallelInterface(GeneralInterface): """Registry of named TP styles. Configs and modeling files reference these by string name. @@ -700,7 +603,8 @@ def apply_tensor_parallel(model, tp_mesh, tp_plan): """Apply tensor parallelism using PyTorch's parallelize_module. Converts the wildcard tp_plan from model config into a concrete plan - for ``parallelize_module``. Plan values is a `TPStyle`` instances + for ``parallelize_module``. Plan values are string names looked up in + ``ALL_PARALLEL_STYLES``. """ if tp_plan is None: return model @@ -732,28 +636,16 @@ def apply_tensor_parallel(model, tp_mesh, tp_plan): if style_value is None: continue - if isinstance(style_value, str): - if style_value not in ALL_PARALLEL_STYLES: - raise ValueError( - f"Unknown TP style {style_value!r} for module {name!r}. " - f"Valid styles: {sorted(ALL_PARALLEL_STYLES)}" - ) - parallelize_plan[name] = ALL_PARALLEL_STYLES[style_value] - elif isinstance(style_value, TPStyle): - dtensor_style = style_value.to_dtensor_style() - parallelize_plan[name] = dtensor_style - # For MoE modules, attach the per-parameter shard plan from TPStyle - # so _partition_fn can create DTensors with the correct placements. - if isinstance(dtensor_style, MoEExpertsParallel) and style_value.shard_plan: - dtensor_style._moe_shard_plan = style_value.shard_plan - elif isinstance(style_value, ParallelStyle): - parallelize_plan[name] = style_value - else: + if not isinstance(style_value, str): raise TypeError( f"Unsupported plan value for '{name}': {style_value!r} (type {type(style_value).__name__}). " - f"TP plan values must be strings (looked up in ALL_PARALLEL_STYLES), TPStyle, " - f"or ParallelStyle instances." + f"TP plan values must be strings looked up in ALL_PARALLEL_STYLES." + ) + if style_value not in ALL_PARALLEL_STYLES: + raise ValueError( + f"Unknown TP style {style_value!r} for module {name!r}. Valid styles: {sorted(ALL_PARALLEL_STYLES)}" ) + parallelize_plan[name] = ALL_PARALLEL_STYLES[style_value] parallelize_module(model, tp_mesh, parallelize_plan) @@ -783,10 +675,7 @@ def _inject_sp_metadata(mod, args, kwargs): # 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" or (isinstance(v, TPStyle) and v.comm == "loss_parallel") - for v in tp_plan.values() - ) + has_loss_parallel = any(v == "colwise_loss_parallel" for v in tp_plan.values()) if has_loss_parallel: from torch.distributed.tensor.parallel import loss_parallel From 8e0f60c16ae6d6df544b1ee81887442a4d478469 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 28 Apr 2026 10:46:41 +0000 Subject: [PATCH 055/116] fix use_local_output defaults for SequenceParallel and PrepareModuleInput in registry --- src/transformers/integrations/tensor_parallel.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index c6b2d43f7fc9..fde48fe0400e 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -576,13 +576,19 @@ class ParallelInterface(GeneralInterface): "vocab_allreduce": RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()), "vocab_reduce_scatter": RowwiseParallel(input_layouts=Replicate(), output_layouts=Shard(1)), # Activation / norm (sequence-parallel passthrough) - "activation": SequenceParallel(), - "activation_seq_dim_2": SequenceParallel(sequence_dim=2), - # Module-level prepare-input - "module_allgather": PrepareModuleInput(input_layouts=(Shard(1),), desired_input_layouts=(Replicate(),)), + # use_local_output=True: torch defaults to False here, but downstream modeling + # code expects plain tensors, not DTensors. + "activation": SequenceParallel(use_local_output=True), + "activation_seq_dim_2": SequenceParallel(sequence_dim=2, use_local_output=True), + # Module-level prepare-input. Same use_local_output=True override as above — + # torch's default is False, our modeling code expects plain tensors downstream. + "module_allgather": PrepareModuleInput( + input_layouts=(Shard(1),), desired_input_layouts=(Replicate(),), use_local_output=True + ), "module_allgather_hidden_states": PrepareModuleInput( input_kwarg_layouts={"hidden_states": Shard(1)}, desired_input_kwarg_layouts={"hidden_states": Replicate()}, + use_local_output=True, ), "module_allgather_split": PrepareModuleInputOutput(), # MoE — canonical shard_plan baked in (only variant in use across configs) From 5b336bd9db9befb39487a0aacf0bf352ebea8afa Mon Sep 17 00:00:00 2001 From: 3outeille Date: Thu, 30 Apr 2026 03:10:42 +0000 Subject: [PATCH 056/116] use parallel style from torch --- .../integrations/tensor_parallel.py | 18 ------------------ 1 file changed, 18 deletions(-) diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index fde48fe0400e..2a82a4e55067 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -15,7 +15,6 @@ import contextlib import re -from abc import ABC, abstractmethod from ..utils import logging from ..utils.generic import GeneralInterface @@ -24,11 +23,9 @@ if is_torch_available(): import torch - import torch.nn as nn if is_torch_available() and is_torch_greater_or_equal("2.5"): import torch.distributed as dist - from torch.distributed.device_mesh import DeviceMesh from torch.distributed.tensor import DTensor, Partial, Replicate, Shard, distribute_tensor from torch.distributed.tensor.parallel import ( ColwiseParallel, @@ -251,21 +248,6 @@ def verify_tp_plan(expected_keys: list[str], tp_plan: dict[str, str] | None): logger.warning(f"The following layers were not sharded: {', '.join(unsharded_layers)}") -class ParallelStyle(ABC): - """ - Import from torch.distributed.tensor.parallel.style.ParallelStyle to avoid import guarding every class that inherits from it. - The parallel style contract defines how the module or submodule should be parallelized. - - It only defines the ``apply`` method for ``parallelize_module`` to use, this allows maximum - flexibility for different kind of style implementations. - """ - - src_data_rank: int | None = 0 - - @abstractmethod - def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module: ... - - class PrepareModuleInputOutput(ParallelStyle): """Allgather input (Shard(1) → Replicate) + local split output (Replicate → Shard(1)). From 465d029563033928ee48a8e4054ff4fbce81f36c Mon Sep 17 00:00:00 2001 From: 3outeille Date: Fri, 1 May 2026 03:16:13 +0000 Subject: [PATCH 057/116] revert changes in weight converter --- src/transformers/core_model_loading.py | 47 +++++++------------------- 1 file changed, 12 insertions(+), 35 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 7e0d14a21237..60addb1b46fc 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -775,41 +775,18 @@ def convert( pass if hf_quantizer is not None and self.quantization_operation is not None: - if len(collected_tensors) > 1 and model is not None: - quantized_tensors = {} - for target_key, tensor in collected_tensors.items(): - if not hf_quantizer.param_needs_quantization(model, target_key): - quantized_tensors[target_key] = tensor - continue - quantize_input = tensor if isinstance(tensor, list) else [tensor] - with log_conversion_errors( - target_key, loading_info, (len(quantize_input), target_key), self.quantization_operation - ): - quantized_tensors.update( - self.quantization_operation.convert( - {target_key: quantize_input}, - source_patterns=self.source_patterns, - target_patterns=[target_key], - full_layer_name=target_key, - config=config, - model=model, - missing_keys=loading_info.missing_keys if loading_info else None, - ) - ) - collected_tensors = quantized_tensors - else: - with log_conversion_errors( - layer_name, loading_info, (len(collected_tensors), layer_name), self.quantization_operation - ): - collected_tensors = self.quantization_operation.convert( - collected_tensors, - source_patterns=self.source_patterns, - target_patterns=self.target_patterns, - full_layer_name=layer_name, - config=config, - model=model, - missing_keys=loading_info.missing_keys if loading_info else None, - ) + with log_conversion_errors( + layer_name, loading_info, (len(collected_tensors), layer_name), self.quantization_operation + ): + collected_tensors = self.quantization_operation.convert( + collected_tensors, + source_patterns=self.source_patterns, + target_patterns=self.target_patterns, + full_layer_name=layer_name, + config=config, + model=model, + missing_keys=loading_info.missing_keys if loading_info else None, + ) return collected_tensors From bc6d6f9feb3cd1745b0930208f2bbe16fb15bbbf Mon Sep 17 00:00:00 2001 From: 3outeille Date: Fri, 1 May 2026 03:37:33 +0000 Subject: [PATCH 058/116] remove dead code in set_param_for_module --- src/transformers/core_model_loading.py | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 60addb1b46fc..7ac4b8fe5383 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -1187,19 +1187,12 @@ def set_param_for_module( loading_info.missing_keys.discard(target_name) if isinstance(ref, DTensor): - local_shape, global_offset = compute_local_shape_and_global_offset( - ref.shape, ref.device_mesh, ref.placements - ) + local_shape, _ = compute_local_shape_and_global_offset(ref.shape, ref.device_mesh, ref.placements) expected_shape = torch.Size(local_shape) else: expected_shape = ref.shape - # When a WeightConverter produces the full global tensor, slice it to the local DTensor shard. - if isinstance(ref, DTensor) and param_value.shape == ref.shape and param_value.shape != expected_shape: - slices = [slice(global_offset[d], global_offset[d] + local_shape[d]) for d in range(param_value.ndim)] - param_value = param_value[tuple(slices)].contiguous() - - if ref is not None and param_value.shape != expected_shape and hf_quantizer is None: + if 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): From f305f92480c6939d9677af54af6ceae5e67d2500 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Fri, 1 May 2026 03:50:43 +0000 Subject: [PATCH 059/116] remove dead code --- src/transformers/core_model_loading.py | 47 -------------------------- 1 file changed, 47 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 7ac4b8fe5383..68bd985bac25 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -83,21 +83,6 @@ def build_glob_alternation( return alternation, src_group_to_glob, tgt_group_to_glob -def resolve_target_wildcards(source_pattern: str, target_pattern: str, source_key: str) -> str: - if "*" not in target_pattern or "*" not in source_pattern: - return target_pattern - - wildcard_regex = re.escape(source_pattern).replace(r"\*", r"(.*?)") - match = re.fullmatch(wildcard_regex, source_key) - if match is None: - return target_pattern - - resolved_target = target_pattern - for wildcard_value in match.groups(): - resolved_target = resolved_target.replace("*", wildcard_value, 1) - return resolved_target - - class ConversionOps: """Base class for weight conversion operations.""" @@ -620,7 +605,6 @@ def rename_source_key(self, source_key: str) -> tuple[str, str | None]: source_pattern_that_matched = self.source_patterns[int(matching_group_name[1:])] # If we matched, we always replace with the first target pattern, in case we have several (one to many transform) replacement = self.target_patterns[0] - replacement = resolve_target_wildcards(source_pattern_that_matched, replacement, source_key) # Allow capturing groups in patterns, i.e. to add a prefix to all keys (e.g. timm_wrapper, sam3) if r"\1" in replacement: # The index of the internal group we need to replace is the index of the matched named group as it comes @@ -1282,30 +1266,6 @@ def rename_source_key( return renamed_key, source_pattern -def concretize_target_patterns( - converter: WeightConverter, - source_key: str, - source_pattern: str, - prefix: str | None, - meta_state_dict: dict | None, -) -> WeightConverter: - concrete_targets = [] - for target_pattern in converter.target_patterns: - concrete_target = resolve_target_wildcards(source_pattern, target_pattern, source_key) - if prefix is not None and meta_state_dict is not None: - if ( - concrete_target.startswith(prefix) - and meta_state_dict.get(re.sub(f"^{prefix}.", "", concrete_target, count=1)) is not None - ): - concrete_target = re.sub(f"^{prefix}.", "", concrete_target, count=1) - elif meta_state_dict.get(f"{prefix}.{concrete_target}") is not None: - concrete_target = f"{prefix}.{concrete_target}" - concrete_targets.append(concrete_target) - - object.__setattr__(converter, "target_patterns", concrete_targets) - return converter - - def convert_and_load_state_dict_in_model( model: PreTrainedModel, state_dict: dict[str, Any], @@ -1479,13 +1439,6 @@ def convert_and_load_state_dict_in_model( # If we enter here, we have a WeightConverter operation to perform if source_pattern is not None: new_converter = deepcopy(pattern_to_converter[source_pattern]) - new_converter = concretize_target_patterns( - new_converter, - original_key, - source_pattern, - prefix, - meta_model_state_dict, - ) # each target key gets its own converter instance mapping = param_name_to_load.setdefault(renamed_key, new_converter) # Otherwise, only potential renaming From 39db8c17022e307e29c43f27f412c74b307988b9 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Fri, 1 May 2026 04:01:13 +0000 Subject: [PATCH 060/116] cleaning again --- src/transformers/core_model_loading.py | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 68bd985bac25..2144c36e52dd 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -1427,15 +1427,7 @@ def convert_and_load_state_dict_in_model( # 2. finally, collect the tensor into the proper converter if renamed_key in meta_model_state_dict: - empty_param = meta_model_state_dict[renamed_key] - try: - empty_param = model.get_parameter_or_buffer(renamed_key) - except (AttributeError, KeyError): - if getattr(model, "_is_fsdp_managed_module", False): - raise RuntimeError( - f"FSDP shard-on-read requires the live parameter for {renamed_key!r}, " - f"but get_parameter_or_buffer() failed." - ) + empty_param = meta_model_state_dict.get(renamed_key) # If we enter here, we have a WeightConverter operation to perform if source_pattern is not None: new_converter = deepcopy(pattern_to_converter[source_pattern]) From 951d4ae351b515c2478b272b4383b1b391202c6d Mon Sep 17 00:00:00 2001 From: 3outeille Date: Fri, 1 May 2026 04:33:09 +0000 Subject: [PATCH 061/116] cleaning --- src/transformers/core_model_loading.py | 13 ++++--------- 1 file changed, 4 insertions(+), 9 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 2144c36e52dd..6627e6f9ea5f 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -1170,18 +1170,14 @@ def set_param_for_module( # Remove from missing keys (it's either mismatched, or all good) loading_info.missing_keys.discard(target_name) - if isinstance(ref, DTensor): - local_shape, _ = compute_local_shape_and_global_offset(ref.shape, ref.device_mesh, ref.placements) - expected_shape = torch.Size(local_shape) - else: - expected_shape = ref.shape + expected_shape = ref.to_local().shape if isinstance(ref, DTensor) else ref.shape if 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 - fsdp_param = DTensor.from_local( + dtensor_param = DTensor.from_local( local_param.contiguous(), ref.device_mesh, ref.placements, @@ -1191,10 +1187,9 @@ def set_param_for_module( ) with torch.no_grad(): if ref.is_meta: - fsdp_param = torch.nn.Parameter(fsdp_param, requires_grad=ref.requires_grad) - torch.utils.swap_tensors(ref, fsdp_param) + torch.utils.swap_tensors(ref, torch.nn.Parameter(dtensor_param, requires_grad=ref.requires_grad)) else: - ref.copy_(fsdp_param) + ref.copy_(dtensor_param) ref._is_hf_initialized = True else: # super important otherwise _init_weight will re-init the param From 1b040ef75dcbc614e786ca72c252a04246dc5320 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Fri, 1 May 2026 04:41:21 +0000 Subject: [PATCH 062/116] revert change --- src/transformers/core_model_loading.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 6627e6f9ea5f..a6500d079816 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -1172,7 +1172,7 @@ def set_param_for_module( expected_shape = ref.to_local().shape if isinstance(ref, DTensor) else ref.shape - if param_value.shape != expected_shape and hf_quantizer is None: + 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): From 85ef27c6b5a6221c102ce7c50b5b8752e4b7f191 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Fri, 1 May 2026 04:45:24 +0000 Subject: [PATCH 063/116] linting --- src/transformers/core_model_loading.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index a6500d079816..f270c441d340 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -1187,7 +1187,9 @@ def set_param_for_module( ) with torch.no_grad(): if ref.is_meta: - torch.utils.swap_tensors(ref, torch.nn.Parameter(dtensor_param, requires_grad=ref.requires_grad)) + torch.utils.swap_tensors( + ref, torch.nn.Parameter(dtensor_param, requires_grad=ref.requires_grad) + ) else: ref.copy_(dtensor_param) ref._is_hf_initialized = True From 1fd7b1d043e34a2a330f46bd67a88aa894087fd7 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Mon, 4 May 2026 05:23:31 +0000 Subject: [PATCH 064/116] refactor dtensor shard ops --- src/transformers/core_model_loading.py | 223 +++++-------------------- tests/utils/test_core_model_loading.py | 20 ++- 2 files changed, 62 insertions(+), 181 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index f270c441d340..1416cced8d41 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -816,120 +816,68 @@ def _job(): class DtensorShardOperation: - """Read only this rank's local shard out of a checkpoint tensor. - We first need to classify the source tensor (checkpoint tensor) relative to - the destination `param` (the Dtensor sharded parameter we are loading into). - This decides how we must slice it before writing into the local - DTensor shard. - - (a) Full weight — source.ndim == param.ndim - The checkpoint tensor has the same rank as the model param. - Example: param = (8, 8), source = (8, 8). - → Apply placements directly via the generic interval loop. - - (b) One expert — source.ndim == param.ndim - 1 - MoE models stack experts along a leading axis (E, ...) in the - model, but checkpoints store each expert in its own file - (e.g. `experts.2.w1.weight`). The source is missing that - leading expert axis; `tensor_idx` names which expert it is. - Example: param = (E=4, H=8, I=4), source = (H=8, I=4), - tensor_idx = 2. - → If Shard(0) (the expert axis) is in placements, first decide - whether this rank owns expert `tensor_idx`: - - not owned: return None (file belongs to other ranks); - - owned: drop Shard(0) from placements and fall through - to the generic loop for the remaining inner - dims (e.g. Shard(1) on H). - - (c) One half of a pack — source smaller than param on the packed axis - Some params are built by concatenating two checkpoint tensors - (gate+up → gate_up, Q+K+V → qkv). Each half is loaded on its own, - so source is smaller than param on that axis. - Example: param = (2H, D), source = (H, D) for just `w1`. - → _StridedShard can't stride a pack that doesn't exist yet, so on - the packed axis it falls back to a plain contiguous cut. The - WeightConverter concatenates the halves afterwards. - - (b) + (c) co-occurring - MoE checkpoint that is both per-expert and pre-pack (e.g. - `experts.2.w1.weight`). Resolve the expert axis first (b); the - generic loop then handles the remaining dims with (c) behavior. - """ def __init__(self, param: DTensor): self.device_mesh = param.device_mesh self.placements = tuple(param.placements) - self.param_shape = tuple(param.shape) self.param_ndim = param.ndim - local_shape, _ = compute_local_shape_and_global_offset(param.shape, self.device_mesh, self.placements) - self.local_shape = tuple(local_shape) + local_shape, offsets = compute_local_shape_and_global_offset( + param.shape, self.device_mesh, self.placements + ) + self._first_owned_expert = offsets[0] + self._owned_experts_count = local_shape[0] def shard_tensor( self, source: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None ) -> torch.Tensor | None: 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")] - source_missing_leading_axis = self.param_ndim > len(source_shape) - - # Collect placements that actually split a dim. - # _StridedShard.is_shard() returns False in PyTorch, so also accept - # any non-Replicate placement that exposes a `dim` attribute. - placements = [ - (mesh_dim, p) - for mesh_dim, p in enumerate(self.placements) - if p.is_shard() or (hasattr(p, "dim") and not p.is_replicate()) - ] - if not placements: - return source[...].to(device=device, dtype=dtype) # no sharding → full copy - - source_is_one_expert = tensor_idx is not None and self.param_ndim == len(source_shape) + 1 - if source_is_one_expert and any(self._norm_dim(p.dim) == 0 for _, p in placements): - if not self._owns_expert(tensor_idx): # Case (b) -> Not owned - return None - # Case (b) -> Owned, drop expert axis and continue with the generic interval loop - placements = [(mesh_dim, p) for mesh_dim, p in placements if self._norm_dim(p.dim) != 0] + if tensor_idx is None: + # Source already has the param's shape -> slice every dim directly. if not placements: - # Expert axis was the only sharding → keep the whole expert tensor. return source[...].to(device=device, dtype=dtype) - - # Example - case (a) full weight - # input (source_shape=(8, 16), placements=[Shard(0)], world_size=2): - # intervals -> [[(0, 8)], [(0, 16)]] # whole range on every source dim - # output: - # rank 0: intervals -> [[(0, 4)], [(0, 16)]] # → source[0:4, 0:16] - # rank 1: intervals -> [[(4, 8)], [(0, 16)]] # → source[4:8, 0:16] - intervals: list[list[tuple[int, int]]] = [[(0, size)] for size in source_shape] - for mesh_dim, placement in placements: - source_dim = self._param_dim_to_source_dim(placement.dim, source_shape) + has_strided = any(not p.is_shard() for _, p in placements) + intervals: list[list[tuple[int, int]]] = [[(0, size)] for size in source_shape] + for mesh_dim, placement in placements: + 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) + + # Source is one of N pieces. Ownership check on the leading axis; + # MergeModulelist / Concatenate will assemble the param from kept pieces. + # Inner dims only use _contiguous_intervals, which yields one piece per dim, + # so we can slice directly -- no cat needed. + shards_expert_axis = any(self._norm_dim(p.dim) == 0 for _, p in placements) + owns_expert = self._first_owned_expert <= tensor_idx < self._first_owned_expert + self._owned_experts_count + if shards_expert_axis and not owns_expert: + return None + 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] + 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() - - is_interleaved = not placement.is_shard() and not source_missing_leading_axis - if is_interleaved: - 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) - - # Read the source with those intervals (concatenating along a multi-interval dim if any). - return self._slice_and_cat(source, intervals, device, dtype) + source_dim = self._norm_dim(placement.dim) - 1 # exactly one leading dim missing + 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]]: - """Split each interval into `split_factor` groups, then shard contiguously within each group. - - Produces one sub-interval per group (up to `split_factor` disjoint pieces - per input interval), matching _StridedShard's packed-axis layout. - - Example: - intervals = [(0, 4)], world_size = 2, split_factor = 2 - group 0 = (0, 2), group 1 = (2, 4). - Within each group each rank owns half: - rank 0 → [(0, 1), (2, 3)] (first half of each group) - rank 1 → [(1, 2), (3, 4)] (second half of each group) - """ narrowed = [] for start, end in intervals: group_size = math.ceil((end - start) / split_factor) @@ -946,21 +894,6 @@ def _strided_intervals( def _contiguous_intervals( self, intervals: list[tuple[int, int]], rank: int, world_size: int ) -> list[tuple[int, int]]: - """Shard semantics: treat `intervals` as one flat logical sequence and take this rank's contiguous chunk of it. - - Examples: - Single interval, plain split: - intervals = [(0, 8)], world_size = 2 - rank 0 → [(0, 4)] - rank 1 → [(4, 8)] - - Multiple intervals from a prior _StridedShard, chunk stays - inside one input interval: - intervals = [(0, 1), (2, 3)], world_size = 2 - flat view = [row 0, row 2], rank 0 owns first half - rank 0 → [(0, 1)] (row 0 only) - rank 1 → [(2, 3)] (row 2 only) - """ 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: @@ -984,28 +917,6 @@ def _contiguous_intervals( return out def _slice_and_cat(self, source, intervals, device, dtype): - """Read `source` with per-source-dim intervals; concat along the sole multi-interval dim if any. - - Each entry of ``intervals`` holds the (start, end) pieces this rank - owns on that source dim. At most one dim may have more than one - piece (from _StridedShard); two such dims would require a 2D outer - product of reads and are rejected. - - Examples: - 1) Single interval per dim — one contiguous read: - source shape = [8, 4] - intervals = [[(0, 4)], [(0, 4)]] - → source[0:4, 0:4] - - 2) Multi-interval on one dim — read each piece, concat on that dim: - source shape = [8, 4] - intervals = [[(0, 2), (4, 6)], [(0, 4)]] - → cat([source[0:2, 0:4], source[4:6, 0:4]], dim=0) - - 3) Multi-interval on two dims — rejected: - intervals = [[(0, 2), (4, 6)], [(0, 1), (2, 3)]] - → ValueError (would require a 2D outer product of reads) - """ multi_interval_dim: int | None = None slices: list[slice] = [] for source_dim, pieces in enumerate(intervals): @@ -1030,56 +941,14 @@ def _slice_and_cat(self, source, intervals, device, dtype): pieces_read.append(source[tuple(piece_slices)]) return torch.cat(pieces_read, dim=multi_interval_dim).to(device=device, dtype=dtype) - def _owns_expert(self, expert_idx: int) -> bool: - """True when this rank's shard of the expert axis (param dim 0) contains expert_idx.""" - _, offsets = compute_local_shape_and_global_offset( - torch.Size(self.param_shape), self.device_mesh, self.placements - ) - first_owned_expert = offsets[0] - return first_owned_expert <= expert_idx < first_owned_expert + self.local_shape[0] - def _get_sub_mesh(self, mesh_dim: int): - if self.device_mesh.ndim > 1: - return self.device_mesh[self.device_mesh.mesh_dim_names[mesh_dim]] - return self.device_mesh + 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: return dim if dim >= 0 else self.param_ndim + dim - def _param_dim_to_source_dim(self, placement_dim: int, source_shape) -> int: - """Map a placement's param dim onto the corresponding source-tensor dim. - - When the source has fewer dims than the param (expert or pre-pack - source), leading axes are absent: any placement dim past the missing - prefix shifts down by the difference. - - Examples: - 1) No missing dims (common case) — source and param align, no shift: - param shape = [out, in], source shape = [out, in] - placement_dim = 1 → source_dim = 1 - - 2) Stacked experts — source is one expert (missing `num_experts` axis): - param shape = [8, 4096, 2048], source shape = [4096, 2048] - missing_leading_dims = 1 - placement_dim = 1 (out) → source_dim = 0 - placement_dim = 2 (in) → source_dim = 1 - - 3) Packed QKV — source is one of Q/K/V (missing pack axis): - param shape = [3, 4096, 4096], source shape = [4096, 4096] - missing_leading_dims = 1 - placement_dim = 1 → source_dim = 0 - - 4) Two missing leading dims (stacked experts + packed QKV): - param shape = [8, 3, 4096, 4096], source shape = [4096, 4096] - missing_leading_dims = 2 - placement_dim = 3 → source_dim = 1 - """ - dim = self._norm_dim(placement_dim) - missing_leading_dims = self.param_ndim - len(source_shape) - if missing_leading_dims > 0 and dim >= missing_leading_dims: - dim -= missing_leading_dims - return dim - def dot_natural_key(s: str): """Sort key for state-dict names: split on ``"."`` and sort digits numerically diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py index 4f0e26b40b5e..f2bb07f3b95b 100644 --- a/tests/utils/test_core_model_loading.py +++ b/tests/utils/test_core_model_loading.py @@ -261,13 +261,23 @@ def __getitem__(self, name): def _make_dtensor_shard_op(mesh, placements, param_shape, local_shape): - """Build a DtensorShardOperation without requiring a real DTensor / distributed init.""" + """Build a DtensorShardOperation without requiring a real DTensor / distributed init. + + The expert-axis 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_shape = tuple(param_shape) op.param_ndim = len(param_shape) - op.local_shape = tuple(local_shape) + op._first_owned_expert = 0 + op._owned_experts_count = 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._first_owned_expert = sub.get_local_rank() * local_shape[0] + break return op @@ -1060,7 +1070,9 @@ def test_moe_expert_owned_with_inner_tp(self): } 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)) + op = _make_dtensor_shard_op( + mesh, [Shard(0), Shard(1)], param_shape=(4, 4, 2), local_shape=(2, 2, 2) + ) shard = op.shard_tensor(tensor, tensor_idx=1) if expected[rank] is None: self.assertIsNone(shard) From 4547eb33d0f66afc6680dc018575aa564d209866 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Mon, 4 May 2026 07:47:04 +0000 Subject: [PATCH 065/116] revert some stuff in core model loading --- src/transformers/core_model_loading.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 1416cced8d41..e10740663556 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -519,6 +519,8 @@ class WeightTransform: source_patterns: str | list[str] = field(init=True) target_patterns: str | list[str] = field(init=True) compiled_sources: re.Pattern = field(init=False) + _original_source_patterns: list[str] = field(init=False) + _original_target_patterns: list[str] = field(init=False) quantization_operation: ConversionOps | None = None @@ -541,10 +543,13 @@ def __post_init__(self): # when instantiating the reverse mapping (i.e. the targets become sources, and sources become targets) # The issues lie in the sources usually, so here we need to check the targets for the reversed mapping + # We need to copy the exact original patterns to later reverse (before processing may change them) + self._original_source_patterns = self.source_patterns.copy() + self._original_target_patterns = self.target_patterns.copy() + # Process target_patterns: detect capturing groups and replace with \1 # Store the original capturing group patterns for reverse mapping target_capturing_groups: list[str] = [] - unprocess_targets = self.target_patterns.copy() for i, pattern in enumerate(self.target_patterns): self.target_patterns[i], captured_group = process_target_pattern(pattern) if captured_group is not None: @@ -573,7 +578,7 @@ def __post_init__(self): pattern = pattern.replace(r"\1", unique_capturing_group, 1) # Potentially process a bit more for consistency - only if they are consistent pairs, i.e. the length is the same if len(self.source_patterns) == len(self.target_patterns): - pattern = process_source_pattern(pattern, unprocess_targets[i]) + pattern = process_source_pattern(pattern, self._original_target_patterns[i]) self.source_patterns[i] = pattern # Construct the regex we will use to rename keys from the sources to the targets @@ -627,7 +632,9 @@ def reverse_transform(self) -> WeightTransform: kwargs["operations"] = [op.reverse_op for op in self.operations[::-1]] reverse_transform = self.__class__( - source_patterns=self.target_patterns, target_patterns=self.source_patterns, **kwargs + source_patterns=self._original_target_patterns, + target_patterns=self._original_source_patterns, + **kwargs, ) return reverse_transform @@ -838,7 +845,7 @@ def shard_tensor( if not placements: return source[...].to(device=device, dtype=dtype) has_strided = any(not p.is_shard() for _, p in placements) - intervals: list[list[tuple[int, int]]] = [[(0, size)] for size in source_shape] + intervals = [[(0, size)] for size in source_shape] for mesh_dim, placement in placements: sub_mesh = self._get_sub_mesh(mesh_dim) rank, world_size = sub_mesh.get_local_rank(), sub_mesh.size() From 43086d33cae1598ddfdcb3b039447925ddcfc02e Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 5 May 2026 07:54:10 +0000 Subject: [PATCH 066/116] core model loading clean --- src/transformers/core_model_loading.py | 67 +++- tests/utils/test_core_model_loading.py | 474 +++++++++++-------------- 2 files changed, 265 insertions(+), 276 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index e10740663556..8675272a5802 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -823,30 +823,57 @@ def _job(): class DtensorShardOperation: - 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 - ) - self._first_owned_expert = offsets[0] - self._owned_experts_count = local_shape[0] + 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: - # Source already has the param's shape -> slice every dim directly. 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: + 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) @@ -862,22 +889,29 @@ def shard_tensor( slices = tuple(slice(*(pieces[0] if pieces else (0, 0))) for pieces in intervals) return source[slices].to(device=device, dtype=dtype) - # Source is one of N pieces. Ownership check on the leading axis; - # MergeModulelist / Concatenate will assemble the param from kept pieces. - # Inner dims only use _contiguous_intervals, which yields one piece per dim, - # so we can slice directly -- no cat needed. - shards_expert_axis = any(self._norm_dim(p.dim) == 0 for _, p in placements) - owns_expert = self._first_owned_expert <= tensor_idx < self._first_owned_expert + self._owned_experts_count - if shards_expert_axis and not owns_expert: + # 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() - source_dim = self._norm_dim(placement.dim) - 1 # exactly one leading dim missing + 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) @@ -954,6 +988,7 @@ def _get_sub_mesh(self, mesh_dim: int): 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 diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py index f2bb07f3b95b..605bf7badacf 100644 --- a/tests/utils/test_core_model_loading.py +++ b/tests/utils/test_core_model_loading.py @@ -263,7 +263,7 @@ def __getitem__(self, name): def _make_dtensor_shard_op(mesh, placements, param_shape, local_shape): """Build a DtensorShardOperation without requiring a real DTensor / distributed init. - The expert-axis ownership cache is computed by mimicking + 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. """ @@ -271,12 +271,12 @@ def _make_dtensor_shard_op(mesh, placements, param_shape, local_shape): op.device_mesh = mesh op.placements = tuple(placements) op.param_ndim = len(param_shape) - op._first_owned_expert = 0 - op._owned_experts_count = local_shape[0] + 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._first_owned_expert = sub.get_local_rank() * local_shape[0] + op._axis0_offset = sub.get_local_rank() * local_shape[0] break return op @@ -885,90 +885,97 @@ def test_ernie4_5_vl_moe_conversion_reversed(self): class TestDtensorShardOperation(unittest.TestCase): - """Unit tests for DtensorShardOperation.shard_tensor. - - Each test mirrors a real transformer-layer scenario that produces one - of the placement patterns shard_tensor must handle. The docstring of - every test walks the interval-narrowing loop step by step so the - expected output is traceable by hand. - - Scenarios Placement pattern - --------------------------------------------------------------------------------------------- - Row-parallel layer (o_proj, down_proj) [Shard(0), Shard(1)] - Column-parallel layer (q/k/v_proj, gate_up_proj, lm_head) [_StridedShard(0, sf=TP), Shard(0)] - _StridedShard alone (TP on input dim with group pattern) [Shard(0), _StridedShard(1, sf=2)] - MoE expert, not owned (mixtral experts) [Shard(0)] on expert dim - MoE expert, owned (mixtral experts) [Shard(0)] on expert dim - MoE expert + inner TP (experts + TP) [Shard(0), Shard(1)] - MoE TP without expert- (experts + TP, expert axis [Shard(1)] on inner dim - axis shard replicated) - Pre-pack half (one of gate_up halves) _StridedShard on missing packed axis - Replicate only (biases, norms) [Replicate()] - - Edge cases: - - uneven shard division (5 rows / 2 ranks) - - negative dim index normalization (Shard(-1)) - - Internal _slice_and_cat helper: - - fast path (one interval per dim) - - rejection of two multi-range dims + """Unit tests for DtensorShardOperation. + + 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. """ - # -------------------------------------------------------------- - # Row-parallel: FSDP and TP shard different dims (no collision) - # -------------------------------------------------------------- - def test_row_parallel_layer_shards_different_dims(self): - """Row-parallel (o_proj / down_proj) on a 2×2 mesh [FSDP, TP]. + 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}") - param = Linear.weight, shape (out=8, in=8). - placements = [Shard(0), Shard(1)] — FSDP on output rows, TP on input cols. + 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}") - Walk for rank (FSDP=0, TP=0): - init → dim 0: [(0, 8)] dim 1: [(0, 8)] - after Shard(0) → dim 0: [(0, 4)] dim 1: [(0, 8)] - after Shard(1) → dim 0: [(0, 4)] dim 1: [(0, 4)] - → source[0:4, 0:4] + 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}") - Each of the 4 ranks owns a disjoint 4×4 quadrant. - """ + def test_2D_shard_different_dims(self): tensor = torch.arange(64).reshape(8, 8).float() expected = { - 0: tensor[:4, :4], # (FSDP=0, TP=0) top-left - 1: tensor[:4, 4:], # (FSDP=0, TP=1) top-right - 2: tensor[4:, :4], # (FSDP=1, TP=0) bottom-left - 3: tensor[4:, 4:], # (FSDP=1, TP=1) bottom-right + 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}") - # -------------------------------------------------------------- - # Column-parallel: FSDP + TP both shard dim 0 (stride resolves) - # -------------------------------------------------------------- - def test_column_parallel_layer_same_dim_collision(self): - """Column-parallel (q/k/v_proj, gate_up_proj, lm_head) on a 2×2 mesh [FSDP, TP]. - - param = Linear.weight, shape (out=4, in=4). - placements = [_StridedShard(0, sf=2), Shard(0)] — both on dim 0. - - Walk for rank (FSDP=0, TP=0): - init → dim 0: [(0, 4)] dim 1: [(0, 4)] - after _StridedShard(0, 2) → dim 0: [(0, 1), (2, 3)] - # groups (0,2) and (2,4); FSDP 0 keeps first half of each → rows {0, 2} - after Shard(0) → dim 0: [(0, 1)] - # view [(0,1),(2,3)] flat → {row 0, row 2}; TP 0 takes first half → row 0 - → source[0:1, :] - - Gather across FSDP: TP 0 sees rows {0,1}, TP 1 sees rows {2,3} — contiguous - chunks as column-parallel kernels require. - """ + 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]], # (FSDP=0, TP=0) - 1: tensor[[2]], # (FSDP=0, TP=1) - 2: tensor[[1]], # (FSDP=1, TP=0) - 3: tensor[[3]], # (FSDP=1, TP=1) + 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) @@ -980,151 +987,58 @@ def test_column_parallel_layer_same_dim_collision(self): ) torch.testing.assert_close(op.shard_tensor(tensor), expected[rank], msg=f"rank {rank}") - # -------------------------------------------------------------- - # _StridedShard alone (on its own dim): multi-interval + concat - # -------------------------------------------------------------- - def test_strided_shard_alone_produces_disjoint_intervals(self): - """_StridedShard on its own dim (different from Shard's dim) yields multi-interval reads. - - param shape (8, 8), placements = [Shard(0), _StridedShard(1, sf=2)]. - - Walk for rank (0, 0): - init → dim 0: [(0, 8)] dim 1: [(0, 8)] - after Shard(0) → dim 0: [(0, 4)] dim 1: [(0, 8)] - after _StridedShard(1, sf=2) → dim 0: [(0, 4)] dim 1: [(0, 2), (4, 6)] - # groups (0,4) and (4,8); rank 0 keeps first half of each → cols {0-1, 4-5} - → _slice_and_cat reads source[:4, 0:2] and source[:4, 4:6], - concatenates along dim 1. - """ - tensor = torch.arange(64).reshape(8, 8).float() + def test_2D_strided_shard_different_dims(self): + tensor = torch.arange(16).reshape(4, 4).float() expected = { - 0: torch.cat([tensor[:4, :2], tensor[:4, 4:6]], dim=1), - 1: torch.cat([tensor[:4, 2:4], tensor[:4, 6:8]], dim=1), - 2: torch.cat([tensor[4:, :2], tensor[4:, 4:6]], dim=1), - 3: torch.cat([tensor[4:, 2:4], tensor[4:, 6:8]], dim=1), + 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=(8, 8), - local_shape=(4, 4), + param_shape=(4, 4), + local_shape=(2, 2), ) torch.testing.assert_close(op.shard_tensor(tensor), expected[rank], msg=f"rank {rank}") - # -------------------------------------------------------------- - # MoE experts (case b): source.ndim == param.ndim - 1 - # -------------------------------------------------------------- - def test_moe_expert_not_owned_returns_none(self): - """Expert file belongs to another rank → return None (skip the file). - - param (E=4, H=2, I=2), placements = [Shard(0)] on expert axis. - 2-rank mesh (FSDP=2). Rank 1 owns experts {2, 3} (offset=2, size=2). - Loading expert_idx=0 on rank 1 → the file is for rank 0, skip. - """ - mesh = FakeMesh(shape=(2,), rank=1) - op = _make_dtensor_shard_op(mesh, [Shard(0)], param_shape=(4, 2, 2), local_shape=(2, 2, 2)) - expert_tensor = torch.ones(2, 2) - self.assertIsNone(op.shard_tensor(expert_tensor, tensor_idx=0)) - - def test_moe_expert_owned_without_inner_sharding(self): - """Expert owned and no inner sharding → keep the whole expert tensor. - - Same param / placements / mesh as above. Loading expert_idx=2 on rank 1: - source_is_one_expert = True - Shard(0) dim normalizes to 0 → enters the expert branch. - _owns_expert(2) = True → drop Shard(0) from placements. - Remaining placements = [] → early return of the full source tensor. - """ - mesh = FakeMesh(shape=(2,), rank=1) - op = _make_dtensor_shard_op(mesh, [Shard(0)], param_shape=(4, 2, 2), local_shape=(2, 2, 2)) - expert_tensor = torch.ones(2, 2) - torch.testing.assert_close(op.shard_tensor(expert_tensor, tensor_idx=2), expert_tensor) - - def test_moe_expert_owned_with_inner_tp(self): - """MoE expert sharded on expert axis (FSDP) and inner dim (TP). - - param (E=4, H=4, I=2), placements = [Shard(0), Shard(1)]: - Shard(0) on 2-way FSDP → expert-axis shard - Shard(1) on 2-way TP → inner hidden-dim shard - source = (H=4, I=2) for one expert, tensor_idx=1. - - Walk per rank (FSDP, TP) with tensor_idx=1: - expert branch: rank owns expert 1 only if FSDP==0 (offset 0, size 2) - FSDP=1 ranks → return None - FSDP=0 ranks → drop Shard(0); continue with Shard(1) - remaining placements = [(1, Shard(1))] - source_dim = _source_dim(1, [4, 2]) = 1 - missing_leading(1) = 0 - init → dim 0: [(0, 4)] dim 1: [(0, 2)] - after Shard(1)→src 0: - TP=0 rank → dim 0: [(0, 2)] → source[:2] - TP=1 rank → dim 0: [(2, 4)] → source[2:] - """ - tensor = torch.arange(8).reshape(4, 2).float() + def test_moe_1D_shard_filters_by_axis0_ownership(self): + source = torch.ones(2, 2) expected = { - 0: tensor[:2], # (FSDP=0, TP=0) — owns expert 1, inner rows 0-1 - 1: tensor[2:], # (FSDP=0, TP=1) — owns expert 1, inner rows 2-3 - 2: None, # (FSDP=1, TP=0) — does not own expert 1 - 3: None, # (FSDP=1, TP=1) — does not own expert 1 + 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(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(tensor, tensor_idx=1) - if expected[rank] is None: - self.assertIsNone(shard) - else: - torch.testing.assert_close(shard, expected[rank], msg=f"rank {rank}") - - def test_moe_expert_tp_only_no_expert_axis_shard(self): - """MoE param with TP on inner dim but no expert-axis sharding. - - param (E=4, H=4, I=2), placements = [Shard(1)] — TP only, 2-rank mesh. - source = (H=4, I=2), tensor_idx=0. - - source_is_one_expert = True, BUT Shard(1).dim normalizes to 1 ≠ 0, - so no placement targets the expert axis → skip the expert branch and - fall into the generic loop with missing_leading_dims=1. - - Walk: - source_dim = _source_dim(1, [4, 2]) = 1 - 1 = 0 - init → dim 0: [(0, 4)] dim 1: [(0, 2)] - after Shard(1)→src 0: - rank 0 → dim 0: [(0, 2)] → source[:2] - rank 1 → dim 0: [(2, 4)] → source[2:] - """ - tensor = torch.arange(8).reshape(4, 2).float() - for rank, expected in [(0, tensor[:2]), (1, tensor[2:])]: + for rank in range(2): mesh = FakeMesh(shape=(2,), rank=rank) - op = _make_dtensor_shard_op(mesh, [Shard(1)], param_shape=(4, 4, 2), local_shape=(4, 2, 2)) - torch.testing.assert_close(op.shard_tensor(tensor, tensor_idx=0), expected, msg=f"rank {rank}") - - # -------------------------------------------------------------- - # Pre-pack half (case c): _StridedShard on missing packed axis - # -------------------------------------------------------------- - def test_prepack_half_strided_degrades_to_contiguous(self): - """Pre-concat w1 / w3 tensor: _StridedShard on the packed axis degrades. - - param = packed gate_up per-expert, shape (E=8, 2H=8, D=2). - source = (H=4, D=2) — single w1 half for one expert, tensor_idx=0. - placements = [_StridedShard(dim=1, sf=2)] — would stride the packed 2H dim. - - source_missing_leading_axis = True (source ndim 2 < param ndim 3). - is_interleaved requires same ndim → False. - → _StridedShard falls through to _contiguous_intervals. - The packing is rebuilt later by the WeightConverter's Concatenate op. - - Walk for rank 0 (2-rank mesh): - source_dim = _source_dim(1, [4, 2]) = 1 - 1 = 0 - init → dim 0: [(0, 4)] dim 1: [(0, 2)] - after _StridedShard→src 0 → dim 0: [(0, 2)] dim 1: [(0, 2)] - → source[:2] - """ - tensor = torch.arange(8).reshape(4, 2).float() - for rank, expected in [(0, tensor[:2]), (1, tensor[2:])]: + 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, @@ -1132,69 +1046,109 @@ def test_prepack_half_strided_degrades_to_contiguous(self): param_shape=(8, 8, 2), local_shape=(8, 4, 2), ) - torch.testing.assert_close(op.shard_tensor(tensor, tensor_idx=0), expected, msg=f"rank {rank}") - - # -------------------------------------------------------------- - # Replicate only (biases, norms): no narrowing - # -------------------------------------------------------------- - def test_replicate_only_returns_full_tensor(self): - """All placements are Replicate → placements filter drops everything - → early return with a full-tensor copy.""" - mesh = FakeMesh(shape=(2,), rank=0) - op = _make_dtensor_shard_op(mesh, [Replicate()], param_shape=(4, 4), local_shape=(4, 4)) - tensor = torch.arange(16).reshape(4, 4).float() - torch.testing.assert_close(op.shard_tensor(tensor), tensor) + torch.testing.assert_close(op.shard_tensor(source, tensor_idx=0), expected[rank], msg=f"rank {rank}") - # -------------------------------------------------------------- - # Edge cases for _contiguous_intervals - # -------------------------------------------------------------- - def test_contiguous_shard_uneven_division(self): - """Size-5 axis sharded across 2 ranks: rank 0 gets 3 rows, rank 1 gets 2. + 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}") - _contiguous_intervals calls Shard.local_shard_size_and_offset which - rounds up the per-rank share; the last rank takes whatever remains. - """ - tensor = torch.arange(20).reshape(5, 4).float() - expected = {0: tensor[:3], 1: tensor[3:]} - for rank in range(2): - mesh = FakeMesh(shape=(2,), rank=rank) - local_rows = 3 if rank == 0 else 2 - op = _make_dtensor_shard_op(mesh, [Shard(0)], param_shape=(5, 4), local_shape=(local_rows, 4)) - torch.testing.assert_close(op.shard_tensor(tensor), 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_negative_dim_normalization(self): - """Shard(-1) on a 2D tensor shards the last dim (dim 1). + 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) - _norm_dim(-1) with param_ndim=2 → 2 + (-1) = 1. - """ - tensor = torch.arange(16).reshape(4, 4).float() - for rank, expected in [(0, tensor[:, :2]), (1, tensor[:, 2:])]: - mesh = FakeMesh(shape=(2,), rank=rank) - op = _make_dtensor_shard_op(mesh, [Shard(-1)], param_shape=(4, 4), local_shape=(4, 2)) - torch.testing.assert_close(op.shard_tensor(tensor), expected, msg=f"rank {rank}") - - # -------------------------------------------------------------- - # Internal helper: _slice_and_cat - # -------------------------------------------------------------- - def test_slice_and_cat_fast_path_single_interval_per_dim(self): - """Every dim has exactly one interval → fast path: single slice read, no concat.""" - tensor = torch.arange(64).reshape(8, 8).float() + 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, 8), local_shape=(4, 4)) - intervals = [[(0, 4)], [(2, 6)]] - result = op._slice_and_cat(tensor, intervals, None, None) - torch.testing.assert_close(result, tensor[0:4, 2:6]) - - def test_slice_and_cat_rejects_two_multi_interval_dims(self): - """Two dims with multiple disjoint ranges would require a 2D outer-product - of reads. Not supported → ValueError. - """ + 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, 4)) - intervals = [[(0, 2), (4, 6)], [(0, 2), (4, 6)]] + 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, intervals, None, None) + 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 TestConversionMapping(unittest.TestCase): From 1b7ebe1b006c0e7d29fdb9adea0bcf5171b6b1b9 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Sun, 10 May 2026 22:28:09 +0000 Subject: [PATCH 067/116] guarding import --- src/transformers/distributed/configuration_utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/transformers/distributed/configuration_utils.py b/src/transformers/distributed/configuration_utils.py index 0f3abaccc352..9d94050909bb 100644 --- a/src/transformers/distributed/configuration_utils.py +++ b/src/transformers/distributed/configuration_utils.py @@ -16,7 +16,9 @@ import os from dataclasses import asdict, dataclass -import torch +from ..utils import is_torch_available +if is_torch_available(): + import torch @dataclass From 6d867469264892b38c557360ac17eb6a0344681a Mon Sep 17 00:00:00 2001 From: 3outeille Date: Sun, 10 May 2026 22:38:21 +0000 Subject: [PATCH 068/116] better separation tensor parall and generic utils --- run_compare.sh | 17 ++- .../distributed/configuration_utils.py | 2 + src/transformers/distributed/utils.py | 133 +++++++++++++++++- src/transformers/integrations/__init__.py | 6 - .../integrations/tensor_parallel.py | 129 ----------------- src/transformers/modeling_utils.py | 3 +- tests/test_tensor_parallel_mixin.py | 3 +- tmp_generate.py | 2 +- train_fsdp_tp.py | 8 +- 9 files changed, 156 insertions(+), 147 deletions(-) diff --git a/run_compare.sh b/run_compare.sh index 75e06e8dae96..52200d431676 100755 --- a/run_compare.sh +++ b/run_compare.sh @@ -4,6 +4,7 @@ set -euo pipefail SCRIPT="train_fsdp_tp.py" LOG_FSDP_TP="log.txt" LOG_FSDP_ONLY="ref.txt" +LOG_DIFF="diff.txt" MODEL_NAME="${MODEL_NAME:-hf-internal-testing/tiny-random-MixtralForCausalLM}" COMMON_ARGS="--model_name $MODEL_NAME --lr 3e-4 --seed 42" @@ -47,10 +48,18 @@ echo "FSDP+TP PID=$PID1 | FSDP-only PID=$PID2" wait $PID1 && echo "Phase 2 FSDP+TP done" || { echo "Phase 2 FSDP+TP failed (exit $?)"; cat "${LOG_FSDP_TP}.phase2"; exit 1; } wait $PID2 && echo "Phase 2 FSDP-only done" || { echo "Phase 2 FSDP-only failed (exit $?)"; cat "${LOG_FSDP_ONLY}.phase2"; exit 1; } -# Combine phase logs -cat "${LOG_FSDP_TP}.phase1" "${LOG_FSDP_TP}.phase2" > "$LOG_FSDP_TP" -cat "${LOG_FSDP_ONLY}.phase1" "${LOG_FSDP_ONLY}.phase2" > "$LOG_FSDP_ONLY" +# Combine phase logs, keeping only signal lines (loss/grad steps + checkpoint markers). +# Drops every kind of warning/progress noise: rank warnings, torchrun banners, +# transformers deprecations, tqdm progress bars, ProcessGroup teardown warnings, etc. +strip_warnings() { + grep -E '^(Step |Resumed |Saved )' +} +strip_warnings < "${LOG_FSDP_TP}.phase1" > "$LOG_FSDP_TP" +strip_warnings < "${LOG_FSDP_TP}.phase2" >> "$LOG_FSDP_TP" +strip_warnings < "${LOG_FSDP_ONLY}.phase1" > "$LOG_FSDP_ONLY" +strip_warnings < "${LOG_FSDP_ONLY}.phase2" >> "$LOG_FSDP_ONLY" echo "" echo "=== Full Loss & Grad Diff (steps 0-19) ===" -git diff --no-index --color --word-diff=color "$LOG_FSDP_TP" "$LOG_FSDP_ONLY" || true \ No newline at end of file +git diff --no-index --color --word-diff=color "$LOG_FSDP_TP" "$LOG_FSDP_ONLY" | tee "$LOG_DIFF" || true +echo "Diff written to $LOG_DIFF" \ No newline at end of file diff --git a/src/transformers/distributed/configuration_utils.py b/src/transformers/distributed/configuration_utils.py index 9d94050909bb..d3f76b72442e 100644 --- a/src/transformers/distributed/configuration_utils.py +++ b/src/transformers/distributed/configuration_utils.py @@ -17,6 +17,8 @@ from dataclasses import asdict, dataclass from ..utils import is_torch_available + + if is_torch_available(): import torch diff --git a/src/transformers/distributed/utils.py b/src/transformers/distributed/utils.py index fa4d554f130d..5d2f0d3b0a54 100644 --- a/src/transformers/distributed/utils.py +++ b/src/transformers/distributed/utils.py @@ -27,8 +27,8 @@ if is_torch_available(): import torch import torch.distributed.checkpoint as dcp - - from ..integrations.tensor_parallel import convert_strided_to_shard, restore_strided_from_shard + from torch.distributed.tensor import DTensor, Replicate, Shard + from torch.distributed.tensor.placement_types import _StridedShard def is_fsdp_enabled() -> bool: @@ -124,6 +124,135 @@ def init_device_mesh(distributed_config: DistributedConfig) -> torch.distributed return mesh +def _to_cpu_fresh(tensor: torch.Tensor) -> torch.Tensor: + """Plain tensor → contiguous CPU tensor with fresh storage for safetensors.""" + if tensor.device.type == "meta": + return tensor + t = tensor.detach() + if t.device.type != "cpu": + t = t.to(device="cpu") + out = torch.empty(t.shape, dtype=t.dtype, device="cpu") + out.copy_(t) + return out.contiguous() + + +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 ``{}``. + """ + tp_size = model.tp_size + is_rank0 = torch.distributed.get_rank() == 0 + + # Get state dict — FSDP unshard if needed (returns DTensors, not full tensors) + if getattr(model, "_is_fsdp_managed_module", False): + from torch.distributed.checkpoint.state_dict import get_model_state_dict + + state_dict = get_model_state_dict(model) + else: + state_dict = model.state_dict() + + # No TP — materialize on rank 0 only + if tp_size is None: + if is_rank0: + return {k: _to_cpu_fresh(v) for k, v in state_dict.items()} + return {} + + # Stream: gather one param at a time, only rank 0 keeps the CPU copy + result = {} + for key, tensor in state_dict.items(): + if isinstance(tensor, DTensor): + # All ranks participate in the collective, only rank 0 keeps the result + with torch.no_grad(): + full = _replicate_dtensor(tensor).to_local() + if is_rank0: + result[key] = _to_cpu_fresh(full) + del full + elif is_rank0: + result[key] = _to_cpu_fresh(tensor) + + return result + + +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 + + So we bypass ``redistribute`` and call each placement's low-level + ``_to_replicate_tensor`` (manual all-gather + interleaved reorder). + + We process mesh dims **right-to-left** (innermost first). Under TP+FSDP + the 2D mesh is ``(fsdp, tp)`` and both dims can shard the same tensor dim:: + + placements = (_StridedShard(dim=0), Shard(dim=0)) + local shape = [64, 1024] (global [256, 1024], fsdp=2, tp=2) + + Right-to-left means TP is gathered first (local grows to [128, 1024]), + then FSDP (grows to [256, 1024]). Each step must pass the correct + intermediate logical shape — the global shape divided by the mesh sizes + of dims not yet gathered (to the left). + """ + mesh = tensor.device_mesh + replicate_all = tuple(Replicate() for _ in range(mesh.ndim)) + with torch.no_grad(): + if any(isinstance(p, _StridedShard) for p in tensor.placements): + local = tensor._local_tensor + placements = tensor.placements + for i in reversed(range(mesh.ndim)): + p = placements[i] + if p.is_replicate(): + continue + # Compute the logical shape seen at this step: dims to the left + # (not yet gathered) still divide their tensor dimension. + logical_shape = list(tensor.shape) + for j in range(i): + pj = placements[j] + if not pj.is_replicate(): + logical_shape[pj.dim] //= mesh.size(j) + local = p._to_replicate_tensor(local, mesh, i, logical_shape) + return DTensor.from_local(local, mesh, replicate_all, run_check=False) + + return tensor.redistribute(placements=replicate_all) + + +def convert_strided_to_shard(state_dict: dict) -> dict[str, tuple]: + # Convert _StridedShard DTensors in a state dict to plain Shard for DCP compatibility. + placement_map: dict[str, tuple] = {} + for key, value in state_dict.items(): + if isinstance(value, dict): + nested = convert_strided_to_shard(value) + for nk, nv in nested.items(): + placement_map[f"{key}.{nk}"] = nv + elif isinstance(value, DTensor) and any(isinstance(p, _StridedShard) for p in value.placements): + placement_map[key] = tuple(value.placements) + shard_placements = tuple(Shard(p.dim) if isinstance(p, _StridedShard) else p for p in value.placements) + state_dict[key] = _replicate_dtensor(value).redistribute(placements=shard_placements) + return placement_map + + +def restore_strided_from_shard(state_dict: dict, placement_map: dict[str, tuple]) -> None: + # Restore _StridedShard placements after dcp.load. + def _resolve(d, dotted_key): + parts = dotted_key.split(".", 1) + if len(parts) == 2 and parts[0] in d and isinstance(d[parts[0]], dict): + return _resolve(d[parts[0]], parts[1]) + return d, dotted_key + + for key, original_placements in placement_map.items(): + container, leaf_key = _resolve(state_dict, key) + if leaf_key in container and isinstance(container[leaf_key], DTensor): + container[leaf_key] = _replicate_dtensor(container[leaf_key]).redistribute(placements=original_placements) + + def save_optimizer(optimizer, checkpoint_dir: str) -> None: # Save optimizer state via DCP, handling _StridedShard placements transparently. osd = optimizer.state_dict() diff --git a/src/transformers/integrations/__init__.py b/src/transformers/integrations/__init__.py index 4be557be670a..4d4e43958f3c 100755 --- a/src/transformers/integrations/__init__.py +++ b/src/transformers/integrations/__init__.py @@ -163,9 +163,6 @@ _import_structure["tensor_parallel"] = [ "ALL_PARALLEL_STYLES", "apply_tensor_parallel", - "convert_strided_to_shard", - "gather_full_state_dict", - "restore_strided_from_shard", "verify_tp_plan", ] try: @@ -301,9 +298,6 @@ from .tensor_parallel import ( ALL_PARALLEL_STYLES, apply_tensor_parallel, - convert_strided_to_shard, - gather_full_state_dict, - restore_strided_from_shard, verify_tp_plan, ) from .vptq import replace_with_vptq_linear diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index 2a82a4e55067..457d0678d9e0 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -79,135 +79,6 @@ def _get_parameter_tp_plan(parameter_name: str, tp_plan: dict[str, str], is_weig # ============================================================================= -def _to_cpu_fresh(tensor: torch.Tensor) -> torch.Tensor: - """Plain tensor → contiguous CPU tensor with fresh storage for safetensors.""" - if tensor.device.type == "meta": - return tensor - t = tensor.detach() - if t.device.type != "cpu": - t = t.to(device="cpu") - out = torch.empty(t.shape, dtype=t.dtype, device="cpu") - out.copy_(t) - return out.contiguous() - - -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 ``{}``. - """ - tp_size = model.tp_size - is_rank0 = dist.get_rank() == 0 - - # Get state dict — FSDP unshard if needed (returns DTensors, not full tensors) - if getattr(model, "_is_fsdp_managed_module", False): - from torch.distributed.checkpoint.state_dict import get_model_state_dict - - state_dict = get_model_state_dict(model) - else: - state_dict = model.state_dict() - - # No TP — materialize on rank 0 only - if tp_size is None: - if is_rank0: - return {k: _to_cpu_fresh(v) for k, v in state_dict.items()} - return {} - - # Stream: gather one param at a time, only rank 0 keeps the CPU copy - result = {} - for key, tensor in state_dict.items(): - if isinstance(tensor, DTensor): - # All ranks participate in the collective, only rank 0 keeps the result - with torch.no_grad(): - full = _replicate_dtensor(tensor).to_local() - if is_rank0: - result[key] = _to_cpu_fresh(full) - del full - elif is_rank0: - result[key] = _to_cpu_fresh(tensor) - - return result - - -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 - - So we bypass ``redistribute`` and call each placement's low-level - ``_to_replicate_tensor`` (manual all-gather + interleaved reorder). - - We process mesh dims **right-to-left** (innermost first). Under TP+FSDP - the 2D mesh is ``(fsdp, tp)`` and both dims can shard the same tensor dim:: - - placements = (_StridedShard(dim=0), Shard(dim=0)) - local shape = [64, 1024] (global [256, 1024], fsdp=2, tp=2) - - Right-to-left means TP is gathered first (local grows to [128, 1024]), - then FSDP (grows to [256, 1024]). Each step must pass the correct - intermediate logical shape — the global shape divided by the mesh sizes - of dims not yet gathered (to the left). - """ - mesh = tensor.device_mesh - replicate_all = tuple(Replicate() for _ in range(mesh.ndim)) - with torch.no_grad(): - if any(isinstance(p, _StridedShard) for p in tensor.placements): - local = tensor._local_tensor - placements = tensor.placements - for i in reversed(range(mesh.ndim)): - p = placements[i] - if p.is_replicate(): - continue - # Compute the logical shape seen at this step: dims to the left - # (not yet gathered) still divide their tensor dimension. - logical_shape = list(tensor.shape) - for j in range(i): - pj = placements[j] - if not pj.is_replicate(): - logical_shape[pj.dim] //= mesh.size(j) - local = p._to_replicate_tensor(local, mesh, i, logical_shape) - return DTensor.from_local(local, mesh, replicate_all, run_check=False) - - return tensor.redistribute(placements=replicate_all) - - -def convert_strided_to_shard(state_dict: dict) -> dict[str, tuple]: - # Convert _StridedShard DTensors in a state dict to plain Shard for DCP compatibility. - placement_map: dict[str, tuple] = {} - for key, value in state_dict.items(): - if isinstance(value, dict): - nested = convert_strided_to_shard(value) - for nk, nv in nested.items(): - placement_map[f"{key}.{nk}"] = nv - elif isinstance(value, DTensor) and any(isinstance(p, _StridedShard) for p in value.placements): - placement_map[key] = tuple(value.placements) - shard_placements = tuple(Shard(p.dim) if isinstance(p, _StridedShard) else p for p in value.placements) - state_dict[key] = _replicate_dtensor(value).redistribute(placements=shard_placements) - return placement_map - - -def restore_strided_from_shard(state_dict: dict, placement_map: dict[str, tuple]) -> None: - # Restore _StridedShard placements after dcp.load. - def _resolve(d, dotted_key): - parts = dotted_key.split(".", 1) - if len(parts) == 2 and parts[0] in d and isinstance(d[parts[0]], dict): - return _resolve(d[parts[0]], parts[1]) - return d, dotted_key - - for key, original_placements in placement_map.items(): - container, leaf_key = _resolve(state_dict, key) - if leaf_key in container and isinstance(container[leaf_key], DTensor): - container[leaf_key] = _replicate_dtensor(container[leaf_key]).redistribute(placements=original_placements) - - def verify_tp_plan(expected_keys: list[str], tp_plan: dict[str, str] | None): """ Verify the TP plan of the model, log a warning if the layers that were not sharded and the rules that were not applied. diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 2096d98d7275..420fa01a541d 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -52,7 +52,7 @@ revert_weight_conversion, ) from .distributed import DistributedConfig -from .distributed.utils import init_device_mesh, is_fsdp_enabled +from .distributed.utils import gather_full_state_dict, init_device_mesh, is_fsdp_enabled 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 @@ -78,7 +78,6 @@ from .integrations.tensor_parallel import ( _get_parameter_tp_plan, apply_tensor_parallel, - gather_full_state_dict, verify_tp_plan, ) from .loss.loss_utils import LOSS_MAPPING diff --git a/tests/test_tensor_parallel_mixin.py b/tests/test_tensor_parallel_mixin.py index 5d102d2ff0ec..9c5ab3fba9af 100644 --- a/tests/test_tensor_parallel_mixin.py +++ b/tests/test_tensor_parallel_mixin.py @@ -17,7 +17,8 @@ from transformers import TorchAoConfig, set_seed from transformers.distributed import DistributedConfig -from transformers.integrations.tensor_parallel import _get_parameter_tp_plan, _replicate_dtensor +from transformers.distributed.utils import _replicate_dtensor +from transformers.integrations.tensor_parallel import _get_parameter_tp_plan from transformers.testing_utils import ( is_tensor_parallel_test, is_torch_available, diff --git a/tmp_generate.py b/tmp_generate.py index 9685bed643ed..8f2118771406 100644 --- a/tmp_generate.py +++ b/tmp_generate.py @@ -22,7 +22,7 @@ @record def main(args): - distributed_config = DistributedConfig(tp_size=4, tp_plan="auto") + distributed_config = DistributedConfig(tp_size=8, tp_plan="auto") model = AutoModelForCausalLM.from_pretrained(model_id, distributed_config=distributed_config, dtype=torch.bfloat16) # model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.bfloat16, device_map="auto") tokenizer = AutoTokenizer.from_pretrained(model_id) diff --git a/train_fsdp_tp.py b/train_fsdp_tp.py index ab0737a9d0b6..96c52b2498b2 100644 --- a/train_fsdp_tp.py +++ b/train_fsdp_tp.py @@ -10,6 +10,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer from transformers.distributed import DistributedConfig from transformers.distributed.utils import load_optimizer, save_optimizer +from transformers.integrations.tensor_parallel import _replicate_dtensor def build_packed_dataset(dataset_name, tokenizer, seq_len, dp_rank, dp_world_size): """Stream + tokenize + greedy-pack documents into fixed-length (input, label) windows.""" @@ -104,9 +105,12 @@ def build_fixed_batches(dp_rank): loss = model(input_ids, labels=labels).loss loss.backward() - # Custom grad clip: convert DTensor grads to local to avoid mixed-mesh torch.stack + # Custom grad clip: convert DTensor grads to local to avoid mixed-mesh torch.stack. + # Use _replicate_dtensor (not full_tensor) — full_tensor calls redistribute(), which + # cannot normalize FSDP+TP placements like (Shard(0), _StridedShard(1, sf=2)) from + # packed_colwise MoE experts. grads = [p.grad for p in model.parameters() if p.grad is not None] - local_grads = [g.full_tensor() if isinstance(g, DTensor) else g for g in grads] + 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=2.0) torch.nn.utils.clip_grads_with_norm_(grads, max_norm=1.0, total_norm=total_norm) optimizer.step() From ff493468fde1e2f4b56f09fedbfd7a13acdbf685 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Sun, 10 May 2026 22:57:51 +0000 Subject: [PATCH 069/116] isolate DtensorShardOperation into a separate file --- src/transformers/core_model_loading.py | 174 +-------- src/transformers/distributed/model_loading.py | 199 ++++++++++ tests/utils/test_core_model_loading.py | 335 +---------------- tests/utils/test_distributed_model_loading.py | 354 ++++++++++++++++++ train_fsdp_tp.py | 3 +- 5 files changed, 557 insertions(+), 508 deletions(-) create mode 100644 src/transformers/distributed/model_loading.py create mode 100644 tests/utils/test_distributed_model_loading.py diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 8675272a5802..49496facac25 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -31,6 +31,7 @@ import torch +from .distributed.model_loading import DtensorShardOperation from .integrations.accelerate import get_device, offload_weight from .utils import is_env_variable_true from .utils.loading_report import LoadStateDictInfo @@ -44,9 +45,6 @@ from .quantizers import HfQuantizer elif _torch_distributed_available: from torch.distributed.tensor import DTensor - from torch.distributed.tensor._utils import compute_local_shape_and_global_offset - from torch.distributed.tensor.placement_types import Shard - logger = get_logger(__name__) @@ -822,176 +820,6 @@ def _job(): return _job -class DtensorShardOperation: - 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 dot_natural_key(s: str): """Sort key for state-dict names: split on ``"."`` and sort digits numerically and strings alphabetically. We emit a tuple at each point to sort ints diff --git a/src/transformers/distributed/model_loading.py b/src/transformers/distributed/model_loading.py new file mode 100644 index 000000000000..c91fe6df5acc --- /dev/null +++ b/src/transformers/distributed/model_loading.py @@ -0,0 +1,199 @@ +# 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.tensor._utils import compute_local_shape_and_global_offset + from torch.distributed.tensor.placement_types import Shard + + +class DtensorShardOperation: + 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 diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py index 605bf7badacf..b0156b47c5c4 100644 --- a/tests/utils/test_core_model_loading.py +++ b/tests/utils/test_core_model_loading.py @@ -16,14 +16,13 @@ import torch import torch.nn as nn -from torch.distributed.tensor.placement_types import Replicate, Shard, _StridedShard +from torch.distributed.tensor.placement_types import Shard from transformers import PretrainedConfig from transformers.conversion_mapping import get_checkpoint_conversion_mapping, register_checkpoint_conversion_mapping from transformers.core_model_loading import ( Chunk, Concatenate, - DtensorShardOperation, ErnieFuseAndSplitTextVisionExperts, MergeModulelist, PermuteForRope, @@ -39,6 +38,7 @@ from transformers.utils.import_utils import is_triton_available from ..test_modeling_common import compare_state_dicts +from .test_distributed_model_loading import FakeMesh, _make_dtensor_shard_op class TestWeightGlobMatching(unittest.TestCase): @@ -217,70 +217,6 @@ def __init__(self, add_extra_moe=False): self.mlp = DummyMLP() -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 TestConvertAndLoadStateDict(unittest.TestCase): def test_dtensor_shard_aware_mixtral_conversion_uses_only_local_experts(self): """Integration test: FSDP-sharded expert loading + WeightConverter. @@ -884,273 +820,6 @@ def test_ernie4_5_vl_moe_conversion_reversed(self): self.assertTrue(compare_state_dicts(reversed_state_dict, state_dict)) -class TestDtensorShardOperation(unittest.TestCase): - """Unit tests for DtensorShardOperation. - - 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. - """ - - 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 TestConversionMapping(unittest.TestCase): def test_register_checkpoint_conversion_mapping(self): register_checkpoint_conversion_mapping( diff --git a/tests/utils/test_distributed_model_loading.py b/tests/utils/test_distributed_model_loading.py new file mode 100644 index 000000000000..2dfdb026a8cb --- /dev/null +++ b/tests/utils/test_distributed_model_loading.py @@ -0,0 +1,354 @@ +# 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 unittest + +import torch +from torch.distributed.tensor.placement_types import Replicate, Shard, _StridedShard + +from transformers.distributed.model_loading import DtensorShardOperation + + +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. + + 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. + """ + + 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) + + +if __name__ == "__main__": + unittest.main() diff --git a/train_fsdp_tp.py b/train_fsdp_tp.py index 96c52b2498b2..f4c3f84ebe4b 100644 --- a/train_fsdp_tp.py +++ b/train_fsdp_tp.py @@ -9,8 +9,7 @@ from torch.utils.data import DataLoader from transformers import AutoModelForCausalLM, AutoTokenizer from transformers.distributed import DistributedConfig -from transformers.distributed.utils import load_optimizer, save_optimizer -from transformers.integrations.tensor_parallel import _replicate_dtensor +from transformers.distributed.utils import load_optimizer, save_optimizer, _replicate_dtensor def build_packed_dataset(dataset_name, tokenizer, seq_len, dp_rank, dp_world_size): """Stream + tokenize + greedy-pack documents into fixed-length (input, label) windows.""" From a806b3db2353b0e5c5d9e2b9a42f5179e7ea63a3 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Sun, 10 May 2026 23:48:09 +0000 Subject: [PATCH 070/116] no need to patch rotary --- src/transformers/distributed/patches.py | 80 ------------------- .../integrations/tensor_parallel.py | 6 -- 2 files changed, 86 deletions(-) delete mode 100644 src/transformers/distributed/patches.py diff --git a/src/transformers/distributed/patches.py b/src/transformers/distributed/patches.py deleted file mode 100644 index 3ffe15a37afe..000000000000 --- a/src/transformers/distributed/patches.py +++ /dev/null @@ -1,80 +0,0 @@ -# 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. -""" -Monkey-patches for DTensor-aware operations. - -These patches are applied at model-loading time (during ``from_pretrained``) -so that modeling files stay free of DTensor-specific code. -""" - -from __future__ import annotations - -import inspect -import sys -from functools import wraps - -from torch.distributed.tensor import DTensor, Replicate - - -def _make_dtensor_rotary_wrapper(original_fn): - """Return a wrapper that promotes cos/sin to replicated DTensors. - - Models use two ``apply_rotary_pos_emb`` signatures: - - ``(q, k, cos, sin, ...)`` — most models - - ``(x, cos, sin, ...)`` — gemma3n, gemma4, glm_moe_dsa - - We detect which one at patch time via parameter count and create - the matching wrapper. - """ - params = inspect.signature(original_fn).parameters - n_required = sum(1 for p in params.values() if p.default is inspect.Parameter.empty) - - if n_required >= 4: - - @wraps(original_fn) - def _wrapper(q, k, cos, sin, *args, **kwargs): - if isinstance(q, DTensor) and not isinstance(cos, DTensor): - replicate = (Replicate(),) * q.device_mesh.ndim - cos = DTensor.from_local(cos, q.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, q.device_mesh, replicate, run_check=False) - return original_fn(q, k, cos, sin, *args, **kwargs) - else: - - @wraps(original_fn) - def _wrapper(x, cos, sin, *args, **kwargs): - if isinstance(x, DTensor) and not isinstance(cos, DTensor): - replicate = (Replicate(),) * x.device_mesh.ndim - cos = DTensor.from_local(cos, x.device_mesh, replicate, run_check=False) - sin = DTensor.from_local(sin, x.device_mesh, replicate, run_check=False) - return original_fn(x, cos, sin, *args, **kwargs) - - return _wrapper - - -def patch_dtensor_ops(model): - """Monkey-patch DTensor-aware wrappers onto the model's modeling module. - - Finds the Python module where the model class is defined and wraps - ``apply_rotary_pos_emb`` (if present) so that cos/sin tensors are - automatically promoted to replicated DTensors when the query is a DTensor. - - Called from ``apply_tensor_parallel`` after ``parallelize_module``. - """ - model_module = sys.modules.get(type(model).__module__) - if model_module is None: - return - - original_fn = getattr(model_module, "apply_rotary_pos_emb", None) - if original_fn is not None: - model_module.apply_rotary_pos_emb = _make_dtensor_rotary_wrapper(original_fn) diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index 457d0678d9e0..0b9529466876 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -37,8 +37,6 @@ from torch.distributed.tensor.parallel.style import ParallelStyle from torch.distributed.tensor.placement_types import _StridedShard - from ..distributed.patches import patch_dtensor_ops - # Cache this result has it's a C FFI call which can be pretty time-consuming _torch_distributed_available = torch.distributed.is_available() @@ -508,10 +506,6 @@ def apply_tensor_parallel(model, tp_mesh, tp_plan): parallelize_module(model, tp_mesh, parallelize_plan) - # Patch DTensor-aware operations (e.g. rotary embeddings) onto the - # model's modeling module so modeling files stay free of DTensor code. - patch_dtensor_ops(model) - # 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 From 98d2dc563f74a16bd94d0bcbf84c7669bf799d64 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Mon, 11 May 2026 00:01:01 +0000 Subject: [PATCH 071/116] better seperation --- src/transformers/core_model_loading.py | 2 +- .../{model_loading.py => sharding_utils.py} | 80 +++++++++++++++++- src/transformers/distributed/utils.py | 82 +------------------ tests/test_tensor_parallel_mixin.py | 2 +- tests/utils/test_core_model_loading.py | 2 +- ....py => test_distributed_sharding_utils.py} | 2 +- 6 files changed, 86 insertions(+), 84 deletions(-) rename src/transformers/distributed/{model_loading.py => sharding_utils.py} (71%) rename tests/utils/{test_distributed_model_loading.py => test_distributed_sharding_utils.py} (99%) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 49496facac25..dff8e8403835 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -31,7 +31,7 @@ import torch -from .distributed.model_loading import DtensorShardOperation +from .distributed.sharding_utils import DtensorShardOperation from .integrations.accelerate import get_device, offload_weight from .utils import is_env_variable_true from .utils.loading_report import LoadStateDictInfo diff --git a/src/transformers/distributed/model_loading.py b/src/transformers/distributed/sharding_utils.py similarity index 71% rename from src/transformers/distributed/model_loading.py rename to src/transformers/distributed/sharding_utils.py index c91fe6df5acc..4ff48f687de6 100644 --- a/src/transformers/distributed/model_loading.py +++ b/src/transformers/distributed/sharding_utils.py @@ -25,8 +25,9 @@ if is_torch_available(): import torch + 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 + from torch.distributed.tensor.placement_types import Shard, _StridedShard class DtensorShardOperation: @@ -197,3 +198,80 @@ def _get_sub_mesh(self, mesh_dim: int): 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 _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 + + So we bypass ``redistribute`` and call each placement's low-level + ``_to_replicate_tensor`` (manual all-gather + interleaved reorder). + + We process mesh dims **right-to-left** (innermost first). Under TP+FSDP + the 2D mesh is ``(fsdp, tp)`` and both dims can shard the same tensor dim:: + + placements = (_StridedShard(dim=0), Shard(dim=0)) + local shape = [64, 1024] (global [256, 1024], fsdp=2, tp=2) + + Right-to-left means TP is gathered first (local grows to [128, 1024]), + then FSDP (grows to [256, 1024]). Each step must pass the correct + intermediate logical shape — the global shape divided by the mesh sizes + of dims not yet gathered (to the left). + """ + mesh = tensor.device_mesh + replicate_all = tuple(Replicate() for _ in range(mesh.ndim)) + with torch.no_grad(): + if any(isinstance(p, _StridedShard) for p in tensor.placements): + local = tensor._local_tensor + placements = tensor.placements + for i in reversed(range(mesh.ndim)): + p = placements[i] + if p.is_replicate(): + continue + # Compute the logical shape seen at this step: dims to the left + # (not yet gathered) still divide their tensor dimension. + logical_shape = list(tensor.shape) + for j in range(i): + pj = placements[j] + if not pj.is_replicate(): + logical_shape[pj.dim] //= mesh.size(j) + local = p._to_replicate_tensor(local, mesh, i, logical_shape) + return DTensor.from_local(local, mesh, replicate_all, run_check=False) + + return tensor.redistribute(placements=replicate_all) + + +def convert_strided_to_shard(state_dict: dict) -> dict[str, tuple]: + # Convert _StridedShard DTensors in a state dict to plain Shard for DCP compatibility. + placement_map: dict[str, tuple] = {} + for key, value in state_dict.items(): + if isinstance(value, dict): + nested = convert_strided_to_shard(value) + for nk, nv in nested.items(): + placement_map[f"{key}.{nk}"] = nv + elif isinstance(value, DTensor) and any(isinstance(p, _StridedShard) for p in value.placements): + placement_map[key] = tuple(value.placements) + shard_placements = tuple(Shard(p.dim) if isinstance(p, _StridedShard) else p for p in value.placements) + state_dict[key] = _replicate_dtensor(value).redistribute(placements=shard_placements) + return placement_map + + +def restore_strided_from_shard(state_dict: dict, placement_map: dict[str, tuple]) -> None: + # Restore _StridedShard placements after dcp.load. + def _resolve(d, dotted_key): + parts = dotted_key.split(".", 1) + if len(parts) == 2 and parts[0] in d and isinstance(d[parts[0]], dict): + return _resolve(d[parts[0]], parts[1]) + return d, dotted_key + + for key, original_placements in placement_map.items(): + container, leaf_key = _resolve(state_dict, key) + if leaf_key in container and isinstance(container[leaf_key], DTensor): + container[leaf_key] = _replicate_dtensor(container[leaf_key]).redistribute(placements=original_placements) diff --git a/src/transformers/distributed/utils.py b/src/transformers/distributed/utils.py index 5d2f0d3b0a54..dc53b0229d41 100644 --- a/src/transformers/distributed/utils.py +++ b/src/transformers/distributed/utils.py @@ -27,8 +27,9 @@ if is_torch_available(): import torch import torch.distributed.checkpoint as dcp - from torch.distributed.tensor import DTensor, Replicate, Shard - from torch.distributed.tensor.placement_types import _StridedShard + from torch.distributed.tensor import DTensor + + from .sharding_utils import _replicate_dtensor, convert_strided_to_shard, restore_strided_from_shard def is_fsdp_enabled() -> bool: @@ -176,83 +177,6 @@ def gather_full_state_dict(model) -> dict[str, torch.Tensor]: return result -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 - - So we bypass ``redistribute`` and call each placement's low-level - ``_to_replicate_tensor`` (manual all-gather + interleaved reorder). - - We process mesh dims **right-to-left** (innermost first). Under TP+FSDP - the 2D mesh is ``(fsdp, tp)`` and both dims can shard the same tensor dim:: - - placements = (_StridedShard(dim=0), Shard(dim=0)) - local shape = [64, 1024] (global [256, 1024], fsdp=2, tp=2) - - Right-to-left means TP is gathered first (local grows to [128, 1024]), - then FSDP (grows to [256, 1024]). Each step must pass the correct - intermediate logical shape — the global shape divided by the mesh sizes - of dims not yet gathered (to the left). - """ - mesh = tensor.device_mesh - replicate_all = tuple(Replicate() for _ in range(mesh.ndim)) - with torch.no_grad(): - if any(isinstance(p, _StridedShard) for p in tensor.placements): - local = tensor._local_tensor - placements = tensor.placements - for i in reversed(range(mesh.ndim)): - p = placements[i] - if p.is_replicate(): - continue - # Compute the logical shape seen at this step: dims to the left - # (not yet gathered) still divide their tensor dimension. - logical_shape = list(tensor.shape) - for j in range(i): - pj = placements[j] - if not pj.is_replicate(): - logical_shape[pj.dim] //= mesh.size(j) - local = p._to_replicate_tensor(local, mesh, i, logical_shape) - return DTensor.from_local(local, mesh, replicate_all, run_check=False) - - return tensor.redistribute(placements=replicate_all) - - -def convert_strided_to_shard(state_dict: dict) -> dict[str, tuple]: - # Convert _StridedShard DTensors in a state dict to plain Shard for DCP compatibility. - placement_map: dict[str, tuple] = {} - for key, value in state_dict.items(): - if isinstance(value, dict): - nested = convert_strided_to_shard(value) - for nk, nv in nested.items(): - placement_map[f"{key}.{nk}"] = nv - elif isinstance(value, DTensor) and any(isinstance(p, _StridedShard) for p in value.placements): - placement_map[key] = tuple(value.placements) - shard_placements = tuple(Shard(p.dim) if isinstance(p, _StridedShard) else p for p in value.placements) - state_dict[key] = _replicate_dtensor(value).redistribute(placements=shard_placements) - return placement_map - - -def restore_strided_from_shard(state_dict: dict, placement_map: dict[str, tuple]) -> None: - # Restore _StridedShard placements after dcp.load. - def _resolve(d, dotted_key): - parts = dotted_key.split(".", 1) - if len(parts) == 2 and parts[0] in d and isinstance(d[parts[0]], dict): - return _resolve(d[parts[0]], parts[1]) - return d, dotted_key - - for key, original_placements in placement_map.items(): - container, leaf_key = _resolve(state_dict, key) - if leaf_key in container and isinstance(container[leaf_key], DTensor): - container[leaf_key] = _replicate_dtensor(container[leaf_key]).redistribute(placements=original_placements) - - def save_optimizer(optimizer, checkpoint_dir: str) -> None: # Save optimizer state via DCP, handling _StridedShard placements transparently. osd = optimizer.state_dict() diff --git a/tests/test_tensor_parallel_mixin.py b/tests/test_tensor_parallel_mixin.py index 9c5ab3fba9af..e417bc8f3f44 100644 --- a/tests/test_tensor_parallel_mixin.py +++ b/tests/test_tensor_parallel_mixin.py @@ -17,7 +17,7 @@ from transformers import TorchAoConfig, set_seed from transformers.distributed import DistributedConfig -from transformers.distributed.utils import _replicate_dtensor +from transformers.distributed.sharding_utils import _replicate_dtensor from transformers.integrations.tensor_parallel import _get_parameter_tp_plan from transformers.testing_utils import ( is_tensor_parallel_test, diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py index b0156b47c5c4..15c9f79ba74d 100644 --- a/tests/utils/test_core_model_loading.py +++ b/tests/utils/test_core_model_loading.py @@ -38,7 +38,7 @@ from transformers.utils.import_utils import is_triton_available from ..test_modeling_common import compare_state_dicts -from .test_distributed_model_loading import FakeMesh, _make_dtensor_shard_op +from .test_distributed_sharding_utils import FakeMesh, _make_dtensor_shard_op class TestWeightGlobMatching(unittest.TestCase): diff --git a/tests/utils/test_distributed_model_loading.py b/tests/utils/test_distributed_sharding_utils.py similarity index 99% rename from tests/utils/test_distributed_model_loading.py rename to tests/utils/test_distributed_sharding_utils.py index 2dfdb026a8cb..5983d0c2f5fe 100644 --- a/tests/utils/test_distributed_model_loading.py +++ b/tests/utils/test_distributed_sharding_utils.py @@ -16,7 +16,7 @@ import torch from torch.distributed.tensor.placement_types import Replicate, Shard, _StridedShard -from transformers.distributed.model_loading import DtensorShardOperation +from transformers.distributed.sharding_utils import DtensorShardOperation class FakeMesh: From 14e02aacb50ef0c50d8df4fd15bc893c51c07e2b Mon Sep 17 00:00:00 2001 From: 3outeille Date: Mon, 11 May 2026 01:20:32 +0000 Subject: [PATCH 072/116] simplify gather_full_state_dict --- .../distributed/sharding_utils.py | 4 ++ src/transformers/distributed/utils.py | 45 +++++-------------- 2 files changed, 14 insertions(+), 35 deletions(-) diff --git a/src/transformers/distributed/sharding_utils.py b/src/transformers/distributed/sharding_utils.py index 4ff48f687de6..c0fe2eac747b 100644 --- a/src/transformers/distributed/sharding_utils.py +++ b/src/transformers/distributed/sharding_utils.py @@ -25,6 +25,7 @@ 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 @@ -243,6 +244,9 @@ def _replicate_dtensor(tensor: DTensor) -> DTensor: if not pj.is_replicate(): logical_shape[pj.dim] //= mesh.size(j) local = p._to_replicate_tensor(local, mesh, i, logical_shape) + # Drain the async functional collective so downstream storage + # queries (e.g. tied-weight dedup) don't hit "invalid python storage". + local = wait_tensor(local) return DTensor.from_local(local, mesh, replicate_all, run_check=False) return tensor.redistribute(placements=replicate_all) diff --git a/src/transformers/distributed/utils.py b/src/transformers/distributed/utils.py index dc53b0229d41..df2398053583 100644 --- a/src/transformers/distributed/utils.py +++ b/src/transformers/distributed/utils.py @@ -27,6 +27,7 @@ if is_torch_available(): import torch import torch.distributed.checkpoint as dcp + from torch.distributed.checkpoint.state_dict import get_model_state_dict from torch.distributed.tensor import DTensor from .sharding_utils import _replicate_dtensor, convert_strided_to_shard, restore_strided_from_shard @@ -125,18 +126,6 @@ def init_device_mesh(distributed_config: DistributedConfig) -> torch.distributed return mesh -def _to_cpu_fresh(tensor: torch.Tensor) -> torch.Tensor: - """Plain tensor → contiguous CPU tensor with fresh storage for safetensors.""" - if tensor.device.type == "meta": - return tensor - t = tensor.detach() - if t.device.type != "cpu": - t = t.to(device="cpu") - out = torch.empty(t.shape, dtype=t.dtype, device="cpu") - out.copy_(t) - return out.contiguous() - - def gather_full_state_dict(model) -> dict[str, torch.Tensor]: """Gather all sharded params to full plain tensors for saving. @@ -144,36 +133,22 @@ def gather_full_state_dict(model) -> dict[str, torch.Tensor]: Streams one parameter at a time to avoid holding all full tensors on GPU. Only rank 0 accumulates the result; other ranks return ``{}``. """ - tp_size = model.tp_size is_rank0 = torch.distributed.get_rank() == 0 - - # Get state dict — FSDP unshard if needed (returns DTensors, not full tensors) - if getattr(model, "_is_fsdp_managed_module", False): - from torch.distributed.checkpoint.state_dict import get_model_state_dict - - state_dict = get_model_state_dict(model) - else: - state_dict = model.state_dict() - - # No TP — materialize on rank 0 only - if tp_size is None: - if is_rank0: - return {k: _to_cpu_fresh(v) for k, v in state_dict.items()} - return {} + state_dict = get_model_state_dict(model) # Stream: gather one param at a time, only rank 0 keeps the CPU copy result = {} for key, tensor in state_dict.items(): - if isinstance(tensor, DTensor): - # All ranks participate in the collective, only rank 0 keeps the result - with torch.no_grad(): - full = _replicate_dtensor(tensor).to_local() + if not isinstance(tensor, DTensor): if is_rank0: - result[key] = _to_cpu_fresh(full) - del full - elif is_rank0: - result[key] = _to_cpu_fresh(tensor) + 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 From 9acf944eb0444ab2127b4f35732dd7a901ea9e52 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Mon, 11 May 2026 01:35:21 +0000 Subject: [PATCH 073/116] simplify _replicate_dtensor --- .../distributed/sharding_utils.py | 67 +++++++++---------- 1 file changed, 30 insertions(+), 37 deletions(-) diff --git a/src/transformers/distributed/sharding_utils.py b/src/transformers/distributed/sharding_utils.py index c0fe2eac747b..d50e3d58046c 100644 --- a/src/transformers/distributed/sharding_utils.py +++ b/src/transformers/distributed/sharding_utils.py @@ -32,6 +32,13 @@ class DtensorShardOperation: + """ + TODO: add explanation of: + - Different scenario of different placement + - What is StridedShard ? + - How does saving work in nD ? + - How does loading work in nD ? + """ def __init__(self, param: DTensor): self.device_mesh = param.device_mesh self.placements = tuple(param.placements) @@ -202,55 +209,41 @@ def _norm_dim(self, dim: int) -> int: 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:: + """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 + _StridedShard -> redistribute() -> Shard ❌ NotImplementedError Shard -> redistribute() -> Replicate ✅ works - Replicate -> redistribute() -> Shard ✅ works + Replicate -> redistribute() -> Shard ✅ works Replicate -> redistribute() -> _StridedShard ✅ works - So we bypass ``redistribute`` and call each placement's low-level - ``_to_replicate_tensor`` (manual all-gather + interleaved reorder). - - We process mesh dims **right-to-left** (innermost first). Under TP+FSDP - the 2D mesh is ``(fsdp, tp)`` and both dims can shard the same tensor dim:: - - placements = (_StridedShard(dim=0), Shard(dim=0)) - local shape = [64, 1024] (global [256, 1024], fsdp=2, tp=2) + 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]: - Right-to-left means TP is gathered first (local grows to [128, 1024]), - then FSDP (grows to [256, 1024]). Each step must pass the correct - intermediate logical shape — the global shape divided by the mesh sizes - of dims not yet gathered (to the left). + i=1 (tp, Shard(0)): [64, 1024] -> [128, 1024] + i=0 (fsdp, _StridedShard(0)): [128, 1024] -> [256, 1024] """ mesh = tensor.device_mesh replicate_all = tuple(Replicate() for _ in range(mesh.ndim)) - with torch.no_grad(): - if any(isinstance(p, _StridedShard) for p in tensor.placements): - local = tensor._local_tensor - placements = tensor.placements - for i in reversed(range(mesh.ndim)): - p = placements[i] - if p.is_replicate(): - continue - # Compute the logical shape seen at this step: dims to the left - # (not yet gathered) still divide their tensor dimension. - logical_shape = list(tensor.shape) - for j in range(i): - pj = placements[j] - if not pj.is_replicate(): - logical_shape[pj.dim] //= mesh.size(j) - local = p._to_replicate_tensor(local, mesh, i, logical_shape) - # Drain the async functional collective so downstream storage - # queries (e.g. tied-weight dedup) don't hit "invalid python storage". - local = wait_tensor(local) - return DTensor.from_local(local, mesh, replicate_all, run_check=False) + if not any(isinstance(p, _StridedShard) for p in tensor.placements): return tensor.redistribute(placements=replicate_all) + with torch.no_grad(): + local = tensor._local_tensor + shape = list(local.shape) + for i in reversed(range(mesh.ndim)): + p = tensor.placements[i] + if p.is_replicate(): + continue + shape[p.dim] *= mesh.size(i) + local = p._to_replicate_tensor(local, mesh, i, shape) + local = wait_tensor(local) + return DTensor.from_local(local, mesh, replicate_all, run_check=False) + def convert_strided_to_shard(state_dict: dict) -> dict[str, tuple]: # Convert _StridedShard DTensors in a state dict to plain Shard for DCP compatibility. From 20cf4e8d86b5fe5f756619e600f288a456b88062 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Mon, 11 May 2026 02:12:58 +0000 Subject: [PATCH 074/116] fix and clean _replicate_dtensor --- src/transformers/distributed/sharding_utils.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/src/transformers/distributed/sharding_utils.py b/src/transformers/distributed/sharding_utils.py index d50e3d58046c..abf2688da5de 100644 --- a/src/transformers/distributed/sharding_utils.py +++ b/src/transformers/distributed/sharding_utils.py @@ -39,6 +39,7 @@ class DtensorShardOperation: - How does saving work in nD ? - How does loading work in nD ? """ + def __init__(self, param: DTensor): self.device_mesh = param.device_mesh self.placements = tuple(param.placements) @@ -227,20 +228,26 @@ def _replicate_dtensor(tensor: DTensor) -> DTensor: 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 tensor.placements): + if not any(isinstance(p, _StridedShard) for p in placements): return tensor.redistribute(placements=replicate_all) with torch.no_grad(): local = tensor._local_tensor - shape = list(local.shape) for i in reversed(range(mesh.ndim)): - p = tensor.placements[i] + p = placements[i] if p.is_replicate(): continue - shape[p.dim] *= mesh.size(i) - local = p._to_replicate_tensor(local, mesh, i, shape) + 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) From ca6d06b17141e2a1dc366bcb7b66b1d7b21f5ffc Mon Sep 17 00:00:00 2001 From: 3outeille Date: Mon, 11 May 2026 03:00:17 +0000 Subject: [PATCH 075/116] better doc for DtensorShardOperation --- .../distributed/sharding_utils.py | 63 +++++++++++++++++-- .../utils/test_distributed_sharding_utils.py | 48 ++++++++------ 2 files changed, 86 insertions(+), 25 deletions(-) diff --git a/src/transformers/distributed/sharding_utils.py b/src/transformers/distributed/sharding_utils.py index abf2688da5de..bed88fa7b92c 100644 --- a/src/transformers/distributed/sharding_utils.py +++ b/src/transformers/distributed/sharding_utils.py @@ -32,12 +32,63 @@ class DtensorShardOperation: - """ - TODO: add explanation of: - - Different scenario of different placement - - What is StridedShard ? - - How does saving work in nD ? - - How does loading work in nD ? + """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 QKV | [_StridedShard(d, sf=3)] | + | 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).| + + Saving (in utils.py) + -------------------- + During `save_pretrained`, each DTensor parameter must be all-gathered + back to a full tensor on rank 0 so it can be written to safetensors. + + | Placements | Path taken | + |----------------------|-----------------------------------------------------| + | No _StridedShard | `redistribute(Replicate)` | + | Has _StridedShard | Manual right-to-left `_to_replicate_tensor` walk | """ def __init__(self, param: DTensor): diff --git a/tests/utils/test_distributed_sharding_utils.py b/tests/utils/test_distributed_sharding_utils.py index 5983d0c2f5fe..aa67c3695fa4 100644 --- a/tests/utils/test_distributed_sharding_utils.py +++ b/tests/utils/test_distributed_sharding_utils.py @@ -86,25 +86,35 @@ def _make_dtensor_shard_op(mesh, placements, param_shape, local_shape): class TestDtensorShardOperation(unittest.TestCase): """Unit tests for DtensorShardOperation. - 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. + 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): From 7e2115f4bdbb08239c6b52e8434174ec0a9e4e36 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 12 May 2026 09:50:53 +0000 Subject: [PATCH 076/116] fix saving optimizer with DCP for fused weights --- .../distributed/sharding_utils.py | 222 +++++++++++++++--- src/transformers/distributed/utils.py | 58 +++-- .../utils/test_distributed_sharding_utils.py | 100 +++++++- train_fsdp_tp.py | 4 +- 4 files changed, 323 insertions(+), 61 deletions(-) diff --git a/src/transformers/distributed/sharding_utils.py b/src/transformers/distributed/sharding_utils.py index bed88fa7b92c..f58466efbb40 100644 --- a/src/transformers/distributed/sharding_utils.py +++ b/src/transformers/distributed/sharding_utils.py @@ -54,7 +54,6 @@ class DtensorShardOperation: | Scenario | Placements | |---------------------------------------------------|---------------------------------------| | TP-only, non-fused (e.g. q_proj/k_proj/v_proj) | [Shard(d)] | - | TP-only, fused QKV | [_StridedShard(d, sf=3)] | | 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)] | @@ -79,16 +78,6 @@ class DtensorShardOperation: | 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).| - - Saving (in utils.py) - -------------------- - During `save_pretrained`, each DTensor parameter must be all-gathered - back to a full tensor on rank 0 so it can be written to safetensors. - - | Placements | Path taken | - |----------------------|-----------------------------------------------------| - | No _StridedShard | `redistribute(Replicate)` | - | Has _StridedShard | Manual right-to-left `_to_replicate_tensor` walk | """ def __init__(self, param: DTensor): @@ -303,30 +292,187 @@ def _replicate_dtensor(tensor: DTensor) -> DTensor: return DTensor.from_local(local, mesh, replicate_all, run_check=False) -def convert_strided_to_shard(state_dict: dict) -> dict[str, tuple]: - # Convert _StridedShard DTensors in a state dict to plain Shard for DCP compatibility. - placement_map: dict[str, tuple] = {} - for key, value in state_dict.items(): - if isinstance(value, dict): - nested = convert_strided_to_shard(value) - for nk, nv in nested.items(): - placement_map[f"{key}.{nk}"] = nv - elif isinstance(value, DTensor) and any(isinstance(p, _StridedShard) for p in value.placements): - placement_map[key] = tuple(value.placements) - shard_placements = tuple(Shard(p.dim) if isinstance(p, _StridedShard) else p for p in value.placements) - state_dict[key] = _replicate_dtensor(value).redistribute(placements=shard_placements) - return placement_map - - -def restore_strided_from_shard(state_dict: dict, placement_map: dict[str, tuple]) -> None: - # Restore _StridedShard placements after dcp.load. - def _resolve(d, dotted_key): - parts = dotted_key.split(".", 1) - if len(parts) == 2 and parts[0] in d and isinstance(d[parts[0]], dict): - return _resolve(d[parts[0]], parts[1]) - return d, dotted_key - - for key, original_placements in placement_map.items(): - container, leaf_key = _resolve(state_dict, key) - if leaf_key in container and isinstance(container[leaf_key], DTensor): - container[leaf_key] = _replicate_dtensor(container[leaf_key]).redistribute(placements=original_placements) +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] \ No newline at end of file diff --git a/src/transformers/distributed/utils.py b/src/transformers/distributed/utils.py index df2398053583..1aa87868ea74 100644 --- a/src/transformers/distributed/utils.py +++ b/src/transformers/distributed/utils.py @@ -27,10 +27,19 @@ if is_torch_available(): import torch import torch.distributed.checkpoint as dcp - from torch.distributed.checkpoint.state_dict import get_model_state_dict + 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 - from .sharding_utils import _replicate_dtensor, convert_strided_to_shard, restore_strided_from_shard + from .sharding_utils import ( + _replicate_dtensor, + fuse_optimizer_state, + get_fusion_metadata, + unfuse_optimizer_state, + ) def is_fsdp_enabled() -> bool: @@ -152,21 +161,30 @@ def gather_full_state_dict(model) -> dict[str, torch.Tensor]: return result -def save_optimizer(optimizer, checkpoint_dir: str) -> None: - # Save optimizer state via DCP, handling _StridedShard placements transparently. - osd = optimizer.state_dict() - placement_map = convert_strided_to_shard(osd) - dcp.save({"optimizer": osd}, checkpoint_id=checkpoint_dir) - if placement_map and torch.distributed.get_rank() == 0: - torch.save(placement_map, os.path.join(checkpoint_dir, "placement_map.pt")) - - -def load_optimizer(optimizer, checkpoint_dir: str) -> None: - # Load optimizer state via DCP, restoring _StridedShard placements transparently. - osd = optimizer.state_dict() - dcp.load({"optimizer": osd}, checkpoint_id=checkpoint_dir) - pmap_path = os.path.join(checkpoint_dir, "placement_map.pt") - if os.path.exists(pmap_path): - placement_map = torch.load(pmap_path, weights_only=False) - restore_strided_from_shard(osd, placement_map) - optimizer.load_state_dict(osd) +def save_optimizer(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. + """ + 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(model, optimizer, checkpoint_dir: str) -> None: + """Load optimizer state via DCP. + + Symmetric to `save_optimizer`: 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. + """ + 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) diff --git a/tests/utils/test_distributed_sharding_utils.py b/tests/utils/test_distributed_sharding_utils.py index aa67c3695fa4..0358a0c6cafb 100644 --- a/tests/utils/test_distributed_sharding_utils.py +++ b/tests/utils/test_distributed_sharding_utils.py @@ -11,14 +11,32 @@ # 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 +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.""" @@ -360,5 +378,85 @@ def test_slice_and_cat(self): 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/train_fsdp_tp.py b/train_fsdp_tp.py index f4c3f84ebe4b..94eacfacfb7e 100644 --- a/train_fsdp_tp.py +++ b/train_fsdp_tp.py @@ -88,7 +88,7 @@ def build_fixed_batches(dp_rank): optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr) if args.resume_dir: - load_optimizer(optimizer, os.path.join(args.resume_dir, "optimizer")) + load_optimizer(model, optimizer, os.path.join(args.resume_dir, "optimizer")) if rank == 0: print(f"Resumed optimizer from {args.resume_dir}") @@ -120,7 +120,7 @@ def build_fixed_batches(dp_rank): # Save model (HF format) and optimizer (DCP) model.save_pretrained(args.save_dir) - save_optimizer(optimizer, os.path.join(args.save_dir, "optimizer")) + save_optimizer(model, optimizer, os.path.join(args.save_dir, "optimizer")) if rank == 0: print(f"Saved to {args.save_dir}") From 1c6f8484cae872570307e0ee20dbacff0d2a046f Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 13 May 2026 06:07:38 +0000 Subject: [PATCH 077/116] save_pretrained(distributed_checkpoint=true) --- compare_save_reload.sh | 33 +++++ src/transformers/distributed/utils.py | 53 ++++++-- src/transformers/modeling_utils.py | 16 ++- train_save_reload.py | 186 ++++++++++++++++++++++++++ 4 files changed, 274 insertions(+), 14 deletions(-) create mode 100755 compare_save_reload.sh create mode 100644 train_save_reload.py diff --git a/compare_save_reload.sh b/compare_save_reload.sh new file mode 100755 index 000000000000..44de35d36f23 --- /dev/null +++ b/compare_save_reload.sh @@ -0,0 +1,33 @@ +#!/bin/bash +# Compare the loss/grad_norm trajectory of `train_save_reload.py` with and without +# the mid-training save/reload. They should match step-for-step. +# +# Each mode writes to its own dir (`./checkpoints_baseline/`, `./checkpoints_save_reload/`) +# so the artifacts are inspectable side-by-side after the run. +set -euo pipefail + +NPROC="${NPROC:-4}" +BASELINE_LOG="${BASELINE_LOG:-baseline.log}" +SAVE_RELOAD_LOG="${SAVE_RELOAD_LOG:-save_reload.log}" +DIFF_LOG="${DIFF_LOG:-save_reload_diff.log}" + +rm -rf ./checkpoints_baseline ./checkpoints_save_reload + +# Pull out the per-step training lines and the post-reload generation lines so +# the diff is mechanical and covers both the training trajectory and the +# generated tokens/text. +filter_steps() { grep -E '^step |^# gen '; } + +echo "=== Run 1/2: --mode baseline ===" +torchrun --nproc_per_node="$NPROC" train_save_reload.py --mode baseline 2>&1 \ + | tee /dev/stderr | filter_steps > "$BASELINE_LOG" + +echo +echo "=== Run 2/2: --mode save_reload ===" +torchrun --nproc_per_node="$NPROC" train_save_reload.py --mode save_reload 2>&1 \ + | tee /dev/stderr | filter_steps > "$SAVE_RELOAD_LOG" + +echo +echo "=== Diff (baseline vs save_reload) ===" +git diff --no-index --color --word-diff=color "$BASELINE_LOG" "$SAVE_RELOAD_LOG" | tee "$DIFF_LOG" || true +echo "Diff written to $DIFF_LOG" diff --git a/src/transformers/distributed/utils.py b/src/transformers/distributed/utils.py index 1aa87868ea74..13064f093646 100644 --- a/src/transformers/distributed/utils.py +++ b/src/transformers/distributed/utils.py @@ -17,7 +17,13 @@ from typing import TYPE_CHECKING from ..utils import is_torch_available, is_torch_greater_or_equal, strtobool - +from .sharding_utils import ( + _find_strided_shard_placement_from_fused_params, + _replicate_dtensor, + fuse_optimizer_state, + get_fusion_metadata, + unfuse_optimizer_state, +) if TYPE_CHECKING: import torch.nn as nn @@ -33,13 +39,7 @@ set_optimizer_state_dict, ) from torch.distributed.tensor import DTensor - - from .sharding_utils import ( - _replicate_dtensor, - fuse_optimizer_state, - get_fusion_metadata, - unfuse_optimizer_state, - ) + from torch.distributed.checkpoint.hf_storage import HuggingFaceStorageWriter def is_fsdp_enabled() -> bool: @@ -145,7 +145,6 @@ def gather_full_state_dict(model) -> dict[str, torch.Tensor]: is_rank0 = torch.distributed.get_rank() == 0 state_dict = get_model_state_dict(model) - # Stream: gather one param at a time, only rank 0 keeps the CPU copy result = {} for key, tensor in state_dict.items(): if not isinstance(tensor, DTensor): @@ -161,7 +160,36 @@ def gather_full_state_dict(model) -> dict[str, torch.Tensor]: return result -def save_optimizer(model, optimizer, checkpoint_dir: str) -> None: +def save_model_checkpoint(model, checkpoint_dir: str) -> None: + """Save model parameters as standard HF-format sharded safetensors using + DCP + HuggingFaceStorageWriter with consolidation enabled. + + Every rank first writes its own shard in parallel under + `/sharded/`, then a consolidation pass reads those shards + and emits HF-compatible `model-*-of-N.safetensors` (+ index) at + `/`. The result is a directory `from_pretrained` reads + through its normal path — no special flag needed at load time. + + DTensors carrying an uncomposed `_StridedShard` placement (e.g. fused + gate||up MoE weights) are replicated to a full tensor on every rank + before the save, otherwise DCP cannot encode that placement. + """ + state_dict = get_model_state_dict(model) + for key, value in list(state_dict.items()): + if isinstance(value, DTensor) and _find_strided_shard_placement_from_fused_params(value.placements) is not None: + state_dict[key] = _replicate_dtensor(value) + + dcp.save( + state_dict, + storage_writer=HuggingFaceStorageWriter( + path=checkpoint_dir, + save_distributed=True, + enable_consolidation=True, + ), + ) + + +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 @@ -174,11 +202,10 @@ def save_optimizer(model, optimizer, checkpoint_dir: str) -> None: unfuse_optimizer_state(optimizer_state_dict, fusion_metadata) dcp.save({"optimizer": optimizer_state_dict}, checkpoint_id=checkpoint_dir) - -def load_optimizer(model, optimizer, checkpoint_dir: str) -> None: +def load_optimizer_distributed(model, optimizer, checkpoint_dir: str) -> None: """Load optimizer state via DCP. - Symmetric to `save_optimizer`: build the unfused template, let DCP fill + 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. """ diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 420fa01a541d..2e4a45da83a8 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -52,7 +52,7 @@ revert_weight_conversion, ) from .distributed import DistributedConfig -from .distributed.utils import gather_full_state_dict, init_device_mesh, is_fsdp_enabled +from .distributed.utils import gather_full_state_dict, init_device_mesh, is_fsdp_enabled, save_model_checkpoint 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 @@ -1948,6 +1948,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[cls.__module__] # This can happen for a custom model in a jupyter notebook or repl for example - simply do not allow to set it then if not hasattr(class_module, "__file__"): @@ -1967,6 +1969,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[cls.__module__] # This can happen for a custom model in a jupyter notebook or repl for example - simply do not allow to set it then if not hasattr(class_module, "__file__"): @@ -3176,6 +3180,7 @@ def save_pretrained( token: str | bool | None = None, save_peft_format: bool = True, save_original_format: bool = True, + distributed_checkpoint: bool = False, **kwargs, ): """ @@ -3325,6 +3330,15 @@ def save_pretrained( if distributed_config is not None: model_to_save.config.distributed_config = distributed_config + if distributed_checkpoint: + if torch.distributed.is_initialized() and getattr(self, "device_mesh", None) is None: + raise ValueError( + "save_pretrained(distributed_checkpoint=True) requires the model to have been " + "initialized with a distributed_config (device_mesh is None)." + ) + save_model_checkpoint(self, save_directory) + return + # Get the model state_dict (handles FSDP unshard + TP gather in one call) if state_dict is None: if getattr(self, "device_mesh", None) is not None: diff --git a/train_save_reload.py b/train_save_reload.py new file mode 100644 index 000000000000..68a60c25b6df --- /dev/null +++ b/train_save_reload.py @@ -0,0 +1,186 @@ +"""Minimal save/reload demo: FSDP + TP on Isotonic/TinyMixtral-4x248M-MoE. + + torchrun --nproc_per_node=4 train_save_reload.py # save+reload + torchrun --nproc_per_node=4 train_save_reload.py --mode baseline # straight N steps + +The training loop deliberately overfits a single fixed sample so that, after +enough steps, the model memorizes it and `generate()` from a prefix produces +the rest of the sentence verbatim. To verify the save/reload round-trip is +lossless, run both modes and diff the loss / grad_norm logs *and* the final +generated token stream — they should all match step-for-step. +""" + +import argparse +import os + +import torch +from torch.distributed.tensor import DTensor + +from transformers import AutoModelForCausalLM, AutoTokenizer +from transformers.distributed import DistributedConfig +from transformers.distributed.utils import _replicate_dtensor, load_optimizer_distributed, save_optimizer_distributed + + +MODEL_NAME = "Isotonic/TinyMixtral-4x248M-MoE" +TOTAL_STEPS = 30 +HALFWAY = TOTAL_STEPS // 2 +LR = 1e-3 +SEED = 42 +BATCH_SIZE = 1 + +# A single passage long enough to tokenize to at least SEQ_LEN+1 tokens. The +# prompt below is a prefix of this; after overfitting, generation should +# reproduce the continuation verbatim. +OVERFIT_TEXT = ( + "In a quiet village nestled between rolling hills and a slow river, the " + "autumn mornings arrived with mist that hung low over the fields and a sky " + "that turned from grey to pale gold as the sun climbed." +) +GEN_PROMPT = "In a quiet village" + + +def run_phase(model, optimizer, batch_iter, local_rank, rank, start, stop): + for step in range(start, stop): + batch = next(batch_iter) + input_ids = batch["input_ids"].to(f"cuda:{local_rank}") + labels = batch["labels"].to(f"cuda:{local_rank}") + + loss = model(input_ids, labels=labels).loss + loss.backward() + + # Custom grad clip that tolerates DTensor grads with mixed placements: + # _replicate_dtensor handles _StridedShard (which redistribute() can't). + grads = [p.grad for p in model.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=2.0) + torch.nn.utils.clip_grads_with_norm_(grads, max_norm=1.0, total_norm=total_norm) + + optimizer.step() + optimizer.zero_grad() + + if rank == 0: + # Single canonical line per step so `diff` between modes is mechanical. + print(f"step {step:>3d} | loss {loss.item():.6f} | grad_norm {total_norm.item():.6f}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--mode", choices=["save_reload", "baseline"], default="save_reload") + args = parser.parse_args() + + torch.distributed.init_process_group(backend="nccl") + rank = int(os.environ["RANK"]) + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + torch.manual_seed(SEED) + + distributed_config = DistributedConfig( + tp_size=2, + fsdp_size=2, + tp_plan="auto", + fsdp_plan="auto", + enable_sequence_parallel=True, + ) + + # Each mode writes to its own top-level directory so the artifacts of one run + # don't clobber the other and can be inspected side-by-side after the fact. + save_dir = f"./checkpoints_{args.mode}" + intermediate_dir = os.path.join(save_dir, "intermediate") + + if rank == 0: + print(f"# mode = {args.mode} | save_dir = {save_dir}") + + # Build the initial model + optimizer. + model = AutoModelForCausalLM.from_pretrained( + MODEL_NAME, + distributed_config=distributed_config, + torch_dtype=torch.bfloat16, + ) + optimizer = torch.optim.AdamW(model.parameters(), lr=LR) + model.train() + + tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + ids = tokenizer(OVERFIT_TEXT, return_tensors="pt").input_ids[0] + fixed_batch = { + "input_ids": ids.unsqueeze(0).to(f"cuda:{local_rank}"), + "labels": ids.unsqueeze(0).to(f"cuda:{local_rank}"), + } + + def fixed_iter(): + while True: + yield fixed_batch + + batch_iter = fixed_iter() + + if args.mode == "baseline": + run_phase(model, optimizer, batch_iter, local_rank, rank, 0, TOTAL_STEPS) + else: + # 1. Train first half. + run_phase(model, optimizer, batch_iter, local_rank, rank, 0, HALFWAY) + + # 2. Save (model via DCP→HF-format consolidation; optimizer via DCP). + model.save_pretrained(intermediate_dir, distributed_checkpoint=True) + save_optimizer_distributed(model, optimizer, os.path.join(intermediate_dir, "optimizer")) + if rank == 0: + print(f"# saved intermediate to {intermediate_dir}") + + # 3. Tear down + reload from disk. Note: the dataloader iterator stays alive across + # the boundary so batch indices line up with the baseline run. + del model, optimizer + torch.cuda.empty_cache() + model = AutoModelForCausalLM.from_pretrained( + intermediate_dir, + distributed_config=distributed_config, + torch_dtype=torch.bfloat16, + ) + optimizer = torch.optim.AdamW(model.parameters(), lr=LR) + load_optimizer_distributed(model, optimizer, os.path.join(intermediate_dir, "optimizer")) + model.train() + if rank == 0: + print(f"# reloaded model + optimizer from {intermediate_dir}") + + # 4. Train second half. + run_phase(model, optimizer, batch_iter, local_rank, rank, HALFWAY, TOTAL_STEPS) + + # Final save: canonical safetensors for the model + DCP for the optimizer. + model.save_pretrained(save_dir) + save_optimizer_distributed(model, optimizer, os.path.join(save_dir, "optimizer")) + if rank == 0: + print(f"# saved final model + optimizer to {save_dir}") + + del model, optimizer + torch.cuda.empty_cache() + + gen_distributed_config = DistributedConfig( + tp_size=4, + tp_plan="auto", + enable_sequence_parallel=False, + ) + model = AutoModelForCausalLM.from_pretrained( + save_dir, + distributed_config=gen_distributed_config, + torch_dtype=torch.bfloat16, + ) + model.eval() + inputs = tokenizer(GEN_PROMPT, return_tensors="pt").to(f"cuda:{local_rank}") + max_new = ids.numel() - inputs.input_ids.shape[-1] + with torch.no_grad(): + output_ids = model.generate(**inputs, max_new_tokens=max_new, do_sample=False) + + if rank == 0: + tokens = output_ids[0].tolist() + expected = ids.tolist() + print(f"# gen tokens: {tokens}") + print(f"# exp tokens: {expected}") + print(f"# gen text: {tokenizer.decode(tokens, skip_special_tokens=True)!r}") + print(f"# exp text: {tokenizer.decode(expected, skip_special_tokens=True)!r}") + assert tokens == expected, ( + f"generated tokens do not match OVERFIT_TEXT — " + f"first mismatch at index {next((i for i, (g, e) in enumerate(zip(tokens, expected)) if g != e), min(len(tokens), len(expected)))}" + ) + + torch.distributed.destroy_process_group() From 41bc6eb0a3fa99963d0bb03a8df6766bc968c3ce Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 13 May 2026 08:06:29 +0000 Subject: [PATCH 078/116] linting --- .../distributed/sharding_utils.py | 2 +- src/transformers/distributed/utils.py | 9 ++++++-- .../utils/test_distributed_sharding_utils.py | 22 +++++++++---------- 3 files changed, 19 insertions(+), 14 deletions(-) diff --git a/src/transformers/distributed/sharding_utils.py b/src/transformers/distributed/sharding_utils.py index f58466efbb40..295e0a00242c 100644 --- a/src/transformers/distributed/sharding_utils.py +++ b/src/transformers/distributed/sharding_utils.py @@ -475,4 +475,4 @@ def fuse_optimizer_state(optimizer_state_dict, fusion_metadata): optimizer_state[param_name] = merged_fields for unfused_key in metadata["unfused_keys"]: - del optimizer_state[unfused_key] \ No newline at end of file + del optimizer_state[unfused_key] diff --git a/src/transformers/distributed/utils.py b/src/transformers/distributed/utils.py index 13064f093646..328ce6c522a4 100644 --- a/src/transformers/distributed/utils.py +++ b/src/transformers/distributed/utils.py @@ -25,6 +25,7 @@ unfuse_optimizer_state, ) + if TYPE_CHECKING: import torch.nn as nn @@ -33,13 +34,13 @@ if is_torch_available(): import torch import torch.distributed.checkpoint as dcp + from torch.distributed.checkpoint.hf_storage import HuggingFaceStorageWriter from torch.distributed.checkpoint.state_dict import ( get_model_state_dict, get_optimizer_state_dict, set_optimizer_state_dict, ) from torch.distributed.tensor import DTensor - from torch.distributed.checkpoint.hf_storage import HuggingFaceStorageWriter def is_fsdp_enabled() -> bool: @@ -176,7 +177,10 @@ def save_model_checkpoint(model, checkpoint_dir: str) -> None: """ 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: + 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( @@ -202,6 +206,7 @@ def save_optimizer_distributed(model, optimizer, checkpoint_dir: str) -> None: 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. diff --git a/tests/utils/test_distributed_sharding_utils.py b/tests/utils/test_distributed_sharding_utils.py index 0358a0c6cafb..677f13ca2638 100644 --- a/tests/utils/test_distributed_sharding_utils.py +++ b/tests/utils/test_distributed_sharding_utils.py @@ -31,12 +31,14 @@ 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.""" @@ -379,22 +381,21 @@ def test_slice_and_cat(self): 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, + (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(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), + (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, + (_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): @@ -449,7 +450,6 @@ def _optimizer_state_checkpointing_e2e_worker(rank, world_size, port, ckpt_dir): @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: From 27fc8a90c6fd29137ca95446a8ac3c42f2e777df Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 13 May 2026 08:35:01 +0000 Subject: [PATCH 079/116] refactor into a single function _dtensor_from_local_like --- src/transformers/core_model_loading.py | 11 ++--------- src/transformers/distributed/sharding_utils.py | 13 +++++++++++++ src/transformers/modeling_utils.py | 10 ++-------- 3 files changed, 17 insertions(+), 17 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index dff8e8403835..5ed2c851081e 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -31,7 +31,7 @@ import torch -from .distributed.sharding_utils import DtensorShardOperation +from .distributed.sharding_utils import DtensorShardOperation, _dtensor_from_local_like from .integrations.accelerate import get_device, offload_weight from .utils import is_env_variable_true from .utils.loading_report import LoadStateDictInfo @@ -916,14 +916,7 @@ def set_param_for_module( 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( - local_param.contiguous(), - ref.device_mesh, - ref.placements, - run_check=False, - shape=ref.shape, - stride=tuple(ref.stride()), - ) + dtensor_param = _dtensor_from_local_like(local_param, ref) with torch.no_grad(): if ref.is_meta: torch.utils.swap_tensors( diff --git a/src/transformers/distributed/sharding_utils.py b/src/transformers/distributed/sharding_utils.py index 295e0a00242c..097e30216420 100644 --- a/src/transformers/distributed/sharding_utils.py +++ b/src/transformers/distributed/sharding_utils.py @@ -249,6 +249,19 @@ def _norm_dim(self, dim: int) -> int: 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. diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 2e4a45da83a8..a33fe93bf7d5 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -52,6 +52,7 @@ revert_weight_conversion, ) from .distributed import DistributedConfig +from .distributed.sharding_utils import _dtensor_from_local_like from .distributed.utils import gather_full_state_dict, init_device_mesh, is_fsdp_enabled, save_model_checkpoint from .dynamic_module_utils import custom_object_save from .generation import CompileConfig, GenerationConfig @@ -4581,14 +4582,7 @@ def _move_missing_keys_from_meta_to_device( dtype=param.dtype, device=torch.device(param.device_mesh.device_type, torch.cuda.current_device()), ) - new_dtensor = DTensor.from_local( - local_value, - param.device_mesh, - param.placements, - run_check=False, - shape=param.shape, - stride=tuple(param.stride()), - ) + 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) From 1e25f1f064be5f1d3b61750dfb985652b4a83d65 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 13 May 2026 08:38:04 +0000 Subject: [PATCH 080/116] zeros_like instead of empty_like --- src/transformers/modeling_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index a33fe93bf7d5..ec3c84970aee 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -4563,10 +4563,10 @@ def _move_missing_keys_from_meta_to_device( # In this case we need to move everything back if is_fsdp_enabled() and not is_local_dist_rank_0() and not is_quantized: for key, param in self.named_parameters(): - value = torch.empty_like(param, device="cpu") + value = torch.zeros_like(param, device="cpu") _load_parameter_into_model(self, key, value) for key, buffer in self.named_buffers(): - value = torch.empty_like(buffer, device="cpu") + value = torch.zeros_like(buffer, device="cpu") _load_parameter_into_model(self, key, value) return From 7405892b6a4c068746a1332ebac738b4740340d0 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 13 May 2026 08:49:18 +0000 Subject: [PATCH 081/116] move tp and fsdp under distributed --- src/transformers/distributed/__init__.py | 12 ++++++++++++ .../{integrations => distributed}/fsdp.py | 2 +- .../{integrations => distributed}/tensor_parallel.py | 0 src/transformers/generation/utils.py | 2 +- src/transformers/integrations/__init__.py | 12 ------------ src/transformers/modeling_utils.py | 12 ++++++------ .../models/data2vec/modeling_data2vec_audio.py | 2 +- src/transformers/models/dia/generation_dia.py | 2 +- src/transformers/models/hubert/modeling_hubert.py | 2 +- .../models/nllb_moe/modeling_nllb_moe.py | 2 +- .../models/seamless_m4t/modeling_seamless_m4t.py | 2 +- .../seamless_m4t_v2/modeling_seamless_m4t_v2.py | 2 +- src/transformers/models/sew/modeling_sew.py | 2 +- src/transformers/models/sew/modular_sew.py | 2 +- .../models/speecht5/modeling_speecht5.py | 2 +- .../models/unispeech/modeling_unispeech.py | 2 +- .../models/unispeech_sat/modeling_unispeech_sat.py | 2 +- src/transformers/models/vits/modeling_vits.py | 2 +- .../models/wav2vec2/modeling_wav2vec2.py | 2 +- .../models/wav2vec2_bert/modeling_wav2vec2_bert.py | 2 +- .../models/wav2vec2_bert/modular_wav2vec2_bert.py | 2 +- .../modeling_wav2vec2_conformer.py | 2 +- .../wav2vec2_conformer/modular_wav2vec2_conformer.py | 2 +- src/transformers/models/wavlm/modeling_wavlm.py | 2 +- src/transformers/models/wavlm/modular_wavlm.py | 2 +- src/transformers/trainer.py | 2 +- src/transformers/trainer_seq2seq.py | 2 +- tests/test_fsdp_mixin.py | 2 +- tests/test_modeling_common.py | 2 +- tests/test_tensor_parallel_mixin.py | 2 +- 30 files changed, 44 insertions(+), 44 deletions(-) rename src/transformers/{integrations => distributed}/fsdp.py (99%) rename src/transformers/{integrations => distributed}/tensor_parallel.py (100%) diff --git a/src/transformers/distributed/__init__.py b/src/transformers/distributed/__init__.py index ba6db8358d2b..fbb12304576f 100644 --- a/src/transformers/distributed/__init__.py +++ b/src/transformers/distributed/__init__.py @@ -19,6 +19,12 @@ _import_structure = { "configuration_utils": ["DistributedConfig"], + "fsdp": ["is_fsdp_enabled", "is_fsdp_managed_module"], + "tensor_parallel": [ + "ALL_PARALLEL_STYLES", + "apply_tensor_parallel", + "verify_tp_plan", + ], } @@ -26,6 +32,12 @@ from .configuration_utils import ( DistributedConfig, ) + from .fsdp import is_fsdp_enabled, is_fsdp_managed_module + from .tensor_parallel import ( + ALL_PARALLEL_STYLES, + apply_tensor_parallel, + verify_tp_plan, + ) else: import sys diff --git a/src/transformers/integrations/fsdp.py b/src/transformers/distributed/fsdp.py similarity index 99% rename from src/transformers/integrations/fsdp.py rename to src/transformers/distributed/fsdp.py index 128cba7d253f..332554a22af7 100644 --- a/src/transformers/integrations/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -17,9 +17,9 @@ import os from typing import Any, Literal -from ..distributed.utils import is_fsdp_enabled, is_fsdp_managed_module # noqa: F401 from ..utils import is_torch_available, is_torch_greater_or_equal, logging from ..utils.quantization_config import QuantizationMethod +from .utils import is_fsdp_enabled, is_fsdp_managed_module # noqa: F401 if is_torch_available() and is_torch_greater_or_equal("2.5"): diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py similarity index 100% rename from src/transformers/integrations/tensor_parallel.py rename to src/transformers/distributed/tensor_parallel.py diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py index ec47642f5000..6c4a3cb6e76b 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 ( diff --git a/src/transformers/integrations/__init__.py b/src/transformers/integrations/__init__.py index 4d4e43958f3c..f6282cf4609e 100755 --- a/src/transformers/integrations/__init__.py +++ b/src/transformers/integrations/__init__.py @@ -50,7 +50,6 @@ "eetq": ["replace_with_eetq_linear"], "fbgemm_fp8": ["FbgemmFp8Linear", "FbgemmFp8Llama4TextExperts", "replace_with_fbgemm_fp8_linear"], "finegrained_fp8": ["FP8Linear", "replace_with_fp8_linear"], - "fsdp": ["is_fsdp_enabled", "is_fsdp_managed_module"], "ggml": [ "GGUF_CONFIG_DEFAULTS_MAPPING", "GGUF_CONFIG_MAPPING", @@ -160,11 +159,6 @@ "convert_and_export_with_cache", ] -_import_structure["tensor_parallel"] = [ - "ALL_PARALLEL_STYLES", - "apply_tensor_parallel", - "verify_tp_plan", -] try: if not is_torch_greater_or_equal("2.5"): raise OptionalDependencyNotAvailable() @@ -209,7 +203,6 @@ from .eetq import replace_with_eetq_linear from .fbgemm_fp8 import FbgemmFp8Linear, FbgemmFp8Llama4TextExperts, replace_with_fbgemm_fp8_linear from .finegrained_fp8 import FP8Linear, replace_with_fp8_linear - from .fsdp import is_fsdp_enabled, is_fsdp_managed_module from .ggml import ( GGUF_CONFIG_DEFAULTS_MAPPING, GGUF_CONFIG_MAPPING, @@ -295,11 +288,6 @@ from .quanto import replace_with_quanto_layers from .sinq import SinqDeserialize, SinqQuantize from .spqr import replace_with_spqr_linear - from .tensor_parallel import ( - ALL_PARALLEL_STYLES, - apply_tensor_parallel, - verify_tp_plan, - ) from .vptq import replace_with_vptq_linear try: diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index ec3c84970aee..af1232acdb3b 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -52,7 +52,13 @@ revert_weight_conversion, ) from .distributed import DistributedConfig +from .distributed.fsdp import apply_fully_shard_data_parallel from .distributed.sharding_utils import _dtensor_from_local_like +from .distributed.tensor_parallel import ( + _get_parameter_tp_plan, + apply_tensor_parallel, + verify_tp_plan, +) from .distributed.utils import gather_full_state_dict, init_device_mesh, is_fsdp_enabled, save_model_checkpoint from .dynamic_module_utils import custom_object_save from .generation import CompileConfig, GenerationConfig @@ -71,16 +77,10 @@ from .integrations.flash_attention import flash_attention_forward from .integrations.flash_paged import paged_attention_forward from .integrations.flex_attention import flex_attention_forward -from .integrations.fsdp import apply_fully_shard_data_parallel from .integrations.hub_kernels import allow_all_hub_kernels, is_kernel 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 ( - _get_parameter_tp_plan, - apply_tensor_parallel, - verify_tp_plan, -) from .loss.loss_utils import LOSS_MAPPING from .modeling_flash_attention_utils import ( FLASH_ATTENTION_COMPATIBILITY_MATRIX, 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/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/hubert/modeling_hubert.py b/src/transformers/models/hubert/modeling_hubert.py index e3934ba80f68..7ba97733e3ae 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/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/seamless_m4t/modeling_seamless_m4t.py b/src/transformers/models/seamless_m4t/modeling_seamless_m4t.py index e19c81de5fe6..799666f6e48c 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 51a5fd456781..75e6e2924ee1 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/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/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/unispeech/modeling_unispeech.py b/src/transformers/models/unispeech/modeling_unispeech.py index e1ee81f42950..dd3b79819c7d 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 c23fdcf16420..67659802da48 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/vits/modeling_vits.py b/src/transformers/models/vits/modeling_vits.py index b8d318ca4e26..30d1cb79b9eb 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/wav2vec2/modeling_wav2vec2.py b/src/transformers/models/wav2vec2/modeling_wav2vec2.py index 08442dad50b8..6e297a9f3265 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 354146cedb55..f63179db85b5 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 a0bd70a14976..505e9eaa2dd1 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/trainer.py b/src/transformers/trainer.py index 57042d1aad92..e6cb3015d8dd 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 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/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py index 89ccc5f221ed..dba4865f2953 100644 --- a/tests/test_fsdp_mixin.py +++ b/tests/test_fsdp_mixin.py @@ -50,7 +50,7 @@ from torch.nn.parallel import DistributedDataParallel as DDP from transformers.distributed import DistributedConfig - from transformers.integrations.fsdp import ( + from transformers.distributed.fsdp import ( _find_final_norm, apply_fully_shard_data_parallel, get_transformer_block_classes, diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py index 13b81855aaa6..cb0b9a4ee3b2 100755 --- a/tests/test_modeling_common.py +++ b/tests/test_modeling_common.py @@ -126,8 +126,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 e417bc8f3f44..3490d5539a0b 100644 --- a/tests/test_tensor_parallel_mixin.py +++ b/tests/test_tensor_parallel_mixin.py @@ -18,7 +18,7 @@ from transformers import TorchAoConfig, set_seed from transformers.distributed import DistributedConfig from transformers.distributed.sharding_utils import _replicate_dtensor -from transformers.integrations.tensor_parallel import _get_parameter_tp_plan +from transformers.distributed.tensor_parallel import _get_parameter_tp_plan from transformers.testing_utils import ( is_tensor_parallel_test, is_torch_available, From ed45c917a411971e9d090b20828541541dc14c1f Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 13 May 2026 09:03:08 +0000 Subject: [PATCH 082/116] distribute_model --- src/transformers/distributed/fsdp.py | 43 ++++++++++++++++-- src/transformers/distributed/utils.py | 49 +++++++-------------- src/transformers/integrations/accelerate.py | 2 +- src/transformers/modeling_utils.py | 20 ++++----- 4 files changed, 65 insertions(+), 49 deletions(-) diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py index 332554a22af7..9e8813bd2241 100644 --- a/src/transformers/distributed/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -15,15 +15,19 @@ import inspect import os -from typing import Any, Literal +from typing import TYPE_CHECKING, Any, Literal -from ..utils import is_torch_available, is_torch_greater_or_equal, logging +from ..utils import is_torch_available, is_torch_greater_or_equal, logging, strtobool from ..utils.quantization_config import QuantizationMethod -from .utils import is_fsdp_enabled, is_fsdp_managed_module # noqa: F401 -if is_torch_available() and is_torch_greater_or_equal("2.5"): +if TYPE_CHECKING: + import torch.nn as nn + +if is_torch_available(): import torch + +if is_torch_available() and is_torch_greater_or_equal("2.5"): import torch.distributed as dist from torch.distributed._composable.fsdp import fully_shard from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy, OffloadPolicy @@ -31,6 +35,37 @@ logger = logging.get_logger(__name__) +def is_fsdp_enabled() -> bool: + """Check if FSDP is active via Accelerate (env var based) — covers FSDP1 only.""" + if not is_torch_available(): + return False + + return ( + torch.distributed.is_available() + and torch.distributed.is_initialized() + and strtobool(os.environ.get("ACCELERATE_USE_FSDP", "False")) == 1 + and strtobool(os.environ.get("FSDP_CPU_RAM_EFFICIENT_LOADING", "False")) == 1 + ) + + +def is_fsdp_managed_module(module: nn.Module) -> bool: + """Check if a module is managed by FSDP (1 or 2).""" + if not is_torch_available(): + return False + if not torch.distributed.is_available(): + return False + + # FSDP2: attribute set by apply_fsdp2() + if getattr(module, "_is_fsdp_managed_module", False): + return True + # FSDP1: wrapped by FullyShardedDataParallel + try: + from torch.distributed.fsdp import FullyShardedDataParallel + except ImportError: + return False + return isinstance(module, FullyShardedDataParallel) + + def initialize_fsdp( fsdp_plan: dict[str, Any] | None, device_mesh=None, diff --git a/src/transformers/distributed/utils.py b/src/transformers/distributed/utils.py index 328ce6c522a4..ff6fbaddde40 100644 --- a/src/transformers/distributed/utils.py +++ b/src/transformers/distributed/utils.py @@ -16,7 +16,8 @@ import os from typing import TYPE_CHECKING -from ..utils import is_torch_available, is_torch_greater_or_equal, strtobool +from ..utils import is_torch_available, is_torch_greater_or_equal +from .fsdp import apply_fully_shard_data_parallel from .sharding_utils import ( _find_strided_shard_placement_from_fused_params, _replicate_dtensor, @@ -24,6 +25,7 @@ get_fusion_metadata, unfuse_optimizer_state, ) +from .tensor_parallel import apply_tensor_parallel if TYPE_CHECKING: @@ -43,37 +45,6 @@ from torch.distributed.tensor import DTensor -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 _ensure_torch_distributed(device_type: str): """Initialize torch.distributed if not already initialized.""" if not torch.distributed.is_initialized(): @@ -136,6 +107,20 @@ def init_device_mesh(distributed_config: DistributedConfig) -> torch.distributed return mesh +def distribute_model(model, distributed_config: DistributedConfig, device_mesh) -> nn.Module: + """Apply TP and/or FSDP2 to `model` based on the mesh dims in `device_mesh`.""" + model.config.distributed_config = distributed_config + model.device_mesh = device_mesh + mesh_dim_names = device_mesh.mesh_dim_names or () + if "tp" in mesh_dim_names: + tp_mesh = device_mesh["tp"] if device_mesh.ndim > 1 else device_mesh + model = apply_tensor_parallel(model, tp_mesh, distributed_config.tp_plan) + if "fsdp" in mesh_dim_names: + fsdp_mesh = device_mesh["fsdp"] if device_mesh.ndim > 1 else device_mesh + model = apply_fully_shard_data_parallel(model, fsdp_mesh, distributed_config.fsdp_plan) + return model + + def gather_full_state_dict(model) -> dict[str, torch.Tensor]: """Gather all sharded params to full plain tensors for saving. diff --git a/src/transformers/integrations/accelerate.py b/src/transformers/integrations/accelerate.py index a5f9835ad1ca..840e1189ab2a 100644 --- a/src/transformers/integrations/accelerate.py +++ b/src/transformers/integrations/accelerate.py @@ -26,7 +26,7 @@ from safetensors import safe_open from safetensors.torch import save_file -from ..distributed.utils import is_fsdp_enabled +from ..distributed.fsdp import is_fsdp_enabled from ..utils import ( is_accelerate_available, is_torch_available, diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index af1232acdb3b..3ea53621bc05 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -52,14 +52,18 @@ revert_weight_conversion, ) from .distributed import DistributedConfig -from .distributed.fsdp import apply_fully_shard_data_parallel +from .distributed.fsdp import is_fsdp_enabled from .distributed.sharding_utils import _dtensor_from_local_like from .distributed.tensor_parallel import ( _get_parameter_tp_plan, - apply_tensor_parallel, verify_tp_plan, ) -from .distributed.utils import gather_full_state_dict, init_device_mesh, is_fsdp_enabled, save_model_checkpoint +from .distributed.utils import ( + distribute_model, + gather_full_state_dict, + init_device_mesh, + save_model_checkpoint, +) 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 @@ -4146,15 +4150,7 @@ def from_pretrained( weight_conversions = get_model_conversion_mapping(model, key_mapping, hf_quantizer) if distributed_config is not None: - model.config.distributed_config = distributed_config - model.device_mesh = device_mesh - mesh_dim_names = device_mesh.mesh_dim_names or () - if "tp" in mesh_dim_names: - tp_mesh = device_mesh["tp"] if device_mesh.ndim > 1 else device_mesh - model = apply_tensor_parallel(model, tp_mesh, distributed_config.tp_plan) - if "fsdp" in mesh_dim_names: - fsdp_mesh = device_mesh["fsdp"] if device_mesh.ndim > 1 else device_mesh - model = apply_fully_shard_data_parallel(model, fsdp_mesh, distributed_config.fsdp_plan) + model = distribute_model(model, distributed_config, device_mesh) else: # Accelerate path: auto device mapping if device_map is not None: From f97c3a4474d9fba3bdb71d396238c798a607443f Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 13 May 2026 09:41:04 +0000 Subject: [PATCH 083/116] fix deadlock when saving --- src/transformers/distributed/utils.py | 4 + src/transformers/modeling_utils.py | 104 +++++++++++++++----------- 2 files changed, 63 insertions(+), 45 deletions(-) diff --git a/src/transformers/distributed/utils.py b/src/transformers/distributed/utils.py index ff6fbaddde40..4372a7036d04 100644 --- a/src/transformers/distributed/utils.py +++ b/src/transformers/distributed/utils.py @@ -176,6 +176,10 @@ def save_model_checkpoint(model, checkpoint_dir: str) -> None: 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. + if torch.distributed.is_initialized(): + torch.distributed.barrier() def save_optimizer_distributed(model, optimizer, checkpoint_dir: str) -> None: diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 3ea53621bc05..15834bf3dada 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -3261,8 +3261,12 @@ def save_pretrained( return os.makedirs(save_directory, exist_ok=True) + save_on_this_rank = is_main_process + if torch.distributed.is_initialized() and getattr(self, "device_mesh", None) is not None: + # DTensor gather materializes a full plain state dict only on global rank 0. + save_on_this_rank = save_on_this_rank and torch.distributed.get_rank() == 0 - 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.split(os.path.sep)[-1]) create_pr = kwargs.pop("create_pr", False) @@ -3287,7 +3291,7 @@ 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) # Don't persist distributed_config in saved config — it's runtime-only @@ -3299,7 +3303,7 @@ def save_pretrained( # Save the config try: - 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(): @@ -3345,10 +3349,12 @@ def save_pretrained( return # Get the model state_dict (handles FSDP unshard + TP gather in one call) + used_distributed_gather = False if state_dict is None: 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() @@ -3417,54 +3423,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) @@ -3480,6 +3487,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 and torch.distributed.is_initialized(): + torch.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 [] From b59c4bf7b38430752c88ea04848fa62e27ebaf42 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 13 May 2026 10:00:36 +0000 Subject: [PATCH 084/116] clip grad norm function --- src/transformers/distributed/utils.py | 49 ++++++++++++++++++++++++--- src/transformers/modeling_utils.py | 5 +-- 2 files changed, 47 insertions(+), 7 deletions(-) diff --git a/src/transformers/distributed/utils.py b/src/transformers/distributed/utils.py index 4372a7036d04..e7be16b40aaa 100644 --- a/src/transformers/distributed/utils.py +++ b/src/transformers/distributed/utils.py @@ -56,10 +56,16 @@ def _ensure_torch_distributed(device_type: str): backend_map = {"cuda": "nccl", "cpu": "gloo", "xpu": "xccl", "hpu": "hccl"} backend = backend_map.get(device_type) - torch.distributed.init_process_group(backend=backend, rank=rank, world_size=world_size) - current_device = getattr(torch, device_type) + # 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": - current_device.set_device(local_rank) + 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 " @@ -67,6 +73,22 @@ def _ensure_torch_distributed(device_type: str): ) 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`.") @@ -121,6 +143,24 @@ def distribute_model(model, distributed_config: DistributedConfig, device_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. @@ -178,8 +218,7 @@ def save_model_checkpoint(model, checkpoint_dir: str) -> None: ) # Wait for rank 0 to finish writing the HF safetensors so other # ranks don't return (and hit `from_pretrained`) before the files exist. - if torch.distributed.is_initialized(): - torch.distributed.barrier() + _distributed_barrier() def save_optimizer_distributed(model, optimizer, checkpoint_dir: str) -> None: diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 15834bf3dada..414a70514c44 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -59,6 +59,7 @@ verify_tp_plan, ) from .distributed.utils import ( + _distributed_barrier, distribute_model, gather_full_state_dict, init_device_mesh, @@ -3491,8 +3492,8 @@ def save_pretrained( # 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 and torch.distributed.is_initialized(): - torch.distributed.barrier() + if used_distributed_gather: + _distributed_barrier() @wraps(PushToHubMixin.push_to_hub) def push_to_hub(self, *args, **kwargs): From 242e814ff5622eeaae7ddce23510d868b43c2be9 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 13 May 2026 16:42:30 +0000 Subject: [PATCH 085/116] maybe_disable_foreach_and_fused_for_mixed_dtensor_groups --- src/transformers/distributed/utils.py | 41 ++++++++++++++++++++++++++- 1 file changed, 40 insertions(+), 1 deletion(-) diff --git a/src/transformers/distributed/utils.py b/src/transformers/distributed/utils.py index e7be16b40aaa..e6bde120eeee 100644 --- a/src/transformers/distributed/utils.py +++ b/src/transformers/distributed/utils.py @@ -16,7 +16,7 @@ import os from typing import TYPE_CHECKING -from ..utils import is_torch_available, is_torch_greater_or_equal +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, @@ -28,6 +28,9 @@ from .tensor_parallel import apply_tensor_parallel +logger = logging.get_logger(__name__) + + if TYPE_CHECKING: import torch.nn as nn @@ -221,6 +224,39 @@ def save_model_checkpoint(model, checkpoint_dir: str) -> None: _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. @@ -229,6 +265,7 @@ def save_optimizer_distributed(model, optimizer, checkpoint_dir: str) -> None: 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) @@ -242,9 +279,11 @@ def load_optimizer_distributed(model, optimizer, checkpoint_dir: str) -> None: 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) From 8fe831fca2f63fe6bf2e1cbe3f3fe366f97f2683 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Thu, 14 May 2026 05:45:03 +0000 Subject: [PATCH 086/116] better TP api for ease of understanding --- run_compare.sh | 12 +- .../distributed/tensor_parallel.py | 350 +++++++++--------- train_fsdp_tp.py | 55 ++- 3 files changed, 231 insertions(+), 186 deletions(-) diff --git a/run_compare.sh b/run_compare.sh index 52200d431676..4aec50851e10 100755 --- a/run_compare.sh +++ b/run_compare.sh @@ -16,12 +16,13 @@ echo "--- Launching FSDP+TP and FSDP-only in parallel ---" CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nproc_per_node=4 --master_port=29500 \ $SCRIPT $COMMON_ARGS --fsdp_size 2 --tp_size 2 --enable_sp \ - --num_steps 10 --save_dir ./checkpoints_tp > "${LOG_FSDP_TP}.phase1" 2>&1 & + --start_step 0 --stop_step 10 --save_dir ./checkpoints_tp \ + --distributed_checkpoint > "${LOG_FSDP_TP}.phase1" 2>&1 & PID1=$! CUDA_VISIBLE_DEVICES=4,5 torchrun --nproc_per_node=2 --master_port=29501 \ $SCRIPT $COMMON_ARGS --fsdp_size 2 \ - --num_steps 10 --save_dir ./checkpoints_fsdp > "${LOG_FSDP_ONLY}.phase1" 2>&1 & + --start_step 0 --stop_step 10 --save_dir ./checkpoints_fsdp > "${LOG_FSDP_ONLY}.phase1" 2>&1 & PID2=$! echo "FSDP+TP PID=$PID1 | FSDP-only PID=$PID2" @@ -34,13 +35,14 @@ echo "--- Launching FSDP+TP and FSDP-only in parallel ---" CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nproc_per_node=4 --master_port=29500 \ $SCRIPT $COMMON_ARGS --fsdp_size 2 --tp_size 2 --enable_sp \ - --num_steps 10 --start_step 10 \ - --resume_dir ./checkpoints_tp --save_dir ./checkpoints_tp_resumed > "${LOG_FSDP_TP}.phase2" 2>&1 & + --start_step 10 --stop_step 20 \ + --resume_dir ./checkpoints_tp --save_dir ./checkpoints_tp_resumed \ + --distributed_checkpoint > "${LOG_FSDP_TP}.phase2" 2>&1 & PID1=$! CUDA_VISIBLE_DEVICES=4,5 torchrun --nproc_per_node=2 --master_port=29501 \ $SCRIPT $COMMON_ARGS --fsdp_size 2 \ - --num_steps 10 --start_step 10 \ + --start_step 10 --stop_step 20 \ --resume_dir ./checkpoints_fsdp --save_dir ./checkpoints_fsdp_resumed > "${LOG_FSDP_ONLY}.phase2" 2>&1 & PID2=$! diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index 0b9529466876..a5fa0c1073ce 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -117,7 +117,39 @@ 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)}") -class PrepareModuleInputOutput(ParallelStyle): +class TensorParallelStyle(ParallelStyle): + + def shard_param(self, name, param, mesh): + return None + + def transform_inputs_pre_forward(self, module, args, kwargs, mesh): + return args, kwargs + + def context_around_forward(self, module): + return contextlib.nullcontext() + + def transform_output_post_forward(self, module, output, mesh): + return output + + def _apply(self, module, mesh): + for name, param in list(module.named_parameters(recurse=False)): + new_param = self.shard_param(name, param, mesh) + if new_param is not None: + module._parameters[name] = new_param + + original_forward = module.forward + + def tp_forward(*args, **kwargs): + args, kwargs = self.transform_inputs_pre_forward(module, args, kwargs, mesh) + with self.context_around_forward(module): + output = original_forward(*args, **kwargs) + return self.transform_output_post_forward(module, output, mesh) + + module.forward = tp_forward + return module + + +class PrepareModuleInputOutput(TensorParallelStyle): """Allgather input (Shard(1) → Replicate) + local split output (Replicate → Shard(1)). Used for MoE blocks with SP: the input sequence is gathered before routing, @@ -129,31 +161,24 @@ def __init__(self, use_local_output=True): super().__init__() self.use_local_output = use_local_output - def _apply(self, module, device_mesh): - def input_hook(mod, inputs): - x = inputs[0] if isinstance(inputs, tuple) else inputs - if not isinstance(x, DTensor): - x = DTensor.from_local(x, device_mesh, [Shard(1)], run_check=False) - x = x.redistribute(placements=[Replicate()]) - x = x.to_local() - return (x,) + (inputs[1:] if isinstance(inputs, tuple) else ()) - - def output_hook(mod, inputs, output): - if not isinstance(output, DTensor): - output = DTensor.from_local(output, device_mesh, [Replicate()], run_check=False) - output = output.redistribute(placements=[Shard(1)]) - return output.to_local() - - module.register_forward_pre_hook(input_hook) - module.register_forward_hook(output_hook) - return module + def transform_inputs_pre_forward(self, module, args, kwargs, mesh): + x = args[0] + if not isinstance(x, DTensor): + x = DTensor.from_local(x, mesh, [Shard(1)], run_check=False) + x = x.redistribute(placements=[Replicate()]).to_local() + return (x,) + args[1:], kwargs + + def transform_output_post_forward(self, module, output, mesh): + if not isinstance(output, DTensor): + output = DTensor.from_local(output, mesh, [Replicate()], run_check=False) + return output.redistribute(placements=[Shard(1)]).to_local() def _accumulate_local_param_grad(original_param: DTensor, local_grad: torch.Tensor) -> torch.Tensor: """Stitch a local grad back onto the original DTensor parameter. During forward we replace the DTensor param with a detached plain-tensor - leaf (see ``_local_dtensor_params``) because ``grouped_mm`` / fused ops do + leaf (see ``context_around_forward``) because ``grouped_mm`` / fused ops do not accept DTensor inputs. That swap breaks the autograd link between the local leaf's grad and the DTensor param's ``.grad``, so this tensor hook runs on the leaf and copies/accumulates the grad onto the original DTensor. @@ -181,37 +206,7 @@ def _accumulate_local_param_grad(original_param: DTensor, local_grad: torch.Tens return local_grad -@contextlib.contextmanager -def _local_dtensor_params(module): - """Temporarily swap DTensor params for local leaf params during one forward. - - Needed because ``grouped_mm`` / fused ops do not accept DTensor inputs: we - forward through a detached plain-tensor leaf, then rely on - ``_accumulate_local_param_grad`` (registered as a tensor hook on the leaf) - to copy the backward grad onto the original DTensor param. Restores the - DTensor params on exit (even on exception). - """ - shadows = {} - for name, param in list(module.named_parameters(recurse=False)): - if not isinstance(param, DTensor): - continue - shadows[name] = param - local = torch.nn.Parameter(param._local_tensor.detach(), requires_grad=param.requires_grad) - if param.requires_grad: - local.register_hook(lambda g, p=param: _accumulate_local_param_grad(p, g)) - module._parameters.pop(name) - setattr(module, name, local) - - try: - yield - finally: - for name, param in shadows.items(): - if hasattr(module, name): - delattr(module, name) - module.register_parameter(name, param) - - -class PackedColwiseParallel(ParallelStyle): +class PackedColwiseParallel(TensorParallelStyle): """Column-wise parallel style for fused linear weights packed along the output dimension.""" def __init__( @@ -226,69 +221,72 @@ def __init__( self.use_local_output = use_local_output self.split_factor = split_factor - def _partition_linear_fn(self, module, device_mesh): - if getattr(module, "weight", None) is None: - return - + def shard_param(self, name, param, mesh): + if name not in ("weight", "bias"): + return None packed_shard = _StridedShard(dim=0, split_factor=self.split_factor) - module.register_parameter( - "weight", - torch.nn.Parameter( - distribute_tensor(module.weight, device_mesh, [packed_shard], src_data_rank=self.src_data_rank), - requires_grad=module.weight.requires_grad, - ), + return torch.nn.Parameter( + distribute_tensor(param, mesh, [packed_shard], src_data_rank=self.src_data_rank), + requires_grad=param.requires_grad, ) - if getattr(module, "bias", None) is not None: - module.register_parameter( - "bias", - torch.nn.Parameter( - distribute_tensor(module.bias, device_mesh, [packed_shard], src_data_rank=self.src_data_rank), - requires_grad=module.bias.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): + """Swap DTensor params for local leaf params during forward. + + ``grouped_mm`` / fused ops don't accept DTensor inputs, so we forward + through a detached plain-tensor leaf and let ``_accumulate_local_param_grad`` + (a tensor hook on the leaf) copy the backward grad onto the original DTensor. + DTensor params are restored on exit (even on exception). + """ + shadows = {} + for name, param in list(module.named_parameters(recurse=False)): + if not isinstance(param, DTensor): + continue + shadows[name] = param + local = torch.nn.Parameter(param._local_tensor.detach(), requires_grad=param.requires_grad) + if param.requires_grad: + local.register_hook(lambda g, p=param: _accumulate_local_param_grad(p, g)) + module._parameters.pop(name) + setattr(module, name, local) + try: + yield + finally: + for name, param in shadows.items(): + if hasattr(module, name): + delattr(module, name) + module.register_parameter(name, param) + + def transform_output_post_forward(self, module, output, mesh): + if output is None or self.use_local_output: + return output + return DTensor.from_local( + output, mesh, (_StridedShard(dim=-1, split_factor=self.split_factor),), run_check=False + ) - def _apply(self, module, device_mesh): + def _apply(self, module, mesh): if not isinstance(module, torch.nn.Linear): raise NotImplementedError("PackedColwiseParallel currently only supports nn.Linear!") - - self._partition_linear_fn(module, device_mesh) - - input_layouts = self.input_layouts - use_local_output = self.use_local_output - split_factor = self.split_factor - original_forward = module.forward - - def tp_forward(input_tensor, *args, **kwargs): - if not isinstance(input_tensor, DTensor): - input_tensor = DTensor.from_local(input_tensor, device_mesh, input_layouts, run_check=False) - elif input_tensor.placements != input_layouts: - input_tensor = input_tensor.redistribute(placements=input_layouts) - input_tensor = input_tensor.to_local() - - with _local_dtensor_params(module): - output = original_forward(input_tensor, *args, **kwargs) - - if output is None or use_local_output: - return output - return DTensor.from_local( - output, device_mesh, (_StridedShard(dim=-1, split_factor=split_factor),), run_check=False - ) - - module.forward = tp_forward - return module + return super()._apply(module, mesh) def __repr__(self) -> str: - tmpstr = self.__class__.__name__ + "(" - tmpstr += f"input_layouts={self.input_layouts}, " - tmpstr += f"use_local_output={self.use_local_output}, " - tmpstr += f"split_factor={self.split_factor}" - tmpstr += ")" - return tmpstr + return ( + f"{self.__class__.__name__}(input_layouts={self.input_layouts}, " + f"use_local_output={self.use_local_output}, split_factor={self.split_factor})" + ) # Maps string tp_plan entries for MoE experts to DTensor placements. -# Used by MoEExpertsParallel._partition_fn to create DTensors from the config plan. +# Used by MoEExpertsParallel.shard_param to create DTensors from the config plan. _STRING_TO_PLACEMENT = { "packed_colwise": lambda: _StridedShard(dim=-2, split_factor=2), "colwise": lambda: Shard(-2), @@ -318,21 +316,24 @@ def backward(ctx, grad): return grad, None -class MoEExpertsParallel(ParallelStyle): +class MoEExpertsParallel(TensorParallelStyle): """Tensor-parallel style for MoE expert modules. Shards expert weights as DTensors, then wraps the module's ``forward`` so that grouped_mm (which needs plain tensors) works transparently. - The wrapped forward does four things: - 1. Localize inputs — wrap hidden_states as Replicate DTensor then extract - local tensor (gives us an all-reduce on the backward gradient for free). - 2. Fix routing grads — routing weights are the same on all ranks, but their - backward gradient is partial; use allreduce-sum (not divide-by-world-size). - 3. Swap params — temporarily replace DTensor params with local tensors - for grouped_mm, restore them after so save_pretrained sees DTensors. - 4. Reduce output — each rank's output is partial (only its expert shard - contributed); all-reduce to get the complete hidden state. + Lifecycle phases: + 1. shard_param — distribute each expert weight per the shard_plan. + 2. transform_inputs_pre_forward — localize hidden_states (Replicate→local, + gives us an all-reduce on the backward gradient for free), then fix + routing-weight gradients (their backward is partial; use allreduce-sum, + not divide-by-world-size). + 3. context_around_forward — swap DTensor params for local leaves so + grouped_mm sees plain tensors; restored on exit so save_pretrained + still sees DTensors. + 4. transform_output_post_forward — under TP-only each rank's output is + partial (only its expert shard contributed); reduce/redistribute to + output_layouts. """ def __init__(self, output_layouts=None, shard_plan: dict[str, str] | None = None): @@ -340,64 +341,77 @@ def __init__(self, output_layouts=None, shard_plan: dict[str, str] | None = None self.output_layouts = output_layouts or Replicate() self._moe_shard_plan: dict[str, str] = shard_plan or {} - @staticmethod - def _partition_fn(name, module, device_mesh, shard_plan): - for param_name, param in module.named_parameters(recurse=False): - plan_str = shard_plan.get(param_name) - if plan_str is None: - continue - placement_fn = _STRING_TO_PLACEMENT.get(plan_str) - if placement_fn is None: - continue - placement = placement_fn() - dtensor = distribute_tensor(param.data, device_mesh, [placement]) - module._parameters[param_name] = torch.nn.Parameter(dtensor, requires_grad=param.requires_grad) + def shard_param(self, name, param, mesh): + plan_str = self._moe_shard_plan.get(name) + if plan_str is None: + return None + placement_fn = _STRING_TO_PLACEMENT.get(plan_str) + if placement_fn is None: + return None + return torch.nn.Parameter( + distribute_tensor(param.data, mesh, [placement_fn()]), + requires_grad=param.requires_grad, + ) - def _apply(self, module, device_mesh): - self._partition_fn(module.__class__.__name__, module, device_mesh, self._moe_shard_plan) + def transform_inputs_pre_forward(self, module, args, kwargs, mesh): + hidden_states, top_k_index, top_k_weights = args + if not isinstance(hidden_states, DTensor): + hidden_states = DTensor.from_local(hidden_states, mesh, [Replicate()], run_check=False) + hidden_states = hidden_states.to_local() - output_layouts = self.output_layouts - original_forward = module.forward - tp_group = device_mesh.get_group() if device_mesh.ndim == 1 else device_mesh.get_group("tp") - - def tp_forward(hidden_states, top_k_index, top_k_weights): - # --- 1. Localize hidden_states (backward all-reduce via DTensor) --- - if not isinstance(hidden_states, DTensor): - hidden_states = DTensor.from_local(hidden_states, device_mesh, [Replicate()], run_check=False) - hidden_states = hidden_states.to_local() - - # --- 2. Fix routing weight gradients (allreduce-sum, not ÷ world_size) --- - if isinstance(top_k_weights, DTensor): - top_k_weights = top_k_weights.to_local() - top_k_weights = _AllReduceBackward.apply(top_k_weights, tp_group) - - # --- 3. Run forward with local params (grouped_mm needs plain tensors) --- - with _local_dtensor_params(module): - output = original_forward(hidden_states, top_k_index, top_k_weights) - - # --- 4. Reduce partial output --- - 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, device_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 = 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() + if isinstance(top_k_weights, DTensor): + top_k_weights = top_k_weights.to_local() + tp_group = mesh.get_group() if mesh.ndim == 1 else mesh.get_group("tp") + top_k_weights = _AllReduceBackward.apply(top_k_weights, tp_group) - module.forward = tp_forward - return module + return (hidden_states, top_k_index, top_k_weights), kwargs + + @contextlib.contextmanager + def context_around_forward(self, module): + """Swap DTensor params for local leaf params during forward. + + ``grouped_mm`` / fused ops don't accept DTensor inputs, so we forward + through a detached plain-tensor leaf and let ``_accumulate_local_param_grad`` + (a tensor hook on the leaf) copy the backward grad onto the original DTensor. + DTensor params are restored on exit (even on exception). + """ + shadows = {} + for name, param in list(module.named_parameters(recurse=False)): + if not isinstance(param, DTensor): + continue + shadows[name] = param + local = torch.nn.Parameter(param._local_tensor.detach(), requires_grad=param.requires_grad) + if param.requires_grad: + local.register_hook(lambda g, p=param: _accumulate_local_param_grad(p, g)) + module._parameters.pop(name) + setattr(module, name, local) + try: + yield + finally: + for name, param in shadows.items(): + if hasattr(module, name): + delattr(module, name) + module.register_parameter(name, param) + + def transform_output_post_forward(self, module, output, mesh): + if output is None: + return None + # Under TP-only each rank has a partial result; under TP+FSDP the + # weights may be fully gathered by FSDP, making the output complete. + has_sharded_params = any( + isinstance(p, DTensor) and any(not pl.is_replicate() for pl in p.placements) for p in module.parameters() + ) + source = Partial() if has_sharded_params else Replicate() + if not isinstance(output, DTensor): + output = DTensor.from_local(output, mesh, [source], run_check=False) + # MoE output is 2D [tokens, hidden]. For SP, Shard(1) means seq dim + # in 3D but token dim (0) in 2D. + target = self.output_layouts + if output.dim() == 2 and isinstance(target, Shard) and target.dim == 1: + target = Shard(0) + if output.placements != (target,): + output = output.redistribute(placements=(target,)) + return output.to_local() class ParallelInterface(GeneralInterface): diff --git a/train_fsdp_tp.py b/train_fsdp_tp.py index 94eacfacfb7e..a2881c8bfaf5 100644 --- a/train_fsdp_tp.py +++ b/train_fsdp_tp.py @@ -9,7 +9,11 @@ from torch.utils.data import DataLoader from transformers import AutoModelForCausalLM, AutoTokenizer from transformers.distributed import DistributedConfig -from transformers.distributed.utils import load_optimizer, save_optimizer, _replicate_dtensor +from transformers.distributed.utils import ( + _replicate_dtensor, + load_optimizer_distributed, + save_optimizer_distributed, +) def build_packed_dataset(dataset_name, tokenizer, seq_len, dp_rank, dp_world_size): """Stream + tokenize + greedy-pack documents into fixed-length (input, label) windows.""" @@ -36,7 +40,8 @@ def build_fixed_batches(dp_rank): parser = argparse.ArgumentParser() parser.add_argument("--model_name", type=str, default="Qwen/Qwen3-0.6B") - parser.add_argument("--num_steps", type=int, default=20) + parser.add_argument("--start_step", type=int, default=0, help="Inclusive start of the step range to train") + parser.add_argument("--stop_step", type=int, default=20, help="Exclusive end of the step range to train") parser.add_argument("--lr", type=float, default=3e-4) parser.add_argument("--seq_len", type=int, default=512) parser.add_argument("--batch_size", type=int, default=1) @@ -46,8 +51,12 @@ def build_fixed_batches(dp_rank): parser.add_argument("--enable_sp", action="store_true", help="Enable sequence parallelism") parser.add_argument("--seed", type=int, default=42, help="Random seed") parser.add_argument("--fixed_batches", action="store_true", help="Use pre-generated fixed batches instead of C4") - parser.add_argument("--resume_dir", type=str, default=None, help="Resume from this checkpoint directory") - parser.add_argument("--start_step", type=int, default=0, help="Starting step number (for logging)") + parser.add_argument("--resume_dir", type=str, default=None, + help="Resume model + optimizer from a save_pretrained(distributed_checkpoint=True) dir") + parser.add_argument("--save_at_step", type=int, default=None, + help="Save a distributed checkpoint at this step number (inside [start_step, stop_step))") + parser.add_argument("--distributed_checkpoint", action="store_true", + help="Use distributed_checkpoint=True for the final save (per-rank shards via DCP + HF consolidation)") args = parser.parse_args() torch.distributed.init_process_group(backend="nccl") @@ -66,12 +75,16 @@ def build_fixed_batches(dp_rank): dc_kwargs["enable_sequence_parallel"] = True distributed_config = DistributedConfig(**dc_kwargs) + # Both `args.model_name` (HF hub) and `args.resume_dir` (written by save_pretrained, + # canonical or distributed_checkpoint=True) are plain HF-format directories — same load path. load_path = args.resume_dir if args.resume_dir else args.model_name model = AutoModelForCausalLM.from_pretrained( load_path, distributed_config=distributed_config, torch_dtype=torch.bfloat16, ) + if args.resume_dir and rank == 0: + print(f"Resumed model from {args.resume_dir}") dp_rank = model.device_mesh["fsdp"].get_local_rank() if "fsdp" in model.device_mesh.mesh_dim_names else 0 dp_size = model.device_mesh["fsdp"].size() if "fsdp" in model.device_mesh.mesh_dim_names else 1 @@ -88,12 +101,18 @@ def build_fixed_batches(dp_rank): optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr) if args.resume_dir: - load_optimizer(model, optimizer, os.path.join(args.resume_dir, "optimizer")) - if rank == 0: - print(f"Resumed optimizer from {args.resume_dir}") + optim_dir = os.path.join(args.resume_dir, "optimizer") + if os.path.exists(optim_dir): + load_optimizer_distributed(model, optimizer, optim_dir) + if rank == 0: + print(f"Resumed optimizer from {optim_dir}") + elif rank == 0: + print(f"No optimizer state at {optim_dir}; starting from fresh optimizer state") + + intermediate_dir = os.path.join(args.save_dir, "intermediate") model.train() - for step in range(args.start_step, args.start_step + args.num_steps): + for step in range(args.start_step, args.stop_step): if args.fixed_batches: input_ids = fixed[step]["input_ids"].to(f"cuda:{local_rank}") labels = fixed[step]["labels"].to(f"cuda:{local_rank}") @@ -118,11 +137,21 @@ def build_fixed_batches(dp_rank): if rank == 0: print(f"Step {step:>4d} | Loss: {loss.item():.4f} | Grad norm: {total_norm.item():.4f}") - # Save model (HF format) and optimizer (DCP) - model.save_pretrained(args.save_dir) - save_optimizer(model, optimizer, os.path.join(args.save_dir, "optimizer")) - + # Mid-training distributed checkpoint: every rank writes its own shard in parallel via DCP + + # HuggingFaceStorageWriter consolidation, plus DCP optimizer save. The resulting directory + # is still HF-safetensors-compatible, so `from_pretrained(intermediate_dir, ...)` resumes it. + if args.save_at_step is not None and step == args.save_at_step: + model.save_pretrained(intermediate_dir, distributed_checkpoint=True) + save_optimizer_distributed(model, optimizer, os.path.join(intermediate_dir, "optimizer")) + if rank == 0: + print(f"Saved distributed checkpoint at step {step} to {intermediate_dir}") + + # Final save: either canonical single-file HF safetensors (rank-0 gather) or distributed + # per-rank shards (DCP + HuggingFaceStorageWriter consolidation). Both are HF-format dirs + # that `from_pretrained` can resume from. Optimizer always saved via DCP. + model.save_pretrained(args.save_dir, distributed_checkpoint=args.distributed_checkpoint) + save_optimizer_distributed(model, optimizer, os.path.join(args.save_dir, "optimizer")) if rank == 0: - print(f"Saved to {args.save_dir}") + print(f"Saved final checkpoint to {args.save_dir}") torch.distributed.destroy_process_group() \ No newline at end of file From bf30f0a1d6296692f1598c36216c775cbaaabc86 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Thu, 14 May 2026 06:31:56 +0000 Subject: [PATCH 087/116] remove shard_param to make it easier --- .../distributed/tensor_parallel.py | 89 ++++++++++--------- 1 file changed, 47 insertions(+), 42 deletions(-) diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index a5fa0c1073ce..723aaa2692f7 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -118,9 +118,19 @@ def verify_tp_plan(expected_keys: list[str], tp_plan: dict[str, str] | None): class TensorParallelStyle(ParallelStyle): - - def shard_param(self, name, param, mesh): - return None + """Base class for transformers TP styles. Installs the pre / around / post + forward hooks. Subclasses that need to shard params override `_apply` to + wrap them as DTensor placeholders before calling `super()._apply(...)`. + + Param wrapping runs on meta (the model is on meta when `apply_tensor_parallel` + is invoked); `distribute_tensor` on meta builds metadata only — no collective. + Real data flows in later, async, via DtensorShardOperation during load. + + Forward-time hooks (override what you need): + - transform_inputs_pre_forward(module, args, kwargs, mesh) → (args, kwargs) + - context_around_forward(module) → context manager wrapping the call + - transform_output_post_forward(module, output, mesh) → output + """ def transform_inputs_pre_forward(self, module, args, kwargs, mesh): return args, kwargs @@ -132,11 +142,6 @@ def transform_output_post_forward(self, module, output, mesh): return output def _apply(self, module, mesh): - for name, param in list(module.named_parameters(recurse=False)): - new_param = self.shard_param(name, param, mesh) - if new_param is not None: - module._parameters[name] = new_param - original_forward = module.forward def tp_forward(*args, **kwargs): @@ -221,15 +226,6 @@ def __init__( self.use_local_output = use_local_output self.split_factor = split_factor - def shard_param(self, name, param, mesh): - if name not in ("weight", "bias"): - return None - packed_shard = _StridedShard(dim=0, split_factor=self.split_factor) - return torch.nn.Parameter( - distribute_tensor(param, mesh, [packed_shard], src_data_rank=self.src_data_rank), - requires_grad=param.requires_grad, - ) - def transform_inputs_pre_forward(self, module, args, kwargs, mesh): input_tensor = args[0] if not isinstance(input_tensor, DTensor): @@ -276,6 +272,17 @@ def transform_output_post_forward(self, module, output, mesh): def _apply(self, module, mesh): if not isinstance(module, torch.nn.Linear): raise NotImplementedError("PackedColwiseParallel currently only supports nn.Linear!") + # Wrap weight + bias as DTensor placeholders. Runs on meta — + # distribute_tensor builds metadata only, no collective. + placement = _StridedShard(dim=0, split_factor=self.split_factor) + for name in ("weight", "bias"): + meta = module._parameters.get(name) + if meta is None: + continue + module._parameters[name] = torch.nn.Parameter( + distribute_tensor(meta, mesh, [placement], src_data_rank=None), + requires_grad=meta.requires_grad, + ) return super()._apply(module, mesh) def __repr__(self) -> str: @@ -285,15 +292,6 @@ def __repr__(self) -> str: ) -# Maps string tp_plan entries for MoE experts to DTensor placements. -# Used by MoEExpertsParallel.shard_param to create DTensors from the config plan. -_STRING_TO_PLACEMENT = { - "packed_colwise": lambda: _StridedShard(dim=-2, split_factor=2), - "colwise": lambda: Shard(-2), - "rowwise": lambda: Shard(-1), -} - - if is_torch_available() and is_torch_greater_or_equal("2.5"): class _AllReduceBackward(torch.autograd.Function): @@ -323,7 +321,8 @@ class MoEExpertsParallel(TensorParallelStyle): that grouped_mm (which needs plain tensors) works transparently. Lifecycle phases: - 1. shard_param — distribute each expert weight per the shard_plan. + 1. _apply — wrap each expert weight named in shard_plan as a DTensor + placeholder with the declared placement. 2. transform_inputs_pre_forward — localize hidden_states (Replicate→local, gives us an all-reduce on the backward gradient for free), then fix routing-weight gradients (their backward is partial; use allreduce-sum, @@ -336,22 +335,23 @@ class MoEExpertsParallel(TensorParallelStyle): output_layouts. """ - def __init__(self, output_layouts=None, shard_plan: dict[str, str] | None = None): + def __init__(self, output_layouts=None, shard_plan=None): super().__init__() self.output_layouts = output_layouts or Replicate() - self._moe_shard_plan: dict[str, str] = shard_plan or {} + self._moe_shard_plan = shard_plan or {} - def shard_param(self, name, param, mesh): - plan_str = self._moe_shard_plan.get(name) - if plan_str is None: - return None - placement_fn = _STRING_TO_PLACEMENT.get(plan_str) - if placement_fn is None: - return None - return torch.nn.Parameter( - distribute_tensor(param.data, mesh, [placement_fn()]), - requires_grad=param.requires_grad, - ) + def _apply(self, module, mesh): + # Wrap each expert weight as a DTensor placeholder. Runs on meta — + # distribute_tensor builds metadata only, no collective. + for name, placement in self._moe_shard_plan.items(): + meta = module._parameters.get(name) + if meta is None: + continue + module._parameters[name] = torch.nn.Parameter( + distribute_tensor(meta, mesh, [placement], src_data_rank=None), + requires_grad=meta.requires_grad, + ) + return super()._apply(module, mesh) def transform_inputs_pre_forward(self, module, args, kwargs, mesh): hidden_states, top_k_index, top_k_weights = args @@ -456,10 +456,15 @@ class ParallelInterface(GeneralInterface): use_local_output=True, ), "module_allgather_split": PrepareModuleInputOutput(), - # MoE — canonical shard_plan baked in (only variant in use across configs) + # MoE — canonical shard_plan baked in (only variant in use across configs). + # gate_up_proj is packed (gate||up along output dim) so we use _StridedShard + # to interleave; down_proj is plain rowwise on its input dim. "moe_experts_allreduce": MoEExpertsParallel( output_layouts=Replicate(), - shard_plan={"gate_up_proj": "packed_colwise", "down_proj": "rowwise"}, + shard_plan={ + "gate_up_proj": _StridedShard(dim=-2, split_factor=2), + "down_proj": Shard(-1), + }, ), } if is_torch_available() and is_torch_greater_or_equal("2.5") and _torch_distributed_available From 261c59b431413aacc4b0d2df5ece00666c7805bc Mon Sep 17 00:00:00 2001 From: 3outeille Date: Thu, 14 May 2026 07:49:01 +0000 Subject: [PATCH 088/116] fix import in test --- tests/utils/test_modeling_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/utils/test_modeling_utils.py b/tests/utils/test_modeling_utils.py index defc2add35bb..c02a0cbd8b0c 100644 --- a/tests/utils/test_modeling_utils.py +++ b/tests/utils/test_modeling_utils.py @@ -456,7 +456,7 @@ def fake_load_pretrained_model(model, state_dict, checkpoint_files, load_config, with ( patch("transformers.modeling_utils.init_device_mesh", return_value=fake_mesh), - patch("transformers.modeling_utils.apply_fully_shard_data_parallel", side_effect=fake_apply_fsdp), + 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, From f0f5f6744561e63bd77dce251b4591d33c41aa2c Mon Sep 17 00:00:00 2001 From: 3outeille Date: Thu, 14 May 2026 07:56:22 +0000 Subject: [PATCH 089/116] _swap_dtensor_params_for_local --- .../distributed/tensor_parallel.py | 96 ++++++++----------- 1 file changed, 41 insertions(+), 55 deletions(-) diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index 723aaa2692f7..48b9d2e0aa82 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -183,10 +183,15 @@ def _accumulate_local_param_grad(original_param: DTensor, local_grad: torch.Tens """Stitch a local grad back onto the original DTensor parameter. During forward we replace the DTensor param with a detached plain-tensor - leaf (see ``context_around_forward``) because ``grouped_mm`` / fused ops do - not accept DTensor inputs. That swap breaks the autograd link between the - local leaf's grad and the DTensor param's ``.grad``, so this tensor hook + leaf (see ``_swap_dtensor_params_for_local``) because ``grouped_mm`` / fused + ops don't accept DTensor inputs. That swap breaks the autograd link between + the local leaf's grad and the DTensor param's ``.grad``, so this tensor hook runs on the leaf and copies/accumulates the grad onto the original DTensor. + + NOTE: An autograd-aware ``param.to_local()`` swap would let backward stitch + the grad automatically, but DTensor's backward path then redistributes the + resulting grad — and that redistribute does not currently support + ``_StridedShard`` placements (used by ``MoEExpertsParallel`` / ``PackedColwiseParallel``). """ tensor_meta = original_param._spec.tensor_meta detached_grad = local_grad.detach() @@ -198,7 +203,6 @@ def _accumulate_local_param_grad(original_param: DTensor, local_grad: torch.Tens shape=tensor_meta.shape, stride=tensor_meta.stride, ) - with torch.no_grad(): existing_grad = original_param.grad if existing_grad is None: @@ -207,10 +211,40 @@ def _accumulate_local_param_grad(original_param: DTensor, local_grad: torch.Tens existing_grad._local_tensor.add_(detached_grad) else: existing_grad.add_(detached_grad) - return local_grad +@contextlib.contextmanager +def _swap_dtensor_params_for_local(module): + """Temporarily replace DTensor params with local-shard ``Parameter``s for forward. + + ``grouped_mm`` / fused kernels don't accept DTensor inputs, so each DTensor + param is swapped for a detached local ``Parameter``. A tensor hook + (``_accumulate_local_param_grad``) on the local leaf copies the backward + grad back onto the original DTensor. + + The original DTensor params are restored on exit (even on exception) so + save_pretrained / state-dict still see sharded params. + """ + shadows = {} + for name, param in list(module.named_parameters(recurse=False)): + if not isinstance(param, DTensor): + continue + shadows[name] = param + local = torch.nn.Parameter(param._local_tensor.detach(), requires_grad=param.requires_grad) + if param.requires_grad: + local.register_hook(lambda g, p=param: _accumulate_local_param_grad(p, g)) + module._parameters.pop(name) + setattr(module, name, local) + try: + yield + finally: + for name, param in shadows.items(): + if hasattr(module, name): + delattr(module, name) + module.register_parameter(name, param) + + class PackedColwiseParallel(TensorParallelStyle): """Column-wise parallel style for fused linear weights packed along the output dimension.""" @@ -235,32 +269,8 @@ def transform_inputs_pre_forward(self, module, args, kwargs, mesh): input_tensor = input_tensor.to_local() return (input_tensor,) + args[1:], kwargs - @contextlib.contextmanager def context_around_forward(self, module): - """Swap DTensor params for local leaf params during forward. - - ``grouped_mm`` / fused ops don't accept DTensor inputs, so we forward - through a detached plain-tensor leaf and let ``_accumulate_local_param_grad`` - (a tensor hook on the leaf) copy the backward grad onto the original DTensor. - DTensor params are restored on exit (even on exception). - """ - shadows = {} - for name, param in list(module.named_parameters(recurse=False)): - if not isinstance(param, DTensor): - continue - shadows[name] = param - local = torch.nn.Parameter(param._local_tensor.detach(), requires_grad=param.requires_grad) - if param.requires_grad: - local.register_hook(lambda g, p=param: _accumulate_local_param_grad(p, g)) - module._parameters.pop(name) - setattr(module, name, local) - try: - yield - finally: - for name, param in shadows.items(): - if hasattr(module, name): - delattr(module, name) - module.register_parameter(name, param) + return _swap_dtensor_params_for_local(module) def transform_output_post_forward(self, module, output, mesh): if output is None or self.use_local_output: @@ -366,32 +376,8 @@ def transform_inputs_pre_forward(self, module, args, kwargs, mesh): return (hidden_states, top_k_index, top_k_weights), kwargs - @contextlib.contextmanager def context_around_forward(self, module): - """Swap DTensor params for local leaf params during forward. - - ``grouped_mm`` / fused ops don't accept DTensor inputs, so we forward - through a detached plain-tensor leaf and let ``_accumulate_local_param_grad`` - (a tensor hook on the leaf) copy the backward grad onto the original DTensor. - DTensor params are restored on exit (even on exception). - """ - shadows = {} - for name, param in list(module.named_parameters(recurse=False)): - if not isinstance(param, DTensor): - continue - shadows[name] = param - local = torch.nn.Parameter(param._local_tensor.detach(), requires_grad=param.requires_grad) - if param.requires_grad: - local.register_hook(lambda g, p=param: _accumulate_local_param_grad(p, g)) - module._parameters.pop(name) - setattr(module, name, local) - try: - yield - finally: - for name, param in shadows.items(): - if hasattr(module, name): - delattr(module, name) - module.register_parameter(name, param) + return _swap_dtensor_params_for_local(module) def transform_output_post_forward(self, module, output, mesh): if output is None: From c135d0e13d0107858aedbbefc2bef9bdefd113e4 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Thu, 14 May 2026 08:19:56 +0000 Subject: [PATCH 090/116] fix qwen3 nanochat dots1 --- .../models/dots1/modeling_dots1.py | 54 +++++++++++++------ .../models/esm/configuration_esm.py | 4 +- .../models/nanochat/modeling_nanochat.py | 52 +++++++++--------- .../models/qwen3/modeling_qwen3.py | 49 +++++++++++------ .../models/qwen3/modular_qwen3.py | 17 +----- 5 files changed, 102 insertions(+), 74 deletions(-) diff --git a/src/transformers/models/dots1/modeling_dots1.py b/src/transformers/models/dots1/modeling_dots1.py index c6caf4e6fd91..0f8ca45a5f83 100644 --- a/src/transformers/models/dots1/modeling_dots1.py +++ b/src/transformers/models/dots1/modeling_dots1.py @@ -28,7 +28,12 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func +from ...integrations import ( + use_experts_implementation, + use_kernel_forward_from_hub, + use_kernel_func_from_hub, + use_kernelized_func, +) from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -128,6 +133,39 @@ def forward(self, x, position_ids): return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +@use_kernel_func_from_hub("rotary_pos_emb") +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, @@ -165,20 +203,6 @@ def eager_attention_forward( return attn_output, attn_weights -def rotate_half(x): - x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] - return torch.cat((-x2, x1), dim=-1) - - -def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): - cos = cos.unsqueeze(unsqueeze_dim) - sin = sin.unsqueeze(unsqueeze_dim) - q_embed = (q * cos) + (rotate_half(q) * sin) - k_embed = (k * cos) + (rotate_half(k) * sin) - return q_embed, k_embed - - @use_kernelized_func(apply_rotary_pos_emb) class Dots1Attention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" diff --git a/src/transformers/models/esm/configuration_esm.py b/src/transformers/models/esm/configuration_esm.py index 7875d88ecee8..a00dcf8b39e3 100644 --- a/src/transformers/models/esm/configuration_esm.py +++ b/src/transformers/models/esm/configuration_esm.py @@ -159,12 +159,12 @@ class EsmConfig(PreTrainedConfig): mask_token_id (`int`, *optional*): The index of the mask token in the vocabulary. This must be included in the config because of the "mask-dropout" scaling trick, which will scale the inputs depending on the number of masked tokens. - rope_theta (`float`, defaults to 10000.0): - The base period of the RoPE embeddings. Only used when `position_embedding_type` is set to `"rotary"`. position_embedding_type (`str`, *optional*, defaults to `"absolute"`): Type of position embedding. Choose either `"absolute"` or "rotary"`. emb_layer_norm_before (`bool`, *optional*): Whether to apply layer normalization after embeddings but before the main stem of the network. + rope_theta (`float`, defaults to 10000.0): + The base period of the RoPE embeddings. Only used when `position_embedding_type` is set to `"rotary"`. token_dropout (`bool`, defaults to `False`): When this is enabled, masked tokens are treated as if they had been dropped out by input dropout. is_folding_model (`bool`, defaults to `False`): diff --git a/src/transformers/models/nanochat/modeling_nanochat.py b/src/transformers/models/nanochat/modeling_nanochat.py index 4f5e1b7fe3c7..e4ea462797e9 100644 --- a/src/transformers/models/nanochat/modeling_nanochat.py +++ b/src/transformers/models/nanochat/modeling_nanochat.py @@ -122,6 +122,32 @@ def forward(self, x, position_ids): return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) +@use_kernel_func_from_hub("rotary_pos_emb") +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, @@ -159,32 +185,6 @@ def eager_attention_forward( return attn_output, attn_weights -@use_kernel_func_from_hub("rotary_pos_emb") -def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): - """Applies Rotary Position Embedding to the query and key tensors. - - Args: - q (`torch.Tensor`): The query tensor. - k (`torch.Tensor`): The key tensor. - cos (`torch.Tensor`): The cosine part of the rotary embedding. - sin (`torch.Tensor`): The sine part of the rotary embedding. - unsqueeze_dim (`int`, *optional*, defaults to 1): - The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and - sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note - that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and - k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes - cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have - the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. - Returns: - `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. - """ - cos = cos.unsqueeze(unsqueeze_dim) - sin = sin.unsqueeze(unsqueeze_dim) - q_embed = (q * cos) + (rotate_half(q) * sin) - k_embed = (k * cos) + (rotate_half(k) * sin) - return q_embed, k_embed - - def rotate_half(x): """Rotates half the hidden dims of the input with flipped signs for NanoChat.""" x1 = x[..., : x.shape[-1] // 2] diff --git a/src/transformers/models/qwen3/modeling_qwen3.py b/src/transformers/models/qwen3/modeling_qwen3.py index 2cf93f6f8ea2..b2e4a4bea863 100644 --- a/src/transformers/models/qwen3/modeling_qwen3.py +++ b/src/transformers/models/qwen3/modeling_qwen3.py @@ -27,7 +27,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import use_kernel_forward_from_hub, use_kernelized_func +from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -148,6 +148,39 @@ def forward(self, x, position_ids): return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +@use_kernel_func_from_hub("rotary_pos_emb") +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, @@ -185,20 +218,6 @@ def eager_attention_forward( return attn_output, attn_weights -def rotate_half(x): - x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] - return torch.cat((-x2, x1), dim=-1) - - -def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): - cos = cos.unsqueeze(unsqueeze_dim) - sin = sin.unsqueeze(unsqueeze_dim) - q_embed = (q * cos) + (rotate_half(q) * sin) - k_embed = (k * cos) + (rotate_half(k) * sin) - return q_embed, k_embed - - @use_kernelized_func(apply_rotary_pos_emb) class Qwen3Attention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" diff --git a/src/transformers/models/qwen3/modular_qwen3.py b/src/transformers/models/qwen3/modular_qwen3.py index c18e0e030079..73cde6d89a7a 100644 --- a/src/transformers/models/qwen3/modular_qwen3.py +++ b/src/transformers/models/qwen3/modular_qwen3.py @@ -34,6 +34,7 @@ Qwen2ForTokenClassification, Qwen2RMSNorm, Qwen2RotaryEmbedding, + apply_rotary_pos_emb, eager_attention_forward, ) from .configuration_qwen3 import Qwen3Config @@ -44,20 +45,6 @@ _CHECKPOINT_FOR_DOC = "Qwen/Qwen3-8B" -def rotate_half(x): - x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] - return torch.cat((-x2, x1), dim=-1) - - -def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): - cos = cos.unsqueeze(unsqueeze_dim) - sin = sin.unsqueeze(unsqueeze_dim) - q_embed = (q * cos) + (rotate_half(q) * sin) - k_embed = (k * cos) + (rotate_half(k) * sin) - return q_embed, k_embed - - class Qwen3RMSNorm(Qwen2RMSNorm): pass @@ -121,8 +108,6 @@ def forward( class Qwen3ForCausalLM(Qwen2ForCausalLM): - _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} - def forward( self, **super_kwargs: Unpack[TransformersKwargs], From 920ade5c2064849865e6299cef27792b0889bec1 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Thu, 14 May 2026 08:56:10 +0000 Subject: [PATCH 091/116] add tpu --- src/transformers/distributed/fsdp.py | 9 ++++++++- src/transformers/distributed/utils.py | 9 ++++++++- 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py index 9e8813bd2241..b499f1d88907 100644 --- a/src/transformers/distributed/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -103,7 +103,14 @@ def initialize_fsdp( local_rank = int(os.environ["LOCAL_RANK"]) world_size = int(os.environ["WORLD_SIZE"]) - backend_map = {"cuda": "nccl", "cpu": "gloo", "xpu": "xccl", "hpu": "hccl"} + backend_map = { + "cuda": "nccl", + "cpu": "gloo", + "xpu": "xccl", + "hpu": "hccl", + "neuron": "neuron", + "tpu": "tpu_dist", + } backend = backend_map.get(device_type) if device_type == "cpu" and int(os.environ.get("CCL_WORKER_COUNT", "0")): backend = "ccl" diff --git a/src/transformers/distributed/utils.py b/src/transformers/distributed/utils.py index e6bde120eeee..71d409d239b1 100644 --- a/src/transformers/distributed/utils.py +++ b/src/transformers/distributed/utils.py @@ -56,7 +56,14 @@ def _ensure_torch_distributed(device_type: str): local_rank = int(os.environ["LOCAL_RANK"]) world_size = int(os.environ["WORLD_SIZE"]) - backend_map = {"cuda": "nccl", "cpu": "gloo", "xpu": "xccl", "hpu": "hccl"} + 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 From 13646c8203b78d1ef8ff75528b736ab3521c2f2d Mon Sep 17 00:00:00 2001 From: 3outeille Date: Thu, 14 May 2026 09:03:34 +0000 Subject: [PATCH 092/116] move TP refactor experimentation scripts to backup branch Move ad-hoc training / verification / compare scripts off this branch into refactor-tp-dtensor-scripts so the diff stays focused on library changes. Co-Authored-By: Claude Opus 4.7 (1M context) --- compare_save_reload.sh | 33 -------- run_compare.sh | 67 --------------- run_verify_all.sh | 160 ----------------------------------- tmp_generate.py | 63 -------------- train_fsdp_tp.py | 157 ---------------------------------- train_save_reload.py | 186 ----------------------------------------- verify_loading.py | 137 ------------------------------ 7 files changed, 803 deletions(-) delete mode 100755 compare_save_reload.sh delete mode 100755 run_compare.sh delete mode 100755 run_verify_all.sh delete mode 100644 tmp_generate.py delete mode 100644 train_fsdp_tp.py delete mode 100644 train_save_reload.py delete mode 100644 verify_loading.py diff --git a/compare_save_reload.sh b/compare_save_reload.sh deleted file mode 100755 index 44de35d36f23..000000000000 --- a/compare_save_reload.sh +++ /dev/null @@ -1,33 +0,0 @@ -#!/bin/bash -# Compare the loss/grad_norm trajectory of `train_save_reload.py` with and without -# the mid-training save/reload. They should match step-for-step. -# -# Each mode writes to its own dir (`./checkpoints_baseline/`, `./checkpoints_save_reload/`) -# so the artifacts are inspectable side-by-side after the run. -set -euo pipefail - -NPROC="${NPROC:-4}" -BASELINE_LOG="${BASELINE_LOG:-baseline.log}" -SAVE_RELOAD_LOG="${SAVE_RELOAD_LOG:-save_reload.log}" -DIFF_LOG="${DIFF_LOG:-save_reload_diff.log}" - -rm -rf ./checkpoints_baseline ./checkpoints_save_reload - -# Pull out the per-step training lines and the post-reload generation lines so -# the diff is mechanical and covers both the training trajectory and the -# generated tokens/text. -filter_steps() { grep -E '^step |^# gen '; } - -echo "=== Run 1/2: --mode baseline ===" -torchrun --nproc_per_node="$NPROC" train_save_reload.py --mode baseline 2>&1 \ - | tee /dev/stderr | filter_steps > "$BASELINE_LOG" - -echo -echo "=== Run 2/2: --mode save_reload ===" -torchrun --nproc_per_node="$NPROC" train_save_reload.py --mode save_reload 2>&1 \ - | tee /dev/stderr | filter_steps > "$SAVE_RELOAD_LOG" - -echo -echo "=== Diff (baseline vs save_reload) ===" -git diff --no-index --color --word-diff=color "$BASELINE_LOG" "$SAVE_RELOAD_LOG" | tee "$DIFF_LOG" || true -echo "Diff written to $DIFF_LOG" diff --git a/run_compare.sh b/run_compare.sh deleted file mode 100755 index 4aec50851e10..000000000000 --- a/run_compare.sh +++ /dev/null @@ -1,67 +0,0 @@ -#!/bin/bash -set -euo pipefail - -SCRIPT="train_fsdp_tp.py" -LOG_FSDP_TP="log.txt" -LOG_FSDP_ONLY="ref.txt" -LOG_DIFF="diff.txt" - -MODEL_NAME="${MODEL_NAME:-hf-internal-testing/tiny-random-MixtralForCausalLM}" -COMMON_ARGS="--model_name $MODEL_NAME --lr 3e-4 --seed 42" - -rm -rf ./checkpoints_tp ./checkpoints_tp_resumed ./checkpoints_fsdp ./checkpoints_fsdp_resumed - -echo "=== Phase 1: Train steps 0-9, save checkpoint ===" -echo "--- Launching FSDP+TP and FSDP-only in parallel ---" - -CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nproc_per_node=4 --master_port=29500 \ - $SCRIPT $COMMON_ARGS --fsdp_size 2 --tp_size 2 --enable_sp \ - --start_step 0 --stop_step 10 --save_dir ./checkpoints_tp \ - --distributed_checkpoint > "${LOG_FSDP_TP}.phase1" 2>&1 & -PID1=$! - -CUDA_VISIBLE_DEVICES=4,5 torchrun --nproc_per_node=2 --master_port=29501 \ - $SCRIPT $COMMON_ARGS --fsdp_size 2 \ - --start_step 0 --stop_step 10 --save_dir ./checkpoints_fsdp > "${LOG_FSDP_ONLY}.phase1" 2>&1 & -PID2=$! - -echo "FSDP+TP PID=$PID1 | FSDP-only PID=$PID2" -wait $PID1 && echo "Phase 1 FSDP+TP done" || { echo "Phase 1 FSDP+TP failed (exit $?)"; cat "${LOG_FSDP_TP}.phase1"; exit 1; } -wait $PID2 && echo "Phase 1 FSDP-only done" || { echo "Phase 1 FSDP-only failed (exit $?)"; cat "${LOG_FSDP_ONLY}.phase1"; exit 1; } - -echo "" -echo "=== Phase 2: Resume from checkpoint, train steps 10-19, save ===" -echo "--- Launching FSDP+TP and FSDP-only in parallel ---" - -CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nproc_per_node=4 --master_port=29500 \ - $SCRIPT $COMMON_ARGS --fsdp_size 2 --tp_size 2 --enable_sp \ - --start_step 10 --stop_step 20 \ - --resume_dir ./checkpoints_tp --save_dir ./checkpoints_tp_resumed \ - --distributed_checkpoint > "${LOG_FSDP_TP}.phase2" 2>&1 & -PID1=$! - -CUDA_VISIBLE_DEVICES=4,5 torchrun --nproc_per_node=2 --master_port=29501 \ - $SCRIPT $COMMON_ARGS --fsdp_size 2 \ - --start_step 10 --stop_step 20 \ - --resume_dir ./checkpoints_fsdp --save_dir ./checkpoints_fsdp_resumed > "${LOG_FSDP_ONLY}.phase2" 2>&1 & -PID2=$! - -echo "FSDP+TP PID=$PID1 | FSDP-only PID=$PID2" -wait $PID1 && echo "Phase 2 FSDP+TP done" || { echo "Phase 2 FSDP+TP failed (exit $?)"; cat "${LOG_FSDP_TP}.phase2"; exit 1; } -wait $PID2 && echo "Phase 2 FSDP-only done" || { echo "Phase 2 FSDP-only failed (exit $?)"; cat "${LOG_FSDP_ONLY}.phase2"; exit 1; } - -# Combine phase logs, keeping only signal lines (loss/grad steps + checkpoint markers). -# Drops every kind of warning/progress noise: rank warnings, torchrun banners, -# transformers deprecations, tqdm progress bars, ProcessGroup teardown warnings, etc. -strip_warnings() { - grep -E '^(Step |Resumed |Saved )' -} -strip_warnings < "${LOG_FSDP_TP}.phase1" > "$LOG_FSDP_TP" -strip_warnings < "${LOG_FSDP_TP}.phase2" >> "$LOG_FSDP_TP" -strip_warnings < "${LOG_FSDP_ONLY}.phase1" > "$LOG_FSDP_ONLY" -strip_warnings < "${LOG_FSDP_ONLY}.phase2" >> "$LOG_FSDP_ONLY" - -echo "" -echo "=== Full Loss & Grad Diff (steps 0-19) ===" -git diff --no-index --color --word-diff=color "$LOG_FSDP_TP" "$LOG_FSDP_ONLY" | tee "$LOG_DIFF" || true -echo "Diff written to $LOG_DIFF" \ No newline at end of file diff --git a/run_verify_all.sh b/run_verify_all.sh deleted file mode 100755 index 3a1c9c08d8f9..000000000000 --- a/run_verify_all.sh +++ /dev/null @@ -1,160 +0,0 @@ -#!/bin/bash - -GREEN='\033[0;32m' -RED='\033[0;31m' -CYAN='\033[0;36m' -YELLOW='\033[1;33m' -BOLD='\033[1m' -DIM='\033[0;90m' -NC='\033[0m' - -SCRIPT="verify_loading.py" -LOGDIR="$(dirname "$0")/verify_logs" -mkdir -p "$LOGDIR" - -NUM_GPUS=$(nvidia-smi -L | wc -l) - -# Job definitions: "mode nproc_per_node" -declare -a JOBS=( - "single_gpu 1" - "fsdp 2" - "tp 2" - "tp_sp 2" - "tp_fsdp 4" - "tp_sp_fsdp 4" -) -MODE_NAMES=(single_gpu fsdp tp tp_sp tp_fsdp tp_sp_fsdp) - -echo -e "${BOLD}==========================================" -echo -e " Verify Loading (${NUM_GPUS} GPUs available)" -echo -e " Modes: ${MODE_NAMES[*]}" -echo -e " Logs: $LOGDIR/" -echo -e "==========================================${NC}" -echo "" - -# ============================================================ -# Round-robin GPU scheduler -# ============================================================ -NEXT_GPU=0 -MASTER_PORT=29500 -PIDS=() -PID_MODES=() - -for job in "${JOBS[@]}"; do - mode=${job% *} - nproc=${job#* } - - # Wait if not enough GPUs left in this round - if [ $((NEXT_GPU + nproc)) -gt "$NUM_GPUS" ]; then - echo -e "${DIM} (waiting for current round to finish...)${NC}" - for pid in "${PIDS[@]}"; do - wait "$pid" 2>/dev/null - done - PIDS=() - NEXT_GPU=0 - fi - - # Build CUDA_VISIBLE_DEVICES range - GPU_END=$((NEXT_GPU + nproc - 1)) - GPUS="" - for g in $(seq "$NEXT_GPU" "$GPU_END"); do - [ -n "$GPUS" ] && GPUS="${GPUS}," - GPUS="${GPUS}${g}" - done - - echo -e " ${CYAN}[${mode}]${NC} GPUs ${NEXT_GPU}-${GPU_END} (nproc=${nproc})" - - if [ "$nproc" -eq 1 ]; then - CUDA_VISIBLE_DEVICES="$GPUS" python "$SCRIPT" --mode "$mode" \ - > "$LOGDIR/${mode}.log" 2>&1 & - else - CUDA_VISIBLE_DEVICES="$GPUS" torchrun \ - --nproc_per_node="$nproc" --master_port="$MASTER_PORT" \ - "$SCRIPT" --mode "$mode" \ - > "$LOGDIR/${mode}.log" 2>&1 & - ((MASTER_PORT++)) - fi - - PIDS+=($!) - PID_MODES+=("$mode") - NEXT_GPU=$((GPU_END + 1)) -done - -# Wait for remaining jobs -echo "" -echo -e "${BOLD}Waiting for all jobs to finish...${NC}" -for i in "${!PIDS[@]}"; do - mode="${PID_MODES[$i]}" - if wait "${PIDS[$i]}"; then - echo -e " ${GREEN}✓${NC} ${mode}" - else - echo -e " ${RED}✗${NC} ${mode} (exit $?)" - fi -done - -# ============================================================ -# Results -# ============================================================ -echo "" -echo -e "${BOLD}=== Results ===${NC}" -for mode in "${MODE_NAMES[@]}"; do - log="$LOGDIR/$mode.log" - loss_before=$(grep -oP 'loss_before = \K[0-9.]+' "$log" 2>/dev/null) - loss_after=$(grep -oP 'loss_after = \K[0-9.]+' "$log" 2>/dev/null) - if grep -q '^PASS' "$log" 2>/dev/null; then - printf " ${GREEN}%-12s PASS (before=%-10s after=%s)${NC}\n" "$mode" "$loss_before" "$loss_after" - elif [ -n "$loss_before" ]; then - diff=$(grep -oP 'diff = \K[0-9.e+-]+' "$log" 2>/dev/null) - printf " ${RED}%-12s FAIL (before=%-10s after=%-10s diff=%s)${NC}\n" "$mode" "$loss_before" "$loss_after" "$diff" - else - printf " ${RED}%-12s ERROR (see log)${NC}\n" "$mode" - fi -done - -# ============================================================ -# Cross-mode loss comparison -# ============================================================ -echo "" -echo -e "${BOLD}=== Cross-mode loss comparison (PASS modes only) ===${NC}" -REF_LOSS="" -ALL_MATCH=1 -for mode in "${MODE_NAMES[@]}"; do - log="$LOGDIR/$mode.log" - # Only include modes where save/load roundtrip passed - if ! grep -q '^PASS' "$log" 2>/dev/null; then - continue - fi - loss=$(grep -oP 'loss_before = \K[0-9.]+' "$log" 2>/dev/null) - if [ -z "$loss" ]; then - continue - fi - if [ -z "$REF_LOSS" ]; then - REF_LOSS="$loss" - printf " ${GREEN}%-12s %s (reference)${NC}\n" "$mode" "$loss" - elif [ "$loss" = "$REF_LOSS" ]; then - printf " ${GREEN}%-12s %s${NC}\n" "$mode" "$loss" - else - printf " ${YELLOW}%-12s %s (differs from %s)${NC}\n" "$mode" "$loss" "$REF_LOSS" - ALL_MATCH=0 - fi -done -if [ "$ALL_MATCH" -eq 1 ] && [ -n "$REF_LOSS" ]; then - echo -e " ${GREEN}All modes produce the same loss.${NC}" -fi - -# Hints for failures -HAS_FAIL=0 -for mode in "${MODE_NAMES[@]}"; do - if ! grep -q '^PASS' "$LOGDIR/$mode.log" 2>/dev/null; then - HAS_FAIL=1 - fi -done -if [ "$HAS_FAIL" -eq 1 ]; then - echo "" - echo -e "${YELLOW}Some modes failed. Check logs:${NC}" - for mode in "${MODE_NAMES[@]}"; do - if ! grep -q '^PASS' "$LOGDIR/$mode.log" 2>/dev/null; then - echo -e " ${YELLOW}cat $LOGDIR/$mode.log${NC}" - fi - done -fi \ No newline at end of file diff --git a/tmp_generate.py b/tmp_generate.py deleted file mode 100644 index 8f2118771406..000000000000 --- a/tmp_generate.py +++ /dev/null @@ -1,63 +0,0 @@ -import argparse -import os - -import torch -from torch.distributed.elastic.multiprocessing.errors import record - -from transformers import AutoModelForCausalLM, AutoTokenizer -from transformers.distributed import DistributedConfig - -model_id = "mistralai/Mixtral-8x7B-Instruct-v0.1" -# model_id = "Qwen/Qwen3-14B" -# model_id = "Qwen/Qwen3-0.6B" -# model_id = "Qwen/Qwen1.5-MoE-A2.7B-Chat" -# model_id = "Qwen/Qwen3-30B-A3B-Instruct-2507" - -rank = int(os.environ["RANK"]) -world_size = int(os.environ["WORLD_SIZE"]) -device = torch.device(f"cuda:{rank}") -# Need to be initialized explicitly to use the `barrier` before loading -torch.distributed.init_process_group(backend="nccl", rank=rank, world_size=world_size, device_id=rank) - -@record -def main(args): - - distributed_config = DistributedConfig(tp_size=8, tp_plan="auto") - model = AutoModelForCausalLM.from_pretrained(model_id, distributed_config=distributed_config, dtype=torch.bfloat16) - # model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.bfloat16, device_map="auto") - tokenizer = AutoTokenizer.from_pretrained(model_id) - - messages = [ - {"role": "user", "content": "What do you think about life?"}, - ] - inputs = tokenizer.apply_chat_template(messages, return_tensors="pt").to(model.device) - input_size = inputs.input_ids.shape[-1] - - if args.profile: - # Warmup - with torch.no_grad(): - _ = model.generate(**inputs, max_new_tokens=5, do_sample=False) - - with torch.profiler.profile( - activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA], - record_shapes=True, - ) as prof: - output = model.generate(**inputs, max_new_tokens=2, do_sample=False) - - if rank == 0: - print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=30)) - prof.export_chrome_trace("trace.json") - else: - output = model.generate(**inputs, max_new_tokens=100, do_sample=False) - - text = tokenizer.batch_decode(output[:, input_size:])[0] - if rank == 0: - print(text) - -parser = argparse.ArgumentParser() -parser.add_argument("--profile", action="store_true") -args = parser.parse_args() - -main(args) - -torch.distributed.destroy_process_group() \ No newline at end of file diff --git a/train_fsdp_tp.py b/train_fsdp_tp.py deleted file mode 100644 index a2881c8bfaf5..000000000000 --- a/train_fsdp_tp.py +++ /dev/null @@ -1,157 +0,0 @@ -# torchrun --nproc_per_node=4 train_fsdp_tp.py - -import argparse -import os - -import torch -from datasets import load_dataset -from torch.distributed.tensor import DTensor -from torch.utils.data import DataLoader -from transformers import AutoModelForCausalLM, AutoTokenizer -from transformers.distributed import DistributedConfig -from transformers.distributed.utils import ( - _replicate_dtensor, - load_optimizer_distributed, - save_optimizer_distributed, -) - -def build_packed_dataset(dataset_name, tokenizer, seq_len, dp_rank, dp_world_size): - """Stream + tokenize + greedy-pack documents into fixed-length (input, label) windows.""" - ds = load_dataset(dataset_name, name="en", split="train", streaming=True) - ds = ds.shard(num_shards=dp_world_size, index=dp_rank) - buf, w = [], seq_len + 1 - - def pack(batch): - for t in batch["text"]: - buf.extend(tokenizer(t)["input_ids"]) - ids, lbls = [], [] - while len(buf) >= w: - ids.append(buf[:seq_len]); lbls.append(buf[1:w]); del buf[:w] - return {"input_ids": ids, "labels": lbls} - - ds = ds.map(pack, batched=True, remove_columns=ds.column_names) - return ds.with_format("torch") - -def build_fixed_batches(dp_rank): - """Load pre-generated fixed batches for a given DP rank.""" - return torch.load(f"fixed_batches_dp{dp_rank}.pt", weights_only=True) - -if __name__ == "__main__": - - parser = argparse.ArgumentParser() - parser.add_argument("--model_name", type=str, default="Qwen/Qwen3-0.6B") - parser.add_argument("--start_step", type=int, default=0, help="Inclusive start of the step range to train") - parser.add_argument("--stop_step", type=int, default=20, help="Exclusive end of the step range to train") - parser.add_argument("--lr", type=float, default=3e-4) - parser.add_argument("--seq_len", type=int, default=512) - parser.add_argument("--batch_size", type=int, default=1) - parser.add_argument("--save_dir", type=str, default="./checkpoints") - parser.add_argument("--tp_size", type=int, default=0, help="Tensor parallel size (0 = disabled)") - parser.add_argument("--fsdp_size", type=int, default=0, help="FSDP size (0 = disabled)") - parser.add_argument("--enable_sp", action="store_true", help="Enable sequence parallelism") - parser.add_argument("--seed", type=int, default=42, help="Random seed") - parser.add_argument("--fixed_batches", action="store_true", help="Use pre-generated fixed batches instead of C4") - parser.add_argument("--resume_dir", type=str, default=None, - help="Resume model + optimizer from a save_pretrained(distributed_checkpoint=True) dir") - parser.add_argument("--save_at_step", type=int, default=None, - help="Save a distributed checkpoint at this step number (inside [start_step, stop_step))") - parser.add_argument("--distributed_checkpoint", action="store_true", - help="Use distributed_checkpoint=True for the final save (per-rank shards via DCP + HF consolidation)") - args = parser.parse_args() - - torch.distributed.init_process_group(backend="nccl") - rank, local_rank = int(os.environ["RANK"]), int(os.environ["LOCAL_RANK"]) - torch.cuda.set_device(local_rank) - torch.manual_seed(args.seed) - - dc_kwargs = {} - if args.tp_size > 0: - dc_kwargs["tp_size"] = args.tp_size - dc_kwargs["tp_plan"] = "auto" - if args.fsdp_size > 0: - dc_kwargs["fsdp_size"] = args.fsdp_size - dc_kwargs["fsdp_plan"] = "auto" - if args.enable_sp: - dc_kwargs["enable_sequence_parallel"] = True - distributed_config = DistributedConfig(**dc_kwargs) - - # Both `args.model_name` (HF hub) and `args.resume_dir` (written by save_pretrained, - # canonical or distributed_checkpoint=True) are plain HF-format directories — same load path. - load_path = args.resume_dir if args.resume_dir else args.model_name - model = AutoModelForCausalLM.from_pretrained( - load_path, - distributed_config=distributed_config, - torch_dtype=torch.bfloat16, - ) - if args.resume_dir and rank == 0: - print(f"Resumed model from {args.resume_dir}") - - dp_rank = model.device_mesh["fsdp"].get_local_rank() if "fsdp" in model.device_mesh.mesh_dim_names else 0 - dp_size = model.device_mesh["fsdp"].size() if "fsdp" in model.device_mesh.mesh_dim_names else 1 - - if args.fixed_batches: - fixed = build_fixed_batches(dp_rank) - else: - tokenizer = AutoTokenizer.from_pretrained(args.model_name) - if tokenizer.pad_token is None: - tokenizer.pad_token = tokenizer.eos_token - dataset = build_packed_dataset("allenai/c4", tokenizer, args.seq_len, dp_rank, dp_size) - dataloader = iter(DataLoader(dataset, batch_size=args.batch_size)) - - optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr) - - if args.resume_dir: - optim_dir = os.path.join(args.resume_dir, "optimizer") - if os.path.exists(optim_dir): - load_optimizer_distributed(model, optimizer, optim_dir) - if rank == 0: - print(f"Resumed optimizer from {optim_dir}") - elif rank == 0: - print(f"No optimizer state at {optim_dir}; starting from fresh optimizer state") - - intermediate_dir = os.path.join(args.save_dir, "intermediate") - - model.train() - for step in range(args.start_step, args.stop_step): - if args.fixed_batches: - input_ids = fixed[step]["input_ids"].to(f"cuda:{local_rank}") - labels = fixed[step]["labels"].to(f"cuda:{local_rank}") - else: - batch = next(dataloader) - input_ids = batch["input_ids"].to(f"cuda:{local_rank}") - labels = batch["labels"].to(f"cuda:{local_rank}") - loss = model(input_ids, labels=labels).loss - loss.backward() - - # Custom grad clip: convert DTensor grads to local to avoid mixed-mesh torch.stack. - # Use _replicate_dtensor (not full_tensor) — full_tensor calls redistribute(), which - # cannot normalize FSDP+TP placements like (Shard(0), _StridedShard(1, sf=2)) from - # packed_colwise MoE experts. - grads = [p.grad for p in model.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=2.0) - torch.nn.utils.clip_grads_with_norm_(grads, max_norm=1.0, total_norm=total_norm) - optimizer.step() - optimizer.zero_grad() - - if rank == 0: - print(f"Step {step:>4d} | Loss: {loss.item():.4f} | Grad norm: {total_norm.item():.4f}") - - # Mid-training distributed checkpoint: every rank writes its own shard in parallel via DCP + - # HuggingFaceStorageWriter consolidation, plus DCP optimizer save. The resulting directory - # is still HF-safetensors-compatible, so `from_pretrained(intermediate_dir, ...)` resumes it. - if args.save_at_step is not None and step == args.save_at_step: - model.save_pretrained(intermediate_dir, distributed_checkpoint=True) - save_optimizer_distributed(model, optimizer, os.path.join(intermediate_dir, "optimizer")) - if rank == 0: - print(f"Saved distributed checkpoint at step {step} to {intermediate_dir}") - - # Final save: either canonical single-file HF safetensors (rank-0 gather) or distributed - # per-rank shards (DCP + HuggingFaceStorageWriter consolidation). Both are HF-format dirs - # that `from_pretrained` can resume from. Optimizer always saved via DCP. - model.save_pretrained(args.save_dir, distributed_checkpoint=args.distributed_checkpoint) - save_optimizer_distributed(model, optimizer, os.path.join(args.save_dir, "optimizer")) - if rank == 0: - print(f"Saved final checkpoint to {args.save_dir}") - - torch.distributed.destroy_process_group() \ No newline at end of file diff --git a/train_save_reload.py b/train_save_reload.py deleted file mode 100644 index 68a60c25b6df..000000000000 --- a/train_save_reload.py +++ /dev/null @@ -1,186 +0,0 @@ -"""Minimal save/reload demo: FSDP + TP on Isotonic/TinyMixtral-4x248M-MoE. - - torchrun --nproc_per_node=4 train_save_reload.py # save+reload - torchrun --nproc_per_node=4 train_save_reload.py --mode baseline # straight N steps - -The training loop deliberately overfits a single fixed sample so that, after -enough steps, the model memorizes it and `generate()` from a prefix produces -the rest of the sentence verbatim. To verify the save/reload round-trip is -lossless, run both modes and diff the loss / grad_norm logs *and* the final -generated token stream — they should all match step-for-step. -""" - -import argparse -import os - -import torch -from torch.distributed.tensor import DTensor - -from transformers import AutoModelForCausalLM, AutoTokenizer -from transformers.distributed import DistributedConfig -from transformers.distributed.utils import _replicate_dtensor, load_optimizer_distributed, save_optimizer_distributed - - -MODEL_NAME = "Isotonic/TinyMixtral-4x248M-MoE" -TOTAL_STEPS = 30 -HALFWAY = TOTAL_STEPS // 2 -LR = 1e-3 -SEED = 42 -BATCH_SIZE = 1 - -# A single passage long enough to tokenize to at least SEQ_LEN+1 tokens. The -# prompt below is a prefix of this; after overfitting, generation should -# reproduce the continuation verbatim. -OVERFIT_TEXT = ( - "In a quiet village nestled between rolling hills and a slow river, the " - "autumn mornings arrived with mist that hung low over the fields and a sky " - "that turned from grey to pale gold as the sun climbed." -) -GEN_PROMPT = "In a quiet village" - - -def run_phase(model, optimizer, batch_iter, local_rank, rank, start, stop): - for step in range(start, stop): - batch = next(batch_iter) - input_ids = batch["input_ids"].to(f"cuda:{local_rank}") - labels = batch["labels"].to(f"cuda:{local_rank}") - - loss = model(input_ids, labels=labels).loss - loss.backward() - - # Custom grad clip that tolerates DTensor grads with mixed placements: - # _replicate_dtensor handles _StridedShard (which redistribute() can't). - grads = [p.grad for p in model.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=2.0) - torch.nn.utils.clip_grads_with_norm_(grads, max_norm=1.0, total_norm=total_norm) - - optimizer.step() - optimizer.zero_grad() - - if rank == 0: - # Single canonical line per step so `diff` between modes is mechanical. - print(f"step {step:>3d} | loss {loss.item():.6f} | grad_norm {total_norm.item():.6f}") - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("--mode", choices=["save_reload", "baseline"], default="save_reload") - args = parser.parse_args() - - torch.distributed.init_process_group(backend="nccl") - rank = int(os.environ["RANK"]) - local_rank = int(os.environ["LOCAL_RANK"]) - torch.cuda.set_device(local_rank) - torch.manual_seed(SEED) - - distributed_config = DistributedConfig( - tp_size=2, - fsdp_size=2, - tp_plan="auto", - fsdp_plan="auto", - enable_sequence_parallel=True, - ) - - # Each mode writes to its own top-level directory so the artifacts of one run - # don't clobber the other and can be inspected side-by-side after the fact. - save_dir = f"./checkpoints_{args.mode}" - intermediate_dir = os.path.join(save_dir, "intermediate") - - if rank == 0: - print(f"# mode = {args.mode} | save_dir = {save_dir}") - - # Build the initial model + optimizer. - model = AutoModelForCausalLM.from_pretrained( - MODEL_NAME, - distributed_config=distributed_config, - torch_dtype=torch.bfloat16, - ) - optimizer = torch.optim.AdamW(model.parameters(), lr=LR) - model.train() - - tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) - if tokenizer.pad_token is None: - tokenizer.pad_token = tokenizer.eos_token - ids = tokenizer(OVERFIT_TEXT, return_tensors="pt").input_ids[0] - fixed_batch = { - "input_ids": ids.unsqueeze(0).to(f"cuda:{local_rank}"), - "labels": ids.unsqueeze(0).to(f"cuda:{local_rank}"), - } - - def fixed_iter(): - while True: - yield fixed_batch - - batch_iter = fixed_iter() - - if args.mode == "baseline": - run_phase(model, optimizer, batch_iter, local_rank, rank, 0, TOTAL_STEPS) - else: - # 1. Train first half. - run_phase(model, optimizer, batch_iter, local_rank, rank, 0, HALFWAY) - - # 2. Save (model via DCP→HF-format consolidation; optimizer via DCP). - model.save_pretrained(intermediate_dir, distributed_checkpoint=True) - save_optimizer_distributed(model, optimizer, os.path.join(intermediate_dir, "optimizer")) - if rank == 0: - print(f"# saved intermediate to {intermediate_dir}") - - # 3. Tear down + reload from disk. Note: the dataloader iterator stays alive across - # the boundary so batch indices line up with the baseline run. - del model, optimizer - torch.cuda.empty_cache() - model = AutoModelForCausalLM.from_pretrained( - intermediate_dir, - distributed_config=distributed_config, - torch_dtype=torch.bfloat16, - ) - optimizer = torch.optim.AdamW(model.parameters(), lr=LR) - load_optimizer_distributed(model, optimizer, os.path.join(intermediate_dir, "optimizer")) - model.train() - if rank == 0: - print(f"# reloaded model + optimizer from {intermediate_dir}") - - # 4. Train second half. - run_phase(model, optimizer, batch_iter, local_rank, rank, HALFWAY, TOTAL_STEPS) - - # Final save: canonical safetensors for the model + DCP for the optimizer. - model.save_pretrained(save_dir) - save_optimizer_distributed(model, optimizer, os.path.join(save_dir, "optimizer")) - if rank == 0: - print(f"# saved final model + optimizer to {save_dir}") - - del model, optimizer - torch.cuda.empty_cache() - - gen_distributed_config = DistributedConfig( - tp_size=4, - tp_plan="auto", - enable_sequence_parallel=False, - ) - model = AutoModelForCausalLM.from_pretrained( - save_dir, - distributed_config=gen_distributed_config, - torch_dtype=torch.bfloat16, - ) - model.eval() - inputs = tokenizer(GEN_PROMPT, return_tensors="pt").to(f"cuda:{local_rank}") - max_new = ids.numel() - inputs.input_ids.shape[-1] - with torch.no_grad(): - output_ids = model.generate(**inputs, max_new_tokens=max_new, do_sample=False) - - if rank == 0: - tokens = output_ids[0].tolist() - expected = ids.tolist() - print(f"# gen tokens: {tokens}") - print(f"# exp tokens: {expected}") - print(f"# gen text: {tokenizer.decode(tokens, skip_special_tokens=True)!r}") - print(f"# exp text: {tokenizer.decode(expected, skip_special_tokens=True)!r}") - assert tokens == expected, ( - f"generated tokens do not match OVERFIT_TEXT — " - f"first mismatch at index {next((i for i, (g, e) in enumerate(zip(tokens, expected)) if g != e), min(len(tokens), len(expected)))}" - ) - - torch.distributed.destroy_process_group() diff --git a/verify_loading.py b/verify_loading.py deleted file mode 100644 index ba0f60f31fd3..000000000000 --- a/verify_loading.py +++ /dev/null @@ -1,137 +0,0 @@ -# Save/load roundtrip test for distributed models (TP, FSDP, TP+FSDP). -# -# Verifies that save_pretrained → from_pretrained preserves model weights by -# checking that the cross-entropy loss is identical before and after the roundtrip. -# This catches bugs in DTensor gather-on-save and shard-on-read paths. -# -# Usage: -# python verify_loading.py --mode single_gpu -# torchrun --nproc_per_node=2 verify_loading.py --mode fsdp -# torchrun --nproc_per_node=2 verify_loading.py --mode tp -# torchrun --nproc_per_node=4 verify_loading.py --mode tp_fsdp -# MODEL=Qwen/Qwen3-0.6B torchrun --nproc_per_node=2 verify_loading.py --mode tp -import argparse -import os -import shutil - -import torch -from torch.distributed.tensor import DTensor, Replicate - -from transformers import AutoModelForCausalLM, AutoTokenizer -from transformers.distributed import DistributedConfig - - -parser = argparse.ArgumentParser() -parser.add_argument("--mode", choices=["single_gpu", "fsdp", "tp", "tp_sp", "tp_fsdp", "tp_sp_fsdp"], required=True) -parser.add_argument("--model", type=str, default=None, help="Model ID (or set MODEL env var)") -args = parser.parse_args() - -model_id = args.model or os.environ.get("MODEL") or os.environ.get("MODEL_ID") or "hf-internal-testing/tiny-random-MixtralForCausalLM" - -if args.mode != "single_gpu": - torch.distributed.init_process_group(backend="nccl") - rank = int(os.environ["RANK"]) - local_rank = int(os.environ["LOCAL_RANK"]) - torch.cuda.set_device(local_rank) -else: - rank = 0 - local_rank = 0 - torch.cuda.set_device(0) - -configs = { - "single_gpu": lambda: None, - "fsdp": lambda: DistributedConfig(fsdp_size=2, fsdp_plan="auto"), - "tp": lambda: DistributedConfig(tp_size=2, tp_plan="auto"), - "tp_sp": lambda: DistributedConfig(tp_size=2, tp_plan="auto", enable_sequence_parallel=True), - "tp_fsdp": lambda: DistributedConfig(tp_size=2, tp_plan="auto", fsdp_size=2, fsdp_plan="auto"), - "tp_sp_fsdp": lambda: DistributedConfig(tp_size=2, tp_plan="auto", fsdp_size=2, fsdp_plan="auto", enable_sequence_parallel=True), -} - -tokenizer = AutoTokenizer.from_pretrained(model_id) -text = "The capital of France is Paris. The largest ocean is the Pacific." - - -def materialize_full_logits(logits: torch.Tensor) -> torch.Tensor: - if isinstance(logits, DTensor): - with torch.no_grad(): - return logits.redistribute(placements=[Replicate()] * logits.device_mesh.ndim, async_op=False).to_local() - return logits - - -def compute_loss(model): - inputs = tokenizer(text, return_tensors="pt").to(f"cuda:{local_rank}") - input_ids = inputs["input_ids"] - # Pad sequence length to a multiple of tp_size so DTensor Shard(1) splits evenly - # across ranks in SP mode. Always pad (even for non-TP modes) so that all modes - # compute on the same input and losses are directly comparable. - max_tp = 2 # all TP configs use tp_size=2 - seq_len = input_ids.shape[1] - if seq_len % max_tp != 0: - pad_len = max_tp - (seq_len % max_tp) - pad_token_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id - input_ids = torch.cat([input_ids, input_ids.new_full((1, pad_len), pad_token_id)], dim=1) - labels = input_ids.clone() - labels[:, seq_len:] = -100 # ignore padding in loss - position_ids = torch.arange(input_ids.shape[1], device=input_ids.device).unsqueeze(0) - - model.eval() - with torch.no_grad(): - logits = model(input_ids, position_ids=position_ids).logits - logits = materialize_full_logits(logits) - loss = torch.nn.functional.cross_entropy( - logits.flatten(0, 1).float(), - labels.flatten(0, 1), - reduction="mean", - ignore_index=-100, - ) - return loss.item() - - -# --- Step 1: Load original model and compute loss --- -model = AutoModelForCausalLM.from_pretrained(model_id, distributed_config=configs[args.mode](), dtype=torch.float32) -if args.mode == "single_gpu": - model = model.to("cuda:0") - -loss_before = compute_loss(model) -if rank == 0: - print(f"{args.mode}: loss_before = {loss_before:.6f}") - -# --- Step 2: Save to local dir (shared path across ranks) --- -save_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), f"verify_ckpt_{args.mode}") -if rank == 0: - if os.path.exists(save_dir): - shutil.rmtree(save_dir) - os.makedirs(save_dir) -if args.mode != "single_gpu": - torch.distributed.barrier() -model.save_pretrained(save_dir, is_main_process=(rank == 0)) -if rank == 0: - print(f"{args.mode}: saved to {save_dir}") - -# Ensure all ranks see the saved files before reloading -if args.mode != "single_gpu": - torch.distributed.barrier() - -del model -torch.cuda.empty_cache() - -# --- Step 3: Reload from saved checkpoint and compute loss --- -model2 = AutoModelForCausalLM.from_pretrained(save_dir, distributed_config=configs[args.mode](), dtype=torch.float32) -if args.mode == "single_gpu": - model2 = model2.to("cuda:0") - -loss_after = compute_loss(model2) -if rank == 0: - print(f"{args.mode}: loss_after = {loss_after:.6f}") - -# --- Step 4: Compare --- -if rank == 0: - diff = abs(loss_before - loss_after) - print(f"{args.mode}: diff = {diff:.2e}") - if diff < 1e-5: - print("PASS: save/load roundtrip is lossless") - else: - print("FAIL: loss mismatch after save/load roundtrip!") - -if args.mode != "single_gpu": - torch.distributed.destroy_process_group() From c49e9abe710d2f1cb13f69ab8235a423358f7845 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Thu, 14 May 2026 09:10:21 +0000 Subject: [PATCH 093/116] linting --- src/transformers/core_model_loading.py | 2 +- src/transformers/modeling_utils.py | 2 +- .../deepseek_v4/modeling_deepseek_v4.py | 3 +- .../models/eurobert/configuration_eurobert.py | 19 ++++- .../configuration_granite4_vision.py | 19 ++++- .../models/hy_v3/modeling_hy_v3.py | 3 +- .../hyperclovax/configuration_hyperclovax.py | 19 ++++- .../hyperclovax/modeling_hyperclovax.py | 3 +- .../models/laguna/modeling_laguna.py | 3 +- .../models/qwen3_5/modeling_qwen3_5.py | 74 +++++++++---------- .../qwen3_5_moe/modeling_qwen3_5_moe.py | 74 +++++++++---------- .../models/qwen3_moe/modeling_qwen3_moe.py | 54 ++++++++++---- .../models/qwen3_next/modeling_qwen3_next.py | 74 +++++++++---------- .../qwen3_omni_moe/modeling_qwen3_omni_moe.py | 26 ++++++- .../models/qwen3_vl/modeling_qwen3_vl.py | 21 +++++- .../qwen3_vl_moe/modeling_qwen3_vl_moe.py | 26 ++++++- 16 files changed, 281 insertions(+), 141 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 039f3daac9e5..6ce8c43d0d56 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -604,7 +604,7 @@ def __init__(self, source_patterns: str | list[str], target_patterns: str | list self._original_target_patterns = self.target_patterns.copy() # Init fields that will be used during conversion - self.distributed_operation: TensorParallelLayer | None = None + self.distributed_operation: Any = None self.quantization_operation: ConversionOps | None = None self.collected_tensors: dict[str, list[Future]] = defaultdict(list) self.layer_targets: dict[str, set[str]] = defaultdict(set) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index de6adb151b4e..81f5435251e8 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -98,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 diff --git a/src/transformers/models/deepseek_v4/modeling_deepseek_v4.py b/src/transformers/models/deepseek_v4/modeling_deepseek_v4.py index 163a0cee77a3..4f93c0dadea0 100644 --- a/src/transformers/models/deepseek_v4/modeling_deepseek_v4.py +++ b/src/transformers/models/deepseek_v4/modeling_deepseek_v4.py @@ -1395,7 +1395,8 @@ def load_balancing_loss_func( @auto_docstring class DeepseekV4ForCausalLM(DeepseekV4PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/eurobert/configuration_eurobert.py b/src/transformers/models/eurobert/configuration_eurobert.py index f64c4f7e5a11..b4b4a0511c41 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"]), diff --git a/src/transformers/models/granite4_vision/configuration_granite4_vision.py b/src/transformers/models/granite4_vision/configuration_granite4_vision.py index 82c9e6765515..8d9477dad556 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"]), diff --git a/src/transformers/models/hy_v3/modeling_hy_v3.py b/src/transformers/models/hy_v3/modeling_hy_v3.py index f2d64736b4f0..92499a1d9609 100644 --- a/src/transformers/models/hy_v3/modeling_hy_v3.py +++ b/src/transformers/models/hy_v3/modeling_hy_v3.py @@ -544,7 +544,8 @@ def forward( @auto_docstring class HYV3ForCausalLM(HYV3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: HYV3Config): diff --git a/src/transformers/models/hyperclovax/configuration_hyperclovax.py b/src/transformers/models/hyperclovax/configuration_hyperclovax.py index 430a56bf0249..b0b4c64d11bd 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"]), diff --git a/src/transformers/models/hyperclovax/modeling_hyperclovax.py b/src/transformers/models/hyperclovax/modeling_hyperclovax.py index 3608d215bfa9..f314eadb363d 100644 --- a/src/transformers/models/hyperclovax/modeling_hyperclovax.py +++ b/src/transformers/models/hyperclovax/modeling_hyperclovax.py @@ -452,7 +452,8 @@ def forward( @auto_docstring class HyperCLOVAXForCausalLM(HyperCLOVAXPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/laguna/modeling_laguna.py b/src/transformers/models/laguna/modeling_laguna.py index aa4060e77f5f..d3796f6e438e 100644 --- a/src/transformers/models/laguna/modeling_laguna.py +++ b/src/transformers/models/laguna/modeling_laguna.py @@ -672,7 +672,8 @@ def load_balancing_loss_func( @auto_docstring class LagunaForCausalLM(LagunaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/qwen3_5/modeling_qwen3_5.py b/src/transformers/models/qwen3_5/modeling_qwen3_5.py index dbee24588d74..2e9e6448a4ea 100644 --- a/src/transformers/models/qwen3_5/modeling_qwen3_5.py +++ b/src/transformers/models/qwen3_5/modeling_qwen3_5.py @@ -557,43 +557,6 @@ def forward( return output -def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: - """ - This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, - num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) - """ - batch, num_key_value_heads, slen, head_dim = hidden_states.shape - if n_rep == 1: - return hidden_states - hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) - return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) - - -def eager_attention_forward( - module: nn.Module, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attention_mask: torch.Tensor | None, - scaling: float, - dropout: float = 0.0, - **kwargs: Unpack[TransformersKwargs], -): - key_states = repeat_kv(key, module.num_key_value_groups) - value_states = repeat_kv(value, module.num_key_value_groups) - - attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling - if attention_mask is not None: - attn_weights = attn_weights + attention_mask - - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) - attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) - attn_output = torch.matmul(attn_weights, value_states) - attn_output = attn_output.transpose(1, 2).contiguous() - - return attn_output, attn_weights - - def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] @@ -640,6 +603,43 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): return q_embed, k_embed +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + @use_kernelized_func(apply_rotary_pos_emb) class Qwen3_5Attention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" 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 1c03f3ea6df7..09e5a78e6799 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 @@ -553,43 +553,6 @@ def forward( return output -def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: - """ - This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, - num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) - """ - batch, num_key_value_heads, slen, head_dim = hidden_states.shape - if n_rep == 1: - return hidden_states - hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) - return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) - - -def eager_attention_forward( - module: nn.Module, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attention_mask: torch.Tensor | None, - scaling: float, - dropout: float = 0.0, - **kwargs: Unpack[TransformersKwargs], -): - key_states = repeat_kv(key, module.num_key_value_groups) - value_states = repeat_kv(value, module.num_key_value_groups) - - attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling - if attention_mask is not None: - attn_weights = attn_weights + attention_mask - - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) - attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) - attn_output = torch.matmul(attn_weights, value_states) - attn_output = attn_output.transpose(1, 2).contiguous() - - return attn_output, attn_weights - - def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] @@ -636,6 +599,43 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): return q_embed, k_embed +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + @use_kernelized_func(apply_rotary_pos_emb) class Qwen3_5MoeAttention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" diff --git a/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py b/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py index 54b7ffe1d167..404869f16bfb 100644 --- a/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py +++ b/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py @@ -29,7 +29,12 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func +from ...integrations import ( + use_experts_implementation, + use_kernel_forward_from_hub, + use_kernel_func_from_hub, + use_kernelized_func, +) from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( @@ -48,6 +53,39 @@ from .configuration_qwen3_moe import Qwen3MoeConfig +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +@use_kernel_func_from_hub("rotary_pos_emb") +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, @@ -85,20 +123,6 @@ def eager_attention_forward( return attn_output, attn_weights -def rotate_half(x): - x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] - return torch.cat((-x2, x1), dim=-1) - - -def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): - cos = cos.unsqueeze(unsqueeze_dim) - sin = sin.unsqueeze(unsqueeze_dim) - q_embed = (q * cos) + (rotate_half(q) * sin) - k_embed = (k * cos) + (rotate_half(k) * sin) - return q_embed, k_embed - - @use_kernelized_func(apply_rotary_pos_emb) class Qwen3MoeAttention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" diff --git a/src/transformers/models/qwen3_next/modeling_qwen3_next.py b/src/transformers/models/qwen3_next/modeling_qwen3_next.py index 197e71cf052e..834eb6dbc241 100644 --- a/src/transformers/models/qwen3_next/modeling_qwen3_next.py +++ b/src/transformers/models/qwen3_next/modeling_qwen3_next.py @@ -169,43 +169,6 @@ def extra_repr(self): return f"{tuple(self.weight.shape)}, eps={self.eps}" -def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: - """ - This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, - num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) - """ - batch, num_key_value_heads, slen, head_dim = hidden_states.shape - if n_rep == 1: - return hidden_states - hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) - return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) - - -def eager_attention_forward( - module: nn.Module, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attention_mask: torch.Tensor | None, - scaling: float, - dropout: float = 0.0, - **kwargs: Unpack[TransformersKwargs], -): - key_states = repeat_kv(key, module.num_key_value_groups) - value_states = repeat_kv(value, module.num_key_value_groups) - - attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling - if attention_mask is not None: - attn_weights = attn_weights + attention_mask - - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) - attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) - attn_output = torch.matmul(attn_weights, value_states) - attn_output = attn_output.transpose(1, 2).contiguous() - - return attn_output, attn_weights - - def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] @@ -252,6 +215,43 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): return q_embed, k_embed +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + @use_kernelized_func(apply_rotary_pos_emb) class Qwen3NextAttention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" 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 2a9c2c0fd54a..a93e2c9df284 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 @@ -35,7 +35,12 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func +from ...integrations import ( + use_experts_implementation, + use_kernel_forward_from_hub, + use_kernel_func_from_hub, + use_kernelized_func, +) from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -891,6 +896,7 @@ def padded_and_mask_function(self, tensor_list, tensor_len, padding_value=0, pad def rotate_half(x): + """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) @@ -1435,7 +1441,25 @@ def extra_repr(self): return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" +@use_kernel_func_from_hub("rotary_pos_emb") def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) q_embed = (q * cos) + (rotate_half(q) * sin) diff --git a/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py b/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py index dad54d8ae7c7..843d8c8570e8 100644 --- a/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py +++ b/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py @@ -31,7 +31,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import use_kernel_forward_from_hub, use_kernelized_func +from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -126,6 +126,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: def rotate_half(x): + """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) @@ -409,7 +410,25 @@ def extra_repr(self): return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" +@use_kernel_func_from_hub("rotary_pos_emb") def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) q_embed = (q * cos) + (rotate_half(q) * sin) diff --git a/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py b/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py index 33aba14a29da..5886919e459c 100644 --- a/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py +++ b/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py @@ -32,7 +32,12 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func +from ...integrations import ( + use_experts_implementation, + use_kernel_forward_from_hub, + use_kernel_func_from_hub, + use_kernelized_func, +) from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -147,6 +152,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: def rotate_half(x): + """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) @@ -189,7 +195,25 @@ def eager_attention_forward( return attn_output, attn_weights +@use_kernel_func_from_hub("rotary_pos_emb") def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) q_embed = (q * cos) + (rotate_half(q) * sin) From a4c6ba8e6c8ab26edcd2e6a956b59fa08950d65c Mon Sep 17 00:00:00 2001 From: 3outeille Date: Thu, 14 May 2026 09:42:11 +0000 Subject: [PATCH 094/116] register distributed sharding_utils and utils in __init__ Co-Authored-By: Claude Opus 4.7 (1M context) --- src/transformers/__init__.py | 1 + src/transformers/distributed/__init__.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/src/transformers/__init__.py b/src/transformers/__init__.py index e949386f2dc5..bc71ecafd36b 100755 --- a/src/transformers/__init__.py +++ b/src/transformers/__init__.py @@ -104,6 +104,7 @@ "debug_utils": [], "dependency_versions_check": [], "dependency_versions_table": [], + "distributed": [], "dynamic_module_utils": [], "feature_extraction_sequence_utils": ["SequenceFeatureExtractor"], "feature_extraction_utils": ["BatchFeature", "FeatureExtractionMixin"], diff --git a/src/transformers/distributed/__init__.py b/src/transformers/distributed/__init__.py index fbb12304576f..09f81832c8e0 100644 --- a/src/transformers/distributed/__init__.py +++ b/src/transformers/distributed/__init__.py @@ -20,11 +20,13 @@ _import_structure = { "configuration_utils": ["DistributedConfig"], "fsdp": ["is_fsdp_enabled", "is_fsdp_managed_module"], + "sharding_utils": [], "tensor_parallel": [ "ALL_PARALLEL_STYLES", "apply_tensor_parallel", "verify_tp_plan", ], + "utils": [], } From 65b0311a9fe32bd5035946e75c91aa6da41be661 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Thu, 14 May 2026 09:42:21 +0000 Subject: [PATCH 095/116] rename TP plan styles to match new ALL_PARALLEL_STYLES registry Replace pre-refactor names that no longer exist in src/transformers/distributed/tensor_parallel.py: rowwise -> rowwise_allreduce moe_tp_experts -> moe_experts_allreduce replicated_with_grad_allreduce -> activation_seq_dim_2 Co-Authored-By: Claude Opus 4.7 (1M context) --- .../deepseek_v4/configuration_deepseek_v4.py | 2 +- .../models/gemma4/configuration_gemma4.py | 4 ++-- .../models/gpt_oss/configuration_gpt_oss.py | 2 +- .../models/hy_v3/configuration_hy_v3.py | 14 +++++++------- src/transformers/models/hy_v3/modular_hy_v3.py | 14 +++++++------- .../models/laguna/configuration_laguna.py | 14 +++++++------- src/transformers/models/laguna/modular_laguna.py | 14 +++++++------- .../models/llama4/configuration_llama4.py | 4 ++-- .../configuration_openai_privacy_filter.py | 2 +- .../models/qwen3_moe/configuration_qwen3_moe.py | 2 +- .../qwen3_omni_moe/configuration_qwen3_omni_moe.py | 2 +- .../qwen3_omni_moe/modular_qwen3_omni_moe.py | 2 +- .../qwen3_vl_moe/configuration_qwen3_vl_moe.py | 2 +- .../models/qwen3_vl_moe/modular_qwen3_vl_moe.py | 2 +- .../configuration_voxtral_realtime.py | 4 ++-- 15 files changed, 42 insertions(+), 42 deletions(-) diff --git a/src/transformers/models/deepseek_v4/configuration_deepseek_v4.py b/src/transformers/models/deepseek_v4/configuration_deepseek_v4.py index 8f8d818d8ced..3405f03099a1 100644 --- a/src/transformers/models/deepseek_v4/configuration_deepseek_v4.py +++ b/src/transformers/models/deepseek_v4/configuration_deepseek_v4.py @@ -125,7 +125,7 @@ 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", } vocab_size: int = 129280 diff --git a/src/transformers/models/gemma4/configuration_gemma4.py b/src/transformers/models/gemma4/configuration_gemma4.py index de6ffea2658a..ebb6ccab2d76 100644 --- a/src/transformers/models/gemma4/configuration_gemma4.py +++ b/src/transformers/models/gemma4/configuration_gemma4.py @@ -139,11 +139,11 @@ class Gemma4TextConfig(PreTrainedConfig): # EP plan for google/gemma-4-26B-A4B-it: do not tp in attention (num_global_key_value_heads=2 too small to partition) "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", "layers.*.router": "ep_router", "layers.*.experts.gate_up_proj": "grouped_gemm", "layers.*.experts.down_proj": "grouped_gemm", - "layers.*.experts": "moe_tp_experts", + "layers.*.experts": "moe_experts_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/gpt_oss/configuration_gpt_oss.py b/src/transformers/models/gpt_oss/configuration_gpt_oss.py index 47c029a5bca9..3a9ca00c8e9d 100644 --- a/src/transformers/models/gpt_oss/configuration_gpt_oss.py +++ b/src/transformers/models/gpt_oss/configuration_gpt_oss.py @@ -38,7 +38,7 @@ class GptOssConfig(PreTrainedConfig): "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/hy_v3/configuration_hy_v3.py b/src/transformers/models/hy_v3/configuration_hy_v3.py index e2eee94b118a..9ed4c5cc81c6 100644 --- a/src/transformers/models/hy_v3/configuration_hy_v3.py +++ b/src/transformers/models/hy_v3/configuration_hy_v3.py @@ -54,18 +54,18 @@ class HYV3Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts.down_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/hy_v3/modular_hy_v3.py b/src/transformers/models/hy_v3/modular_hy_v3.py index fa0931435197..0f63e5b32f59 100644 --- a/src/transformers/models/hy_v3/modular_hy_v3.py +++ b/src/transformers/models/hy_v3/modular_hy_v3.py @@ -79,18 +79,18 @@ class HYV3Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts.down_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/laguna/configuration_laguna.py b/src/transformers/models/laguna/configuration_laguna.py index 33f939f6db43..dfe403281263 100644 --- a/src/transformers/models/laguna/configuration_laguna.py +++ b/src/transformers/models/laguna/configuration_laguna.py @@ -61,18 +61,18 @@ class LagunaConfig(PreTrainedConfig): "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.g_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts.down_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/laguna/modular_laguna.py b/src/transformers/models/laguna/modular_laguna.py index 945cd40a99b2..7587827acb2c 100644 --- a/src/transformers/models/laguna/modular_laguna.py +++ b/src/transformers/models/laguna/modular_laguna.py @@ -79,18 +79,18 @@ class LagunaConfig(Qwen2MoeConfig): "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.g_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts.down_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", } vocab_size: int = 100352 diff --git a/src/transformers/models/llama4/configuration_llama4.py b/src/transformers/models/llama4/configuration_llama4.py index bf809aeeb6c7..781de8a57fd0 100644 --- a/src/transformers/models/llama4/configuration_llama4.py +++ b/src/transformers/models/llama4/configuration_llama4.py @@ -125,12 +125,12 @@ 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.experts.gate_up_proj": "grouped_gemm", # row because not linear "layers.*.feed_forward.experts.down_proj": "grouped_gemm", # col because not linear "layers.*.feed_forward.gate_proj": "colwise", "layers.*.feed_forward.up_proj": "colwise", - "layers.*.feed_forward.down_proj": "rowwise", + "layers.*.feed_forward.down_proj": "rowwise_allreduce", "layers.*.feed_forward.router": "ep_router", } 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..ae833c2bf514 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 @@ -62,7 +62,7 @@ class OpenAIPrivacyFilterConfig(PreTrainedConfig): "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 diff --git a/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py b/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py index ee81111327d1..7009dfe8f835 100644 --- a/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py +++ b/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py @@ -88,7 +88,7 @@ class Qwen3MoeConfig(PreTrainedConfig): "layers.*.mlp.gate": "ep_router", "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", "layers.*.mlp.experts.down_proj": "grouped_gemm", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts": "moe_experts_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), 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 b7f6fda36ffb..fc9fa220c699 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 @@ -407,7 +407,7 @@ class Qwen3OmniMoeTalkerTextConfig(PreTrainedConfig): "layers.*.mlp.gate": "ep_router", "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", "layers.*.mlp.experts.down_proj": "grouped_gemm", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts": "moe_experts_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), 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 b180d50c43c9..2864f1f325ba 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 @@ -392,7 +392,7 @@ class Qwen3OmniMoeTalkerTextConfig(Qwen3MoeConfig): "layers.*.mlp.gate": "ep_router", "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", "layers.*.mlp.experts.down_proj": "grouped_gemm", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts": "moe_experts_allreduce", } vocab_size: int = 3072 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 54ed15758b2d..0b4573b99942 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 @@ -86,7 +86,7 @@ class Qwen3VLMoeTextConfig(PreTrainedConfig): "layers.*.mlp.gate": "ep_router", "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", "layers.*.mlp.experts.down_proj": "grouped_gemm", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts": "moe_experts_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), 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 5b893a639e6d..555e3d63ccb3 100644 --- a/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py +++ b/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py @@ -95,7 +95,7 @@ class Qwen3VLMoeTextConfig(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", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/voxtral_realtime/configuration_voxtral_realtime.py b/src/transformers/models/voxtral_realtime/configuration_voxtral_realtime.py index b0227b418771..48d1035b7243 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"]), From dbf0c60913274675713eb46546d5e12e18ce4fe8 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Thu, 14 May 2026 09:45:39 +0000 Subject: [PATCH 096/116] enable EP --- src/transformers/distributed/configuration_utils.py | 1 + tests/test_tensor_parallel_mixin.py | 7 ++++++- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/src/transformers/distributed/configuration_utils.py b/src/transformers/distributed/configuration_utils.py index d3f76b72442e..386f66cc4b73 100644 --- a/src/transformers/distributed/configuration_utils.py +++ b/src/transformers/distributed/configuration_utils.py @@ -42,6 +42,7 @@ class DistributedConfig: tp_size: int | None = None tp_plan: str | dict[str, str] | None = None enable_sequence_parallel: bool = False + enable_expert_parallel: bool = False fsdp_size: int | None = None fsdp_plan: str | dict | None = None diff --git a/tests/test_tensor_parallel_mixin.py b/tests/test_tensor_parallel_mixin.py index f6e70cc3ca80..b43796ef4304 100644 --- a/tests/test_tensor_parallel_mixin.py +++ b/tests/test_tensor_parallel_mixin.py @@ -402,9 +402,14 @@ def _test_tp_generation_quantized_impl(_rank, model_path, model_class, max_new_t def _load_ep_and_reference_models(model_path, model_class): """Load EP model and non-EP reference model for comparison.""" + tp_size = dist.get_world_size() model_ep = model_class.from_pretrained( model_path, - distributed_config=DistributedConfig(enable_expert_parallel=True), + distributed_config=DistributedConfig( + tp_size=tp_size, + tp_plan="auto", + enable_expert_parallel=True, + ), ) dist.barrier() From 51068ca99df33a9a232dcdf42e5195c10d13de0d Mon Sep 17 00:00:00 2001 From: 3outeille Date: Thu, 14 May 2026 10:05:04 +0000 Subject: [PATCH 097/116] Add enable_expert_parallel configuration option in test_distributed_config --- tests/test_distributed_config.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_distributed_config.py b/tests/test_distributed_config.py index 6057d79d6dcc..abffeb15d867 100644 --- a/tests/test_distributed_config.py +++ b/tests/test_distributed_config.py @@ -54,6 +54,7 @@ def test_to_dict(self): "tp_size": 2, "tp_plan": "auto", "enable_sequence_parallel": False, + "enable_expert_parallel": False, "fsdp_size": 4, "fsdp_plan": "auto", } From bf0696f5a7d5199e90155a5d59c188e5b314fb48 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Sat, 16 May 2026 09:18:53 +0000 Subject: [PATCH 098/116] no more auto mode --- src/transformers/configuration_utils.py | 4 + .../distributed/configuration_utils.py | 16 +- src/transformers/distributed/fsdp.py | 293 ++++++------------ .../distributed/tensor_parallel.py | 13 +- src/transformers/modeling_utils.py | 27 +- .../deepseek_v4/modeling_deepseek_v4.py | 1 + src/transformers/models/doge/modeling_doge.py | 1 + .../models/dots1/modeling_dots1.py | 1 + .../ernie4_5_moe/modeling_ernie4_5_moe.py | 1 + .../models/flex_olmo/modeling_flex_olmo.py | 1 + .../models/gpt_oss/modeling_gpt_oss.py | 1 + .../models/granitemoe/modeling_granitemoe.py | 1 + .../modeling_granitemoehybrid.py | 1 + .../modeling_granitemoeshared.py | 1 + .../models/jamba/modeling_jamba.py | 1 + .../models/laguna/modeling_laguna.py | 1 + .../models/minimax/modeling_minimax.py | 1 + .../models/minimax_m2/modeling_minimax_m2.py | 1 + .../models/mixtral/configuration_mixtral.py | 7 + .../models/mixtral/modeling_mixtral.py | 1 + .../models/mixtral/modular_mixtral.py | 1 + .../models/olmoe/modeling_olmoe.py | 1 + .../models/phimoe/modeling_phimoe.py | 1 + .../models/qwen2_moe/modeling_qwen2_moe.py | 1 + .../models/qwen3/configuration_qwen3.py | 9 + .../models/qwen3/modeling_qwen3.py | 1 + .../models/qwen3/modular_qwen3.py | 2 + .../models/qwen3_5/modeling_qwen3_5.py | 1 + .../qwen3_5_moe/modeling_qwen3_5_moe.py | 1 + .../models/qwen3_moe/modeling_qwen3_moe.py | 1 + .../models/qwen3_next/modeling_qwen3_next.py | 1 + .../configuration_qwen3_omni_moe.py | 9 + .../qwen3_omni_moe/modeling_qwen3_omni_moe.py | 2 + tests/test_distributed_config.py | 26 +- tests/test_fsdp_mixin.py | 159 ++++------ tests/test_tensor_parallel_mixin.py | 13 +- 36 files changed, 270 insertions(+), 333 deletions(-) diff --git a/src/transformers/configuration_utils.py b/src/transformers/configuration_utils.py index cbc30758d759..862139d6ae21 100755 --- a/src/transformers/configuration_utils.py +++ b/src/transformers/configuration_utils.py @@ -147,6 +147,9 @@ 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_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. @@ -219,6 +222,7 @@ class PreTrainedConfig(PushToHubMixin, RotaryEmbeddingConfigMixin): attribute_map: ClassVar[dict[str, str]] = {} base_model_tp_plan: ClassVar[dict[str, Any] | None] = None base_model_sp_plan: ClassVar[dict[str, Any] | None] = None + base_model_fsdp_plan: ClassVar[dict[Any, str] | None] = None base_model_pp_plan: ClassVar[dict[str, Sequence[list[str]]] | None] = None base_model_ep_plan: ClassVar[dict[str, Sequence[list[str]]] | None] = None _auto_class: ClassVar[str | None] = None diff --git a/src/transformers/distributed/configuration_utils.py b/src/transformers/distributed/configuration_utils.py index 386f66cc4b73..c41dd4978d50 100644 --- a/src/transformers/distributed/configuration_utils.py +++ b/src/transformers/distributed/configuration_utils.py @@ -31,20 +31,20 @@ class DistributedConfig: Args: tp_size (`int`, *optional*): Number of devices for tensor parallelism. If `None` and `fsdp_size` is set, defaults to 1. - tp_plan (`str` or `dict`, *optional*): - Tensor parallel sharding plan. Use `"auto"` for the model's default plan. + tp_plan (`dict`, *optional*): + Tensor parallel sharding plan. Leave as `None` to use the model's default plan. fsdp_size (`int`, *optional*): Number of devices for FSDP (data parallelism). If `None` and `tp_size` is set, defaults to 1. - fsdp_plan (`str` or `dict`, *optional*): - FSDP wrapping plan. Use `"auto"` to wrap each transformer layer + root. + fsdp_plan (`dict`, *optional*): + FSDP wrapping plan. Leave as `None` to wrap each transformer layer + root. """ tp_size: int | None = None - tp_plan: str | dict[str, str] | None = None + tp_plan: dict[str, str] | None = None enable_sequence_parallel: bool = False enable_expert_parallel: bool = False fsdp_size: int | None = None - fsdp_plan: str | dict | None = None + fsdp_plan: dict | None = None def __post_init__(self): if self.tp_size is None and self.fsdp_size is None: @@ -54,10 +54,6 @@ def __post_init__(self): self.tp_size = 1 if self.fsdp_size is None: self.fsdp_size = 1 - if self.tp_plan is None: - self.tp_plan = "auto" - if self.fsdp_plan is None: - self.fsdp_plan = "auto" if torch.distributed.is_available() and torch.distributed.is_initialized(): world_size = torch.distributed.get_world_size() diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py index b499f1d88907..d3672ce71b80 100644 --- a/src/transformers/distributed/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -19,6 +19,7 @@ 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: @@ -76,7 +77,8 @@ def initialize_fsdp( This function is called when the model is loaded and fsdp_plan is set. Args: - fsdp_plan: Optional FSDP config dict with an explicit "mode". + fsdp_plan: Optional FSDP config dict. Manual mode is signaled by the + presence of a ``"modules"`` key; otherwise auto mode is used. device_mesh: Optional pre-created DeviceMesh for FSDP. device_map: Optional device map. @@ -153,70 +155,8 @@ def initialize_fsdp( return device_map, device_mesh, fsdp_size -def get_transformer_block_classes(model): - """ - Identifies transformer block classes in a model for FSDP wrapping. - These are typically the repeated layers that benefit from FSDP sharding. - - Returns a set of module classes that should be wrapped with fully_shard(). - """ - block_classes = set() - - # Common transformer block class names - block_names = { - "DecoderLayer", - "EncoderLayer", - "TransformerBlock", - "Block", - "Layer", - } - - for module in model.modules(): - class_name = module.__class__.__name__ - # Use endswith to avoid false positives (e.g. "Layer" matching "LayerNorm") - for block_name in block_names: - if class_name.endswith(block_name): - block_classes.add(type(module)) - break - - # Filter out nested block classes (e.g. SparseMoeBlock inside DecoderLayer). - # We only want to FSDP-wrap the outermost block classes. If a class like - # MoeBlock only ever appears inside a DecoderLayer, we skip it. - if len(block_classes) > 1: - # Collect the dotted module paths for each candidate class. - # i.e: {DecoderLayer: ["layers.0", "layers.1"], - # MoeBlock: ["layers.0.moe", "layers.1.moe"]} - paths_by_class = {} - for name, module in model.named_modules(): - cls = type(module) - if cls in block_classes: - paths_by_class.setdefault(cls, []).append(name) - - def _is_nested_inside_other_class(cls): - # A class is "inner" if every one of its instances lives under - # an instance of a different candidate class in the module tree. - paths = paths_by_class.get(cls, []) - if not paths: - return False - for path in paths: - has_parent = any( - path.startswith(parent_path + ".") - for other_cls, parent_paths in paths_by_class.items() - if other_cls is not cls - for parent_path in parent_paths - ) - if not has_parent: - return False - return True - - # Keep only the outer (non-nested) classes. - block_classes = {cls for cls in block_classes if not _is_nested_inside_other_class(cls)} - - return block_classes - - -def _get_auto_policy_kwargs(fsdp_plan: dict[str, Any]) -> dict[str, Any]: - """Parse auto-mode fsdp_plan into fully_shard policy kwargs.""" +def _get_policy_kwargs(fsdp_plan: dict[str, Any]) -> dict[str, Any]: + """Parse `cpu_offload` / `mixed_precision` flags from the user fsdp_plan into fully_shard kwargs.""" policy_kwargs = {} if fsdp_plan.get("cpu_offload"): policy_kwargs["offload_policy"] = CPUOffloadPolicy() @@ -229,75 +169,6 @@ def _get_auto_policy_kwargs(fsdp_plan: dict[str, Any]) -> dict[str, Any]: return policy_kwargs -def _auto_shard_input_embedding(input_embed, is_weights_tied: bool, device_mesh, auto_policy_kwargs): - # Shard input embeddings (only when not tied). - # When tied, the shared weight is grouped with the final norm in step 3. - if input_embed is None or is_weights_tied: - return - fully_shard(input_embed, mesh=device_mesh, reshard_after_forward=True, **auto_policy_kwargs) - logger.debug(f"Applied fully_shard to input embeddings ({type(input_embed).__name__})") - - -def _auto_shard_transformer_blocks(model, block_classes, device_mesh, auto_policy_kwargs): - for name, module in model.named_modules(): - if type(module) in block_classes: - fully_shard(module, mesh=device_mesh, reshard_after_forward=True, **auto_policy_kwargs) - logger.debug(f"Applied fully_shard to {name} ({type(module).__name__})") - - -def _find_final_norm(model, decoder_layer_names): - """Find the final normalization layer before the output head. - - Searches only within the base model scope (e.g. ``model.*``) so that - norms inside the output head / prediction head (e.g. ``lm_head.norm``) - are excluded. - """ - base_prefix = model.base_model_prefix # e.g. "model" - final_norm = None - for name, module in model.named_modules(): - if "Norm" not in type(module).__name__: - continue - # Only consider norms inside the base model (skip root-level heads like lm_head) - if base_prefix and not name.startswith(base_prefix + ".") and name != base_prefix: - continue - if any(name.startswith(layer_name + ".") for layer_name in decoder_layer_names): - continue - final_norm = module - return final_norm - - -def _auto_get_tail_modules(model, decoder_layer_names, input_embed, output_embed, is_weights_tied: bool) -> list: - # Group final norm + output head. - # NOTE(3outeille): Small optimization by forcing reshard_after_forward=False for the final norm and output head. - # Otherwise, that would mean reshard/freeing full params after the last forward and immediately re-all-gathering - # them in the backward pass, which is wasteful. Better to keep them gathered for reuse. - # Untied: [final_norm, lm_head] - # Tied: [final_norm, embed_tokens] - embed_tokens.weight IS lm_head.weight. - tail_modules = [] - - final_norm = _find_final_norm(model, decoder_layer_names) - - if final_norm is not None: - tail_modules.append(final_norm) - - if is_weights_tied: - if input_embed is not None: - tail_modules.append(input_embed) - elif output_embed is not None: - tail_modules.append(output_embed) - - return tail_modules - - -def _auto_shard_tail_modules(tail_modules, device_mesh, auto_policy_kwargs): - if len(tail_modules) > 1: - fully_shard(tail_modules, mesh=device_mesh, reshard_after_forward=False, **auto_policy_kwargs) - logger.debug(f"Applied fully_shard to {[type(m).__name__ for m in tail_modules]} grouped (reshard=False)") - elif len(tail_modules) == 1: - fully_shard(tail_modules[0], mesh=device_mesh, reshard_after_forward=False, **auto_policy_kwargs) - logger.debug(f"Applied fully_shard to {type(tail_modules[0]).__name__} (reshard=False)") - - def _parse_manual_plan_entry( entry: list[str], ) -> tuple[bool, MixedPrecisionPolicy | None, CPUOffloadPolicy | None]: @@ -377,20 +248,6 @@ def _iter_manual_plan_targets(model, pattern, name_to_module, already_sharded_na yield name, module -def _parse_fsdp_plan_mode(fsdp_plan: dict[str, Any]) -> Literal["auto", "manual"]: - if isinstance(fsdp_plan, str): - fsdp_plan = {"mode": fsdp_plan} - - if not isinstance(fsdp_plan, dict): - raise ValueError(f"fsdp_plan must be a dict with a 'mode' key, got {type(fsdp_plan)}") - - mode = fsdp_plan.get("mode") - if mode not in {"auto", "manual"}: - raise ValueError("fsdp_plan['mode'] must be either 'auto' or 'manual'.") - - return mode - - def _get_manual_plan_modules(fsdp_plan: dict[str, Any]) -> dict[str, list[str]]: modules = fsdp_plan.get("modules") if not isinstance(modules, dict): @@ -398,33 +255,81 @@ def _get_manual_plan_modules(fsdp_plan: dict[str, Any]) -> dict[str, list[str]]: return modules +def is_tail_pair(entries) -> bool: + """Match the canonical tail pair: one final norm + the output head (or tied embedding).""" + if len(entries) != 2: + return False + names = [name for name, _ in entries] + has_norm = any(n == "norm" or n.endswith(".norm") for n in names) + has_head = any(n in {"lm_head", "embed_tokens"} or n.endswith((".lm_head", ".embed_tokens")) for n in names) + return has_norm and has_head + + +def tied_source_path(model) -> str | None: + """Return the dotted path of the input embedding module (the tied source).""" + input_embed = getattr(model, "get_input_embeddings", lambda: None)() + if input_embed is None: + return None + for name, module in model.named_modules(): + if module is input_embed: + return name + return None + + +def _resolve_plan_key(name_to_module: dict, key: str): + """Resolve a plan key into the matching (name, module) pairs. + + Supports exact module names and tp_plan-style wildcards (via + ``replace_layer_number_by_wildcard``). + """ + if key in name_to_module: + return [(key, name_to_module[key])] + return [(name, mod) for name, mod in name_to_module.items() if replace_layer_number_by_wildcard(name) == key] + + +def _iter_plan_targets(model, plan, is_weights_tied: bool, tied_source: str | None): + """Yield ``(name, module, strategy)`` for every module the plan applies to. + + Expands wildcards via ``_resolve_plan_key`` and pre-applies tying rules: + skips the standalone tied-source entry, and rewrites a keep ``"lm_head"`` + entry to the tied source so the shared parameter is wrapped once. + """ + name_to_module = dict(model.named_modules()) + for key, strategy in plan.items(): + if is_weights_tied and key == tied_source: + continue + if is_weights_tied and key == "lm_head" and strategy == "keep_full_weight" and tied_source is not None: + yield tied_source, name_to_module[tied_source], strategy + continue + for name, module in _resolve_plan_key(name_to_module, key): + yield name, module, strategy + + def apply_fully_shard_data_parallel( model, fsdp_mesh, - fsdp_plan: dict[str, Any] | str | None, + fsdp_plan: dict[str, Any] | None, ): """ - Apply FSDP2 (fully_shard) to a model following TorchTitan's approach. - fsdp_plan: - Explicit FSDP config dict with a required "mode" key, or a string - shorthand such as "auto" or "manual". + Apply FSDP2 (fully_shard) to a model. - Auto mode: - fsdp_plan = "auto" + When ``fsdp_plan`` is ``None`` or doesn't contain a ``"modules"`` key, the + model-declared ``model._fsdp_plan`` drives sharding. Policies (`cpu_offload`, + `mixed_precision`) from ``fsdp_plan`` are applied on top. - Auto mode (equivalent): - fsdp_plan = {"mode": "auto"} + When ``fsdp_plan`` has a ``"modules"`` key, the user fully specifies the + layout (manual mode). - Auto mode with optional policies: - fsdp_plan = {"mode": "auto", "cpu_offload": False, "mixed_precision": True} + Examples: + # Plan-driven (uses model._fsdp_plan). + fsdp_plan = None + fsdp_plan = {"cpu_offload": True, "mixed_precision": True} - Manual mode: + # Manual override. fsdp_plan = { - "mode": "manual", "modules": { "model.embed_tokens": ["free_full_weight"], - "model.layers.0.self_attn": ["free_full_weight", "cpu_offload", "mixed_precision"], - "model.layers.0.mlp": ["free_full_weight"], + "model.layers.0.mlp": ["free_full_weight", "cpu_offload", "mixed_precision"], "model.norm": ["keep_full_weight"], "lm_head": ["keep_full_weight"], }, @@ -437,10 +342,7 @@ def apply_fully_shard_data_parallel( raise OSError("FSDP2 requires torch>=2.5") if fsdp_plan is None: - return model - - if fsdp_plan == "auto": - fsdp_plan = {"mode": fsdp_plan} + fsdp_plan = {} input_embed = getattr(model, "get_input_embeddings", lambda: None)() output_embed = getattr(model, "get_output_embeddings", lambda: None)() @@ -452,39 +354,50 @@ def apply_fully_shard_data_parallel( and input_embed.weight is output_embed.weight ) - mode = _parse_fsdp_plan_mode(fsdp_plan) - - if mode == "auto": - auto_policy_kwargs = _get_auto_policy_kwargs(fsdp_plan) + if not isinstance(fsdp_plan, dict): + raise ValueError(f"fsdp_plan must be a dict, got {type(fsdp_plan)}") + is_manual = "modules" in fsdp_plan - block_classes = get_transformer_block_classes(model) - # Need to collect decoder layer names for norm detection. - decoder_layer_names = {name for name, module in model.named_modules() if type(module) in block_classes} + if not is_manual: + policy_kwargs = _get_policy_kwargs(fsdp_plan) - if not block_classes: - logger.warning( - "Could not auto-detect transformer block classes for FSDP. Applying FSDP only to root module." + plan = getattr(model, "_fsdp_plan", None) or {} + if not plan: + raise ValueError( + f"{type(model).__name__} has no `_fsdp_plan` declared. Either set " + "`base_model_fsdp_plan` on the config and `_fsdp_plan` on the head class, " + "or pass an explicit `fsdp_plan={'modules': {...}}` manual override." ) - else: - _auto_shard_input_embedding(input_embed, is_weights_tied, fsdp_mesh, auto_policy_kwargs) - - _auto_shard_transformer_blocks(model, block_classes, fsdp_mesh, auto_policy_kwargs) - tail_modules = _auto_get_tail_modules( - model, decoder_layer_names, input_embed, output_embed, is_weights_tied - ) - _auto_shard_tail_modules(tail_modules, fsdp_mesh, auto_policy_kwargs) + tied_source = tied_source_path(model) if is_weights_tied else None + keep_buffer: list[tuple[str, Any]] = [] + + for name, module, strategy in _iter_plan_targets(model, plan, is_weights_tied, tied_source): + if strategy == "keep_full_weight": + keep_buffer.append((name, module)) + continue + fully_shard(module, mesh=fsdp_mesh, reshard_after_forward=True, **policy_kwargs) + logger.debug(f"Applied fully_shard to {name} (reshard=True)") + + # Optimization: when the keep buffer is exactly the (final_norm, lm_head/embed) + # tail pair, bundle them into one fully_shard so that we dont need to do all-gather during backward pass. + if is_tail_pair(keep_buffer): + keep_names = [n for n, _ in keep_buffer] + keep_modules = [m for _, m in keep_buffer] + fully_shard(keep_modules, mesh=fsdp_mesh, reshard_after_forward=False, **policy_kwargs) + logger.debug(f"Grouped tail {keep_names} (reshard=False)") + else: + for name, module in keep_buffer: + fully_shard(module, mesh=fsdp_mesh, reshard_after_forward=False, **policy_kwargs) + logger.debug(f"Applied fully_shard to {name} (reshard=False)") # Shard root model - fully_shard(model, mesh=fsdp_mesh, **auto_policy_kwargs) + fully_shard(model, mesh=fsdp_mesh, **policy_kwargs) - logger.info( - f"FSDP2 applied to model: {len(block_classes)} block type(s), {len(decoder_layer_names)} decoder layers" - ) + logger.info(f"FSDP2 applied to model via _fsdp_plan: {len(plan)} entries") else: # fsdp_plan = { - # "mode": "manual", # "modules": { # "model.layers.0.self_attn": ["free_full_weight"], # reshard_after_forward=True # "model.norm": ["keep_full_weight"], # reshard_after_forward=False diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index 48b9d2e0aa82..202c10af6c6e 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -468,15 +468,12 @@ def apply_tensor_parallel(model, tp_mesh, tp_plan): for ``parallelize_module``. Plan values are string names looked up in ``ALL_PARALLEL_STYLES``. """ - if tp_plan is None: - return model - - if tp_plan == "auto": - distributed_config = getattr(model.config, "distributed_config", None) - sp_requested = getattr(distributed_config, "enable_sequence_parallel", False) - sp_supported = getattr(model.config, "base_model_sp_plan", None) is not None + distributed_config = getattr(model.config, "distributed_config", None) + sp_requested = getattr(distributed_config, "enable_sequence_parallel", False) + sp_supported = getattr(model.config, "base_model_sp_plan", None) is not None + enable_sp = sp_requested and sp_supported - enable_sp = sp_requested and sp_supported + if tp_plan is None: if enable_sp: tp_plan = dict(model._sp_plan or {}) else: diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 81f5435251e8..e75706ec4002 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -1379,13 +1379,24 @@ def post_init(self): # 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 {} - self._tp_plan, self._sp_plan, self._ep_plan, self._pp_plan = dict(cls_tp_plan), dict(cls_sp_plan), {}, {} - # If current model is a base model, attach `base_model_tp_plan` and `base_model_pp_plan` from config + cls_fsdp_plan = getattr(self, "_fsdp_plan", None) or {} + self._tp_plan, self._sp_plan, self._ep_plan, self._pp_plan, self._fsdp_plan = ( + dict(cls_tp_plan), + dict(cls_sp_plan), + {}, + {}, + dict(cls_fsdp_plan), + ) + # If current model is a base model, attach `base_model_*_plan` from config if self.base_model is self: self._pp_plan = self.config.base_model_pp_plan.copy() if self.config.base_model_pp_plan is not None else {} self._tp_plan = self.config.base_model_tp_plan.copy() if self.config.base_model_tp_plan is not None else {} self._sp_plan = self.config.base_model_sp_plan.copy() if self.config.base_model_sp_plan is not None else {} self._ep_plan = self.config.base_model_ep_plan.copy() if self.config.base_model_ep_plan is not None else {} + self._fsdp_plan = ( + self.config.base_model_fsdp_plan.copy() if self.config.base_model_fsdp_plan is not None else {} + ) + # Current submodel should register its tied weights self.all_tied_weights_keys = self.get_expanded_tied_weights_keys(all_submodels=False) # Current submodel should register its `_keep_in_fp32_modules` @@ -1411,6 +1422,8 @@ def post_init(self): self._sp_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()}) if plan := getattr(module, "_pp_plan", None): self._pp_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()}) + if plan := getattr(module, "_fsdp_plan", None): + self._fsdp_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()}) # Always attach the keys of the children (if the children's config says to NOT tie, then it's empty) if tied_keys := getattr(module, "all_tied_weights_keys", None): self.all_tied_weights_keys.update({f"{name}.{k}": f"{name}.{v}" for k, v in tied_keys.copy().items()}) @@ -3998,9 +4011,9 @@ def from_pretrained( device placement or dispatch. Launch with `torchrun --nproc_per_node=N script.py`. Accepts `tp_size`, `tp_plan`, `fsdp_size`, `fsdp_plan`. When a size is specified without a - plan, the plan defaults to `"auto"`. `tp_plan="auto"` uses the model's predefined tensor - parallel sharding plan. `fsdp_plan="auto"` wraps each transformer layer individually with - FSDP2 (`fully_shard`). Both plans also accept a `dict` for manual control: `tp_plan` maps + plan, the plan defaults to the model's predefined behavior: TP uses the model's predefined + tensor parallel sharding plan, FSDP wraps each transformer layer individually with FSDP2 + (`fully_shard`). Both plans also accept a `dict` for manual control: `tp_plan` maps parameter names to parallel styles (e.g. `{"model.layers.*.self_attn.q_proj": "colwise"}`), `fsdp_plan` maps module names to wrap (e.g. `{"model.layers.0": {}, "model.layers.1": {}}`). @@ -4312,7 +4325,9 @@ def from_pretrained( device_map = _get_device_map(model, device_map, max_memory, hf_quantizer) # Finalize model weight initialization - active_tp_plan = getattr(model, "_tp_plan", None) if getattr(distributed_config, "tp_plan", None) else None + active_tp_plan = ( + getattr(model, "_tp_plan", None) if getattr(distributed_config, "tp_size", None) is not None else None + ) load_config = LoadStateDictConfig( pretrained_model_name_or_path=pretrained_model_name_or_path, ignore_mismatched_sizes=ignore_mismatched_sizes, diff --git a/src/transformers/models/deepseek_v4/modeling_deepseek_v4.py b/src/transformers/models/deepseek_v4/modeling_deepseek_v4.py index 4f93c0dadea0..664a9bcbc1cd 100644 --- a/src/transformers/models/deepseek_v4/modeling_deepseek_v4.py +++ b/src/transformers/models/deepseek_v4/modeling_deepseek_v4.py @@ -1398,6 +1398,7 @@ class DeepseekV4ForCausalLM(DeepseekV4PreTrainedModel, GenerationMixin): _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/doge/modeling_doge.py b/src/transformers/models/doge/modeling_doge.py index 7d2f91f444cf..1f112465369d 100644 --- a/src/transformers/models/doge/modeling_doge.py +++ b/src/transformers/models/doge/modeling_doge.py @@ -719,6 +719,7 @@ class DogeForCausalLM(DogePreTrainedModel, GenerationMixin): _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/dots1/modeling_dots1.py b/src/transformers/models/dots1/modeling_dots1.py index 858da411d46b..05b1a6d20b4c 100644 --- a/src/transformers/models/dots1/modeling_dots1.py +++ b/src/transformers/models/dots1/modeling_dots1.py @@ -572,6 +572,7 @@ class Dots1ForCausalLM(Dots1PreTrainedModel, GenerationMixin): _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py b/src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py index a594cd4e306e..939201f4ff19 100644 --- a/src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py +++ b/src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py @@ -657,6 +657,7 @@ class Ernie4_5_MoeForCausalLM(Ernie4_5_MoePreTrainedModel, GenerationMixin): _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/flex_olmo/modeling_flex_olmo.py b/src/transformers/models/flex_olmo/modeling_flex_olmo.py index fa1811906643..a9403a4110b0 100644 --- a/src/transformers/models/flex_olmo/modeling_flex_olmo.py +++ b/src/transformers/models/flex_olmo/modeling_flex_olmo.py @@ -600,6 +600,7 @@ class FlexOlmoForCausalLM(FlexOlmoPreTrainedModel, GenerationMixin): _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/gpt_oss/modeling_gpt_oss.py b/src/transformers/models/gpt_oss/modeling_gpt_oss.py index bd9002332568..ebb980bd8a9b 100644 --- a/src/transformers/models/gpt_oss/modeling_gpt_oss.py +++ b/src/transformers/models/gpt_oss/modeling_gpt_oss.py @@ -595,6 +595,7 @@ class GptOssForCausalLM(GptOssPreTrainedModel, GenerationMixin): _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/granitemoe/modeling_granitemoe.py b/src/transformers/models/granitemoe/modeling_granitemoe.py index 6731c00e178a..4f79272f164f 100644 --- a/src/transformers/models/granitemoe/modeling_granitemoe.py +++ b/src/transformers/models/granitemoe/modeling_granitemoe.py @@ -629,6 +629,7 @@ class GraniteMoeForCausalLM(GraniteMoePreTrainedModel, GenerationMixin): _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config: GraniteMoeConfig): super().__init__(config) diff --git a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py index bc1378ba5b55..929a82f5ed82 100644 --- a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py +++ b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py @@ -1310,6 +1310,7 @@ class GraniteMoeHybridForCausalLM(GraniteMoeHybridPreTrainedModel, GenerationMix _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config: GraniteMoeHybridConfig): super().__init__(config) diff --git a/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py b/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py index 0c0a4e92ed7f..8d0bc797a57b 100644 --- a/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py +++ b/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py @@ -698,6 +698,7 @@ class GraniteMoeSharedForCausalLM(GraniteMoeSharedPreTrainedModel, GenerationMix _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config: GraniteMoeSharedConfig): super().__init__(config) diff --git a/src/transformers/models/jamba/modeling_jamba.py b/src/transformers/models/jamba/modeling_jamba.py index 08c8b173a945..ba5a759c4636 100755 --- a/src/transformers/models/jamba/modeling_jamba.py +++ b/src/transformers/models/jamba/modeling_jamba.py @@ -846,6 +846,7 @@ class JambaForCausalLM(JambaPreTrainedModel, GenerationMixin): _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config: JambaConfig): super().__init__(config) diff --git a/src/transformers/models/laguna/modeling_laguna.py b/src/transformers/models/laguna/modeling_laguna.py index d3796f6e438e..410fd9a04258 100644 --- a/src/transformers/models/laguna/modeling_laguna.py +++ b/src/transformers/models/laguna/modeling_laguna.py @@ -675,6 +675,7 @@ class LagunaForCausalLM(LagunaPreTrainedModel, GenerationMixin): _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/minimax/modeling_minimax.py b/src/transformers/models/minimax/modeling_minimax.py index 7540f1ce5329..7da5e5867ce2 100644 --- a/src/transformers/models/minimax/modeling_minimax.py +++ b/src/transformers/models/minimax/modeling_minimax.py @@ -792,6 +792,7 @@ class MiniMaxForCausalLM(MiniMaxPreTrainedModel, GenerationMixin): _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/minimax_m2/modeling_minimax_m2.py b/src/transformers/models/minimax_m2/modeling_minimax_m2.py index 82389018bbbf..370ea88c971e 100644 --- a/src/transformers/models/minimax_m2/modeling_minimax_m2.py +++ b/src/transformers/models/minimax_m2/modeling_minimax_m2.py @@ -591,6 +591,7 @@ class MiniMaxM2ForCausalLM(MiniMaxM2PreTrainedModel, GenerationMixin): _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/mixtral/configuration_mixtral.py b/src/transformers/models/mixtral/configuration_mixtral.py index 242664d356b5..41beb623719c 100644 --- a/src/transformers/models/mixtral/configuration_mixtral.py +++ b/src/transformers/models/mixtral/configuration_mixtral.py @@ -70,6 +70,13 @@ class MixtralConfig(PreTrainedConfig): "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 f76680ace4c4..df71b44ebac2 100644 --- a/src/transformers/models/mixtral/modeling_mixtral.py +++ b/src/transformers/models/mixtral/modeling_mixtral.py @@ -583,6 +583,7 @@ class MixtralForCausalLM(MixtralPreTrainedModel, GenerationMixin): _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/mixtral/modular_mixtral.py b/src/transformers/models/mixtral/modular_mixtral.py index 139e580fbca7..438005523f7e 100644 --- a/src/transformers/models/mixtral/modular_mixtral.py +++ b/src/transformers/models/mixtral/modular_mixtral.py @@ -335,6 +335,7 @@ def forward( class MixtralForCausalLM(MistralForCausalLM): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/olmoe/modeling_olmoe.py b/src/transformers/models/olmoe/modeling_olmoe.py index f69926761833..10950656c462 100644 --- a/src/transformers/models/olmoe/modeling_olmoe.py +++ b/src/transformers/models/olmoe/modeling_olmoe.py @@ -607,6 +607,7 @@ class OlmoeForCausalLM(OlmoePreTrainedModel, GenerationMixin): _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/phimoe/modeling_phimoe.py b/src/transformers/models/phimoe/modeling_phimoe.py index 97a095278065..d933a6ec69d9 100644 --- a/src/transformers/models/phimoe/modeling_phimoe.py +++ b/src/transformers/models/phimoe/modeling_phimoe.py @@ -775,6 +775,7 @@ class PhimoeForCausalLM(PhimoePreTrainedModel, GenerationMixin): _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py b/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py index 1afe8ac3c2d2..f7142b4650fb 100644 --- a/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py +++ b/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py @@ -620,6 +620,7 @@ class Qwen2MoeForCausalLM(Qwen2MoePreTrainedModel, GenerationMixin): _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/qwen3/configuration_qwen3.py b/src/transformers/models/qwen3/configuration_qwen3.py index b3dc0e89a6bd..ff0cd8f63a23 100644 --- a/src/transformers/models/qwen3/configuration_qwen3.py +++ b/src/transformers/models/qwen3/configuration_qwen3.py @@ -83,6 +83,15 @@ class Qwen3Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Values are sharding strategies; policies (cpu_offload, mixed_precision) + # live on DistributedConfig.fsdp_plan. All entries marked `keep_full_weight` are bundled + # into a single fully_shard([...]) call at apply time so they share one all-gather. + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 151936 hidden_size: int = 4096 intermediate_size: int = 22016 diff --git a/src/transformers/models/qwen3/modeling_qwen3.py b/src/transformers/models/qwen3/modeling_qwen3.py index b2e4a4bea863..2ba669225879 100644 --- a/src/transformers/models/qwen3/modeling_qwen3.py +++ b/src/transformers/models/qwen3/modeling_qwen3.py @@ -444,6 +444,7 @@ class Qwen3ForCausalLM(Qwen3PreTrainedModel, GenerationMixin): _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/qwen3/modular_qwen3.py b/src/transformers/models/qwen3/modular_qwen3.py index 73cde6d89a7a..ea6ddb2f7499 100644 --- a/src/transformers/models/qwen3/modular_qwen3.py +++ b/src/transformers/models/qwen3/modular_qwen3.py @@ -108,6 +108,8 @@ def forward( class Qwen3ForCausalLM(Qwen2ForCausalLM): + _fsdp_plan = {"lm_head": "keep_full_weight"} + def forward( self, **super_kwargs: Unpack[TransformersKwargs], diff --git a/src/transformers/models/qwen3_5/modeling_qwen3_5.py b/src/transformers/models/qwen3_5/modeling_qwen3_5.py index 2e9e6448a4ea..1ffcd1b6da8d 100644 --- a/src/transformers/models/qwen3_5/modeling_qwen3_5.py +++ b/src/transformers/models/qwen3_5/modeling_qwen3_5.py @@ -1636,6 +1636,7 @@ class Qwen3_5ForCausalLM(Qwen3_5PreTrainedModel, GenerationMixin): _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} config: Qwen3_5TextConfig _keys_to_ignore_on_load_unexpected = [r"^mtp.*", r"^model.visual.*"] diff --git a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py index 09e5a78e6799..f64ff70a1f70 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 @@ -1838,6 +1838,7 @@ class Qwen3_5MoeForCausalLM(Qwen3_5MoePreTrainedModel, GenerationMixin): _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} config: Qwen3_5MoeTextConfig _keys_to_ignore_on_load_unexpected = [r"^mtp.*", r"^model.visual.*"] diff --git a/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py b/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py index 404869f16bfb..012e0cac7243 100644 --- a/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py +++ b/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py @@ -612,6 +612,7 @@ class Qwen3MoeForCausalLM(Qwen3MoePreTrainedModel, GenerationMixin): _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/qwen3_next/modeling_qwen3_next.py b/src/transformers/models/qwen3_next/modeling_qwen3_next.py index 834eb6dbc241..faca94ba5c43 100644 --- a/src/transformers/models/qwen3_next/modeling_qwen3_next.py +++ b/src/transformers/models/qwen3_next/modeling_qwen3_next.py @@ -1091,6 +1091,7 @@ class Qwen3NextForCausalLM(Qwen3NextPreTrainedModel, GenerationMixin): _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py b/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py index fc9fa220c699..223c8932fab5 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 @@ -303,6 +303,15 @@ class Qwen3OmniMoeTalkerCodePredictorConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Values are sharding strategies; policies (cpu_offload, mixed_precision) + # live on DistributedConfig.fsdp_plan. All entries marked `keep_full_weight` are bundled + # into a single fully_shard([...]) call at apply time so they share one all-gather. + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 2048 hidden_size: int = 1024 intermediate_size: int = 3072 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 a93e2c9df284..4d4b203fd25d 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 @@ -2638,6 +2638,7 @@ class Qwen3OmniMoeTalkerCodePredictorModelForConditionalGeneration(Qwen3OmniMoeP _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} config_class = Qwen3OmniMoeTalkerCodePredictorConfig base_model_prefix = "talker.code_predictor" _can_record_outputs = { @@ -3023,6 +3024,7 @@ class Qwen3OmniMoeTalkerForConditionalGeneration(Qwen3OmniMoeThinkerTextPreTrain _tp_plan = {"codec_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"codec_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} config_class = Qwen3OmniMoeTalkerConfig base_model_prefix = "talker" _no_split_modules = ["Qwen3OmniMoeTalkerCodePredictorModelForConditionalGeneration"] diff --git a/tests/test_distributed_config.py b/tests/test_distributed_config.py index abffeb15d867..3cf5b4a914db 100644 --- a/tests/test_distributed_config.py +++ b/tests/test_distributed_config.py @@ -6,24 +6,24 @@ class TestDistributedConfig: def test_2d_parallelism(self): - dc = DistributedConfig(tp_size=2, tp_plan="auto", fsdp_size=2, fsdp_plan="auto") + dc = DistributedConfig(tp_size=2, fsdp_size=2) assert dc.tp_size == 2 assert dc.fsdp_size == 2 - assert dc.tp_plan == "auto" - assert dc.fsdp_plan == "auto" + assert dc.tp_plan is None + assert dc.fsdp_plan is None def test_tp_only_defaults_fsdp_to_1(self): dc = DistributedConfig(tp_size=4) assert dc.tp_size == 4 assert dc.fsdp_size == 1 - assert dc.tp_plan == "auto" # size given → plan defaults to "auto" + assert dc.tp_plan is None def test_fsdp_only_defaults_tp_to_1(self): dc = DistributedConfig(fsdp_size=4) assert dc.tp_size == 1 assert dc.fsdp_size == 4 - assert dc.fsdp_plan == "auto" # size given → plan defaults to "auto" - assert dc.tp_plan == "auto" # tp_size got set to 1 → plan also defaults + assert dc.fsdp_plan is None + assert dc.tp_plan is None def test_empty_config(self): dc = DistributedConfig() @@ -33,10 +33,10 @@ def test_empty_config(self): assert dc.fsdp_plan is None def test_from_dict(self): - dc = DistributedConfig.from_dict({"tp_size": 2, "fsdp_size": 4, "tp_plan": "auto"}) + dc = DistributedConfig.from_dict({"tp_size": 2, "fsdp_size": 4}) assert dc.tp_size == 2 assert dc.fsdp_size == 4 - assert dc.tp_plan == "auto" + assert dc.tp_plan is None def test_from_dict_ignores_unknown_keys(self): dc = DistributedConfig.from_dict({"tp_size": 2, "unknown_key": 42}) @@ -52,11 +52,11 @@ def test_to_dict(self): d = dc.to_dict() assert d == { "tp_size": 2, - "tp_plan": "auto", + "tp_plan": None, "enable_sequence_parallel": False, "enable_expert_parallel": False, "fsdp_size": 4, - "fsdp_plan": "auto", + "fsdp_plan": None, } def test_to_dict_is_a_copy(self): @@ -73,16 +73,16 @@ def test_to_json_string(self): assert parsed["fsdp_size"] == 2 def test_to_json_file(self): - dc = DistributedConfig(tp_size=4, tp_plan="auto") + 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"] == "auto" + assert parsed["tp_plan"] is None def test_roundtrip_dict(self): - original = DistributedConfig(tp_size=2, tp_plan="auto", fsdp_size=4, fsdp_plan="auto") + original = DistributedConfig(tp_size=2, fsdp_size=4) restored = DistributedConfig.from_dict(original.to_dict()) assert original == restored diff --git a/tests/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py index dba4865f2953..32008c74eefb 100644 --- a/tests/test_fsdp_mixin.py +++ b/tests/test_fsdp_mixin.py @@ -50,12 +50,8 @@ from torch.nn.parallel import DistributedDataParallel as DDP from transformers.distributed import DistributedConfig - from transformers.distributed.fsdp import ( - _find_final_norm, - apply_fully_shard_data_parallel, - get_transformer_block_classes, - initialize_fsdp, - ) + from transformers.distributed.fsdp import apply_fully_shard_data_parallel, initialize_fsdp + from transformers.distributed.tensor_parallel import replace_layer_number_by_wildcard # ============================================================================= @@ -68,28 +64,10 @@ LR = 3e-4 SEED = 42 FSDP_TOP_MODEL_NAMES = { - # Dense - "gpt2", + # FSDP coverage is gated on models declaring `base_model_fsdp_plan` + class-level + # `_fsdp_plan`. Roll out to more models in follow-up PRs. "qwen3", - "phi", - "llama", - "modernbert_decoder", - "olmo3", - "phi3", - "mistral", - "lfm2", - "gemma2", - # MoE - "gpt_oss", - "glm_moe_dsa", - "qwen3_moe", - "glm4_moe_lite", - "qwen3_5_moe", - "deepseek_v2", - "qwen3_next", "mixtral", - "qwen2_moe", - "phimoe", } @@ -231,18 +209,15 @@ def _gather_ddp_state_dict(model): return {k: v.clone().detach().cpu() for k, v in model.module.state_dict().items()} -def _build_manual_fsdp_plan(config, device, policy_options=None): - """Build a default manual FSDP2 plan from model structure.""" - policy_options = policy_options or [] - set_seed(SEED) - model = AutoModelForCausalLM.from_config(config).to(device) - named_modules = dict(model.named_modules()) - id_to_name = {id(module): name for name, module in named_modules.items()} - block_classes = get_transformer_block_classes(model) +def _resolve_fsdp_plan_paths(model): + """Expand model._fsdp_plan into (paths, strategy) entries. - decoder_layer_names = {name for name, module in named_modules.items() if type(module) in block_classes} - layer_prefixes = {".".join(name.split(".")[:-1]) for name in decoder_layer_names} - assert layer_prefixes, "Expected at least one decoder layer prefix for manual FSDP plan." + Wildcard keys are expanded via ``replace_layer_number_by_wildcard``. When + weights are tied, the standalone embed_tokens entry is skipped and any + ``"lm_head"`` keep entry is rewritten to the tied source path (the keep + group will wrap the shared parameter once). + """ + plan = model._fsdp_plan input_embed = model.get_input_embeddings() output_embed = model.get_output_embeddings() @@ -253,27 +228,44 @@ def _build_manual_fsdp_plan(config, device, policy_options=None): and hasattr(output_embed, "weight") and input_embed.weight is output_embed.weight ) - embed_name = id_to_name.get(id(input_embed)) if input_embed is not None else None - output_name = id_to_name.get(id(output_embed)) if output_embed is not None else None - final_norm = _find_final_norm(model, decoder_layer_names) - norm_name = id_to_name.get(id(final_norm)) if final_norm is not None else None + tied_source = None + if weights_tied: + for name, mod in model.named_modules(): + if mod is input_embed: + tied_source = name + break + + name_to_module = dict(model.named_modules()) + entries: list[tuple[list[str], str]] = [] + for key, strategy in plan.items(): + if weights_tied and key == tied_source: + continue + if weights_tied and key == "lm_head" and strategy == "keep_full_weight": + entries.append(([tied_source], strategy)) + continue + + if key in name_to_module: + entries.append(([key], strategy)) + continue + matched = [name for name in name_to_module if replace_layer_number_by_wildcard(name) == key] + if matched: + entries.append((matched, strategy)) + return entries - module_plan = {name: ["free_full_weight", *policy_options] for name in layer_prefixes} - if norm_name: - module_plan[norm_name] = ["keep_full_weight", *policy_options] +def _build_manual_fsdp_plan(config, device, policy_options=None): + """Build a manual FSDP2 plan by expanding model._fsdp_plan.""" + policy_options = policy_options or [] + set_seed(SEED) + model = AutoModelForCausalLM.from_config(config).to(device) - if weights_tied: - if embed_name: - module_plan[embed_name] = ["keep_full_weight", *policy_options] - else: - if embed_name: - module_plan[embed_name] = ["free_full_weight", *policy_options] - if output_name: - module_plan[output_name] = ["keep_full_weight", *policy_options] + module_plan: dict[str, list[str]] = {} + for paths, strategy in _resolve_fsdp_plan_paths(model): + for path in paths: + module_plan[path] = [strategy, *policy_options] del model - return {"mode": "manual", "modules": module_plan} + return {"modules": module_plan} def _save_init_pretrained(rank, config, dtype): @@ -460,7 +452,7 @@ def _test_fsdp2_save_load_impl(rank, config_class, config_dict): batches = _build_repeated_training_batches(config, device, 3) - distributed_config = DistributedConfig(fsdp_size=dist.get_world_size(), fsdp_plan="auto") + distributed_config = DistributedConfig(fsdp_size=dist.get_world_size()) init_tmpdir, init_tmpdir_obj = _save_init_pretrained(rank, config, torch.float32) try: @@ -518,7 +510,7 @@ def _test_fsdp2_save_load_impl(rank, config_class, config_dict): def _test_fsdp2_sharding_structure_impl(rank, config_class, config_dict, tie_word_embeddings): """ - Verify that apply_fully_shard_data_parallel(fsdp_plan={"mode": "auto"}) wraps exactly the right modules. + Verify that apply_fully_shard_data_parallel(fsdp_plan=None) wraps exactly the right modules. Expected FSDP targets: UNTIED TIED @@ -533,45 +525,25 @@ def _test_fsdp2_sharding_structure_impl(rank, config_class, config_dict, tie_wor config = config_class.from_dict(config_dict) config.tie_word_embeddings = tie_word_embeddings - auto_plan = {"mode": "auto"} - device_map, device_mesh, _ = initialize_fsdp(fsdp_plan=auto_plan) + auto_plan = None + device_map, device_mesh, _ = initialize_fsdp(fsdp_plan={}) set_seed(SEED) model = AutoModelForCausalLM.from_config(config).to(device_map) - block_classes = get_transformer_block_classes(model) - assert block_classes, "get_transformer_block_classes found no block classes" - - decoder_layer_names = {name for name, module in model.named_modules() if type(module) in block_classes} - assert len(decoder_layer_names) > 0, "Expected at least one transformer block instance" - - id_to_name = {id(module): name for name, module in model.named_modules()} - - input_embed = model.get_input_embeddings() - output_embed = model.get_output_embeddings() - final_norm = _find_final_norm(model, decoder_layer_names) - weights_tied = ( - input_embed is not None - and output_embed is not None - and hasattr(input_embed, "weight") - and hasattr(output_embed, "weight") - and input_embed.weight is output_embed.weight - ) - - embed_name = id_to_name.get(id(input_embed)) - output_name = id_to_name.get(id(output_embed)) - norm_name = id_to_name.get(id(final_norm)) - - expected_targets = {""} | decoder_layer_names | {embed_name} | {norm_name} - if not weights_tied: - expected_targets |= {output_name} + # Expected FSDP targets come from model._fsdp_plan: every resolved path gets a + # fully_shard call (keep_full_weight entries are bundled into one group, but each + # member still appears as an FSDP-wrapped module in named_modules), plus the root. + expected_targets = {""} + for paths, _strategy in _resolve_fsdp_plan_paths(model): + expected_targets.update(paths) model = apply_fully_shard_data_parallel(model, device_mesh, fsdp_plan=auto_plan) actual_targets = {name for name, module in model.named_modules() if type(module).__name__.startswith("FSDP")} if rank == 0: - logger.debug(f" Weights tied: {weights_tied}") + 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)}") @@ -607,7 +579,6 @@ def _test_fsdp2_plan_vs_ddp_impl( if plan_mode == "auto": fsdp_plan = { - "mode": "auto", "cpu_offload": "cpu_offload" in policy_options, "mixed_precision": "mixed_precision" in policy_options, } @@ -774,28 +745,22 @@ def _run_fsdp2_distributed_test(self, test_name, test_impl, *test_args, **test_k # ========================================================================= @is_fsdp_test - def test_get_transformer_block_classes(self): - """get_transformer_block_classes() finds >= 1 block class for the model.""" + 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_get_transformer_block_classes") + 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) - block_classes = get_transformer_block_classes(model) - self.assertTrue(len(block_classes) > 0, f"No block classes found for {type(config).__name__}") - - for cls in block_classes: - count = sum(1 for m in model.modules() if type(m) is cls) - self.assertGreater(count, 0, f"Block class {cls.__name__} has no instances in model") + 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_get_transformer_block_classes (%.1fs)", status, time.perf_counter() - start_time - ) + logger.info("[FSDP] %s test: test_fsdp_plan_declared (%.1fs)", status, time.perf_counter() - start_time) @is_fsdp_test @require_fsdp diff --git a/tests/test_tensor_parallel_mixin.py b/tests/test_tensor_parallel_mixin.py index b43796ef4304..b891e7cd11c6 100644 --- a/tests/test_tensor_parallel_mixin.py +++ b/tests/test_tensor_parallel_mixin.py @@ -141,9 +141,7 @@ def _load_tp_and_reference_models(model_path, model_class, enable_sequence_paral tuple: (model_tp, model_ref, device) """ tp_size = dist.get_world_size() - distributed_config = DistributedConfig( - tp_size=tp_size, tp_plan="auto", enable_sequence_parallel=enable_sequence_parallel - ) + distributed_config = DistributedConfig(tp_size=tp_size, enable_sequence_parallel=enable_sequence_parallel) model_tp = model_class.from_pretrained( model_path, distributed_config=distributed_config, attn_implementation="sdpa" ) @@ -159,10 +157,6 @@ def _load_tp_and_reference_models(model_path, model_class, enable_sequence_paral def _get_active_tp_plan(model_tp): distributed_config = getattr(model_tp.config, "distributed_config", None) tp_plan = getattr(distributed_config, "tp_plan", None) - - if tp_plan == "auto": - return getattr(model_tp, "_tp_plan", None) or {} - return tp_plan or getattr(model_tp, "_tp_plan", None) or {} @@ -351,7 +345,9 @@ 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, distributed_config=DistributedConfig(tp_plan="auto"), quantization_config=quantization_config + model_path, + distributed_config=DistributedConfig(tp_size=dist.get_world_size()), + quantization_config=quantization_config, ) dist.barrier() @@ -407,7 +403,6 @@ def _load_ep_and_reference_models(model_path, model_class): model_path, distributed_config=DistributedConfig( tp_size=tp_size, - tp_plan="auto", enable_expert_parallel=True, ), ) From d13ab2e8c6b2012657a97eaa30ea41b826db8f4a Mon Sep 17 00:00:00 2001 From: 3outeille Date: Sat, 16 May 2026 09:41:29 +0000 Subject: [PATCH 099/116] edit fsdp plan to every other models --- .../models/afmoe/modeling_afmoe.py | 1 + .../models/afmoe/modular_afmoe.py | 1 + .../models/apertus/configuration_apertus.py | 8 ++++++++ .../models/apertus/modeling_apertus.py | 1 + .../models/arcee/configuration_arcee.py | 8 ++++++++ .../models/arcee/modeling_arcee.py | 1 + .../models/aria/configuration_aria.py | 8 ++++++++ src/transformers/models/aria/modeling_aria.py | 1 + .../models/bamba/modeling_bamba.py | 1 + .../models/cohere/configuration_cohere.py | 8 ++++++++ .../models/cohere/modeling_cohere.py | 1 + .../models/cohere2/configuration_cohere2.py | 8 ++++++++ .../models/cohere2/modeling_cohere2.py | 1 + .../models/cwm/configuration_cwm.py | 8 ++++++++ src/transformers/models/cwm/modeling_cwm.py | 1 + src/transformers/models/dbrx/modeling_dbrx.py | 1 + src/transformers/models/dbrx/modular_dbrx.py | 1 + .../deepseek_v2/configuration_deepseek_v2.py | 8 ++++++++ .../deepseek_v2/modeling_deepseek_v2.py | 1 + .../deepseek_v3/configuration_deepseek_v3.py | 8 ++++++++ .../deepseek_v3/modeling_deepseek_v3.py | 1 + .../deepseek_v4/configuration_deepseek_v4.py | 8 ++++++++ .../models/diffllama/modeling_diffllama.py | 1 + .../models/doge/configuration_doge.py | 8 ++++++++ .../models/dots1/configuration_dots1.py | 8 ++++++++ src/transformers/models/emu3/modeling_emu3.py | 1 + .../models/ernie4_5/configuration_ernie4_5.py | 8 ++++++++ .../models/ernie4_5/modeling_ernie4_5.py | 1 + .../configuration_ernie4_5_moe.py | 8 ++++++++ .../configuration_ernie4_5_vl_moe.py | 8 ++++++++ .../models/eurobert/configuration_eurobert.py | 8 ++++++++ .../models/eurobert/modeling_eurobert.py | 1 + .../models/eurobert/modular_eurobert.py | 1 + .../models/exaone4/configuration_exaone4.py | 8 ++++++++ .../models/exaone4/modeling_exaone4.py | 1 + .../exaone_moe/configuration_exaone_moe.py | 8 ++++++++ .../models/exaone_moe/modeling_exaone_moe.py | 1 + .../models/falcon_h1/modeling_falcon_h1.py | 1 + .../flex_olmo/configuration_flex_olmo.py | 8 ++++++++ .../models/gemma/configuration_gemma.py | 8 ++++++++ .../models/gemma/modeling_gemma.py | 1 + .../models/gemma2/configuration_gemma2.py | 8 ++++++++ .../models/gemma2/modeling_gemma2.py | 1 + .../models/gemma3/configuration_gemma3.py | 8 ++++++++ .../models/gemma3/modeling_gemma3.py | 1 + .../models/gemma3n/configuration_gemma3n.py | 8 ++++++++ .../models/gemma3n/modeling_gemma3n.py | 1 + .../models/gemma4/configuration_gemma4.py | 8 ++++++++ .../models/gemma4/modeling_gemma4.py | 1 + .../modeling_gemma4_assistant.py | 1 + .../models/glm/configuration_glm.py | 8 ++++++++ src/transformers/models/glm/modeling_glm.py | 1 + .../models/glm4/configuration_glm4.py | 8 ++++++++ src/transformers/models/glm4/modeling_glm4.py | 1 + .../models/glm4_moe/configuration_glm4_moe.py | 8 ++++++++ .../models/glm4_moe/modeling_glm4_moe.py | 1 + .../configuration_glm4_moe_lite.py | 8 ++++++++ .../glm4_moe_lite/modeling_glm4_moe_lite.py | 1 + .../models/glm4v/configuration_glm4v.py | 8 ++++++++ .../glm4v_moe/configuration_glm4v_moe.py | 8 ++++++++ .../glm_image/configuration_glm_image.py | 8 ++++++++ .../glm_moe_dsa/configuration_glm_moe_dsa.py | 8 ++++++++ .../glm_moe_dsa/modeling_glm_moe_dsa.py | 1 + .../models/glm_ocr/configuration_glm_ocr.py | 8 ++++++++ .../models/gpt_neox/configuration_gpt_neox.py | 8 ++++++++ .../models/gpt_oss/configuration_gpt_oss.py | 5 +++++ .../models/granite/configuration_granite.py | 8 ++++++++ .../models/granite/modeling_granite.py | 1 + .../configuration_granite4_vision.py | 8 ++++++++ .../models/helium/configuration_helium.py | 8 ++++++++ .../models/helium/modeling_helium.py | 1 + .../configuration_higgs_audio_v2.py | 8 ++++++++ .../modeling_hunyuan_v1_dense.py | 1 + .../hunyuan_v1_moe/modeling_hunyuan_v1_moe.py | 1 + .../models/hy_v3/configuration_hy_v3.py | 8 ++++++++ .../models/hy_v3/modeling_hy_v3.py | 1 + .../hyperclovax/configuration_hyperclovax.py | 8 ++++++++ .../hyperclovax/modeling_hyperclovax.py | 1 + .../models/jais2/configuration_jais2.py | 8 ++++++++ .../models/jais2/modeling_jais2.py | 1 + .../modeling_kyutai_speech_to_text.py | 1 + .../models/laguna/configuration_laguna.py | 8 ++++++++ src/transformers/models/lfm2/modeling_lfm2.py | 1 + .../models/lfm2_moe/modeling_lfm2_moe.py | 1 + .../models/llama/configuration_llama.py | 8 ++++++++ .../models/llama/modeling_llama.py | 1 + .../models/llama4/modeling_llama4.py | 1 + .../configuration_longcat_flash.py | 8 ++++++++ .../longcat_flash/modeling_longcat_flash.py | 1 + .../models/minimax/configuration_minimax.py | 8 ++++++++ .../minimax_m2/configuration_minimax_m2.py | 8 ++++++++ .../ministral/configuration_ministral.py | 8 ++++++++ .../models/ministral/modeling_ministral.py | 1 + .../ministral3/configuration_ministral3.py | 8 ++++++++ .../models/ministral3/modeling_ministral3.py | 1 + .../models/mistral/configuration_mistral.py | 8 ++++++++ .../models/mistral/modeling_mistral.py | 1 + .../models/mistral4/configuration_mistral4.py | 8 ++++++++ .../models/mistral4/modeling_mistral4.py | 1 + .../models/nanochat/modeling_nanochat.py | 1 + .../models/nanochat/modular_nanochat.py | 1 + .../models/olmo/configuration_olmo.py | 8 ++++++++ src/transformers/models/olmo/modeling_olmo.py | 1 + .../models/olmo2/configuration_olmo2.py | 8 ++++++++ .../models/olmo2/modeling_olmo2.py | 1 + .../models/olmo3/configuration_olmo3.py | 8 ++++++++ .../models/olmo3/modeling_olmo3.py | 1 + .../olmo_hybrid/configuration_olmo_hybrid.py | 8 ++++++++ .../olmo_hybrid/modeling_olmo_hybrid.py | 1 + .../models/olmoe/configuration_olmoe.py | 8 ++++++++ .../configuration_paddleocr_vl.py | 8 ++++++++ .../models/phi/configuration_phi.py | 8 ++++++++ src/transformers/models/phi/modeling_phi.py | 1 + .../models/phi3/configuration_phi3.py | 8 ++++++++ src/transformers/models/phi3/modeling_phi3.py | 1 + .../configuration_phi4_multimodal.py | 8 ++++++++ .../modeling_phi4_multimodal.py | 1 + .../models/phimoe/configuration_phimoe.py | 5 +++++ .../models/qwen2/configuration_qwen2.py | 8 ++++++++ .../models/qwen2/modeling_qwen2.py | 1 + .../configuration_qwen2_5_omni.py | 8 ++++++++ .../qwen2_5_vl/configuration_qwen2_5_vl.py | 8 ++++++++ .../qwen2_moe/configuration_qwen2_moe.py | 8 ++++++++ .../models/qwen2_moe/modular_qwen2_moe.py | 1 + .../models/qwen2_vl/configuration_qwen2_vl.py | 8 ++++++++ .../models/qwen3_5/configuration_qwen3_5.py | 8 ++++++++ .../qwen3_5_moe/configuration_qwen3_5_moe.py | 8 ++++++++ .../models/qwen3_5_moe/modular_qwen3_5_moe.py | 1 + .../qwen3_moe/configuration_qwen3_moe.py | 8 ++++++++ .../qwen3_next/configuration_qwen3_next.py | 8 ++++++++ .../configuration_qwen3_vl_moe.py | 8 ++++++++ .../models/seed_oss/configuration_seed_oss.py | 8 ++++++++ .../models/seed_oss/modeling_seed_oss.py | 1 + .../models/smollm3/configuration_smollm3.py | 8 ++++++++ .../models/smollm3/modeling_smollm3.py | 1 + .../solar_open/configuration_solar_open.py | 8 ++++++++ .../models/solar_open/modeling_solar_open.py | 1 + .../starcoder2/configuration_starcoder2.py | 8 ++++++++ .../models/starcoder2/modeling_starcoder2.py | 1 + .../models/t5gemma/configuration_t5gemma.py | 8 ++++++++ .../models/t5gemma2/configuration_t5gemma2.py | 8 ++++++++ .../vaultgemma/configuration_vaultgemma.py | 8 ++++++++ .../models/vaultgemma/modeling_vaultgemma.py | 1 + .../configuration_voxtral_realtime.py | 8 ++++++++ .../modeling_voxtral_realtime.py | 1 + .../models/youtu/configuration_youtu.py | 8 ++++++++ .../models/youtu/modeling_youtu.py | 1 + tests/test_fsdp_mixin.py | 20 +++++++++++++++++-- 148 files changed, 712 insertions(+), 2 deletions(-) diff --git a/src/transformers/models/afmoe/modeling_afmoe.py b/src/transformers/models/afmoe/modeling_afmoe.py index a60ce02c5f19..138df483aa15 100644 --- a/src/transformers/models/afmoe/modeling_afmoe.py +++ b/src/transformers/models/afmoe/modeling_afmoe.py @@ -616,6 +616,7 @@ def forward( class AfmoeForCausalLM(AfmoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/afmoe/modular_afmoe.py b/src/transformers/models/afmoe/modular_afmoe.py index 2200f13ce4ee..0875127a3215 100644 --- a/src/transformers/models/afmoe/modular_afmoe.py +++ b/src/transformers/models/afmoe/modular_afmoe.py @@ -396,6 +396,7 @@ def forward( class AfmoeForCausalLM(LlamaForCausalLM, AfmoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/apertus/configuration_apertus.py b/src/transformers/models/apertus/configuration_apertus.py index 33251e5f4b75..fb707e203dc6 100644 --- a/src/transformers/models/apertus/configuration_apertus.py +++ b/src/transformers/models/apertus/configuration_apertus.py @@ -76,6 +76,14 @@ class ApertusConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 de88a5e0023b..9ae78dfc3459 100644 --- a/src/transformers/models/apertus/modeling_apertus.py +++ b/src/transformers/models/apertus/modeling_apertus.py @@ -423,6 +423,7 @@ def forward( class ApertusForCausalLM(ApertusPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/arcee/configuration_arcee.py b/src/transformers/models/arcee/configuration_arcee.py index b23d249f435d..d1773d528676 100644 --- a/src/transformers/models/arcee/configuration_arcee.py +++ b/src/transformers/models/arcee/configuration_arcee.py @@ -73,6 +73,14 @@ class ArceeConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 a2681c4681a4..c0aea6228193 100644 --- a/src/transformers/models/arcee/modeling_arcee.py +++ b/src/transformers/models/arcee/modeling_arcee.py @@ -425,6 +425,7 @@ def forward( class ArceeForCausalLM(ArceePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/aria/configuration_aria.py b/src/transformers/models/aria/configuration_aria.py index 42380c73387a..986574d47c73 100644 --- a/src/transformers/models/aria/configuration_aria.py +++ b/src/transformers/models/aria/configuration_aria.py @@ -70,6 +70,14 @@ class AriaTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 aea12776c9fe..69f67e5fff9c 100644 --- a/src/transformers/models/aria/modeling_aria.py +++ b/src/transformers/models/aria/modeling_aria.py @@ -759,6 +759,7 @@ def forward( class AriaTextForCausalLM(AriaTextPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/bamba/modeling_bamba.py b/src/transformers/models/bamba/modeling_bamba.py index 7d7f7ff2185e..15bcef1ad696 100644 --- a/src/transformers/models/bamba/modeling_bamba.py +++ b/src/transformers/models/bamba/modeling_bamba.py @@ -1072,6 +1072,7 @@ def _update_mamba_mask(self, attention_mask, past_key_values): class BambaForCausalLM(BambaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/cohere/configuration_cohere.py b/src/transformers/models/cohere/configuration_cohere.py index 5b12131d72ba..da2139b85c60 100644 --- a/src/transformers/models/cohere/configuration_cohere.py +++ b/src/transformers/models/cohere/configuration_cohere.py @@ -80,6 +80,14 @@ class CohereConfig(PreTrainedConfig): "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 ed715930260d..d543e709e8f1 100644 --- a/src/transformers/models/cohere/modeling_cohere.py +++ b/src/transformers/models/cohere/modeling_cohere.py @@ -455,6 +455,7 @@ def forward( class CohereForCausalLM(CoherePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/cohere2/configuration_cohere2.py b/src/transformers/models/cohere2/configuration_cohere2.py index f596e3cee86f..548938b379e5 100644 --- a/src/transformers/models/cohere2/configuration_cohere2.py +++ b/src/transformers/models/cohere2/configuration_cohere2.py @@ -78,6 +78,14 @@ class Cohere2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 8b35fd4ab25a..24b2be0ca578 100644 --- a/src/transformers/models/cohere2/modeling_cohere2.py +++ b/src/transformers/models/cohere2/modeling_cohere2.py @@ -434,6 +434,7 @@ def forward( class Cohere2ForCausalLM(Cohere2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/cwm/configuration_cwm.py b/src/transformers/models/cwm/configuration_cwm.py index 1e086edaaa43..a5db8611589e 100644 --- a/src/transformers/models/cwm/configuration_cwm.py +++ b/src/transformers/models/cwm/configuration_cwm.py @@ -75,6 +75,14 @@ class CwmConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 862f1cfcebe7..1abeff6707a5 100644 --- a/src/transformers/models/cwm/modeling_cwm.py +++ b/src/transformers/models/cwm/modeling_cwm.py @@ -427,6 +427,7 @@ def forward( class CwmForCausalLM(CwmPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/dbrx/modeling_dbrx.py b/src/transformers/models/dbrx/modeling_dbrx.py index f009aa23d2b9..a984b4585a55 100644 --- a/src/transformers/models/dbrx/modeling_dbrx.py +++ b/src/transformers/models/dbrx/modeling_dbrx.py @@ -643,6 +643,7 @@ def load_balancing_loss_func( class DbrxForCausalLM(DbrxPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "transformer.wte.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/dbrx/modular_dbrx.py b/src/transformers/models/dbrx/modular_dbrx.py index 500d8acc5915..548af86ef4f6 100644 --- a/src/transformers/models/dbrx/modular_dbrx.py +++ b/src/transformers/models/dbrx/modular_dbrx.py @@ -431,6 +431,7 @@ def forward( class DbrxForCausalLM(DbrxPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "transformer.wte.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py b/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py index 5ab83e16ad87..e874ad881f96 100644 --- a/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py +++ b/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py @@ -94,6 +94,14 @@ class DeepseekV2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 0eb276fb30f8..bc45461bbb90 100644 --- a/src/transformers/models/deepseek_v2/modeling_deepseek_v2.py +++ b/src/transformers/models/deepseek_v2/modeling_deepseek_v2.py @@ -542,6 +542,7 @@ def forward( class DeepseekV2ForCausalLM(DeepseekV2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py index 5dcad86422f6..be557639b37f 100644 --- a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py @@ -62,6 +62,14 @@ class DeepseekV3Config(PreTrainedConfig): "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index 082ab5135496..d3ee5ae3efac 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -635,6 +635,7 @@ def forward( class DeepseekV3ForCausalLM(DeepseekV3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/deepseek_v4/configuration_deepseek_v4.py b/src/transformers/models/deepseek_v4/configuration_deepseek_v4.py index 3405f03099a1..02cd4feed68d 100644 --- a/src/transformers/models/deepseek_v4/configuration_deepseek_v4.py +++ b/src/transformers/models/deepseek_v4/configuration_deepseek_v4.py @@ -128,6 +128,14 @@ class DeepseekV4Config(PreTrainedConfig): "layers.*.mlp.experts": "moe_experts_allreduce", } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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/diffllama/modeling_diffllama.py b/src/transformers/models/diffllama/modeling_diffllama.py index 1603ba8022a0..fe8130e84116 100644 --- a/src/transformers/models/diffllama/modeling_diffllama.py +++ b/src/transformers/models/diffllama/modeling_diffllama.py @@ -661,6 +661,7 @@ def forward( class DiffLlamaForCausalLM(DiffLlamaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/doge/configuration_doge.py b/src/transformers/models/doge/configuration_doge.py index 4fbbc393fb2a..375a1c8738f7 100644 --- a/src/transformers/models/doge/configuration_doge.py +++ b/src/transformers/models/doge/configuration_doge.py @@ -70,6 +70,14 @@ class DogeConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 9bf2bce8fd4a..b998a8e0bc54 100644 --- a/src/transformers/models/dots1/configuration_dots1.py +++ b/src/transformers/models/dots1/configuration_dots1.py @@ -66,6 +66,14 @@ class Dots1Config(PreTrainedConfig): "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 e676bc0a9ba2..ba27551a74f4 100644 --- a/src/transformers/models/emu3/modeling_emu3.py +++ b/src/transformers/models/emu3/modeling_emu3.py @@ -1276,6 +1276,7 @@ def forward( class Emu3ForCausalLM(Emu3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} config: Emu3TextConfig diff --git a/src/transformers/models/ernie4_5/configuration_ernie4_5.py b/src/transformers/models/ernie4_5/configuration_ernie4_5.py index ac2d6fead084..b1d0f8c47dd8 100644 --- a/src/transformers/models/ernie4_5/configuration_ernie4_5.py +++ b/src/transformers/models/ernie4_5/configuration_ernie4_5.py @@ -76,6 +76,14 @@ class Ernie4_5Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 367ce8c88323..c98444bd38ce 100644 --- a/src/transformers/models/ernie4_5/modeling_ernie4_5.py +++ b/src/transformers/models/ernie4_5/modeling_ernie4_5.py @@ -423,6 +423,7 @@ def forward( class Ernie4_5ForCausalLM(Ernie4_5PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} 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 5967808f5607..f0ea848bff2d 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 @@ -81,6 +81,14 @@ class Ernie4_5_MoeConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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_vl_moe/configuration_ernie4_5_vl_moe.py b/src/transformers/models/ernie4_5_vl_moe/configuration_ernie4_5_vl_moe.py index cd48b6adae67..3cf9dfbe063e 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 @@ -100,6 +100,14 @@ class Ernie4_5_VLMoeTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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/eurobert/configuration_eurobert.py b/src/transformers/models/eurobert/configuration_eurobert.py index b4b4a0511c41..ab818c3ba80d 100644 --- a/src/transformers/models/eurobert/configuration_eurobert.py +++ b/src/transformers/models/eurobert/configuration_eurobert.py @@ -80,6 +80,14 @@ class EuroBertConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 149f0ef247d1..655d10ab123d 100644 --- a/src/transformers/models/eurobert/modeling_eurobert.py +++ b/src/transformers/models/eurobert/modeling_eurobert.py @@ -409,6 +409,7 @@ def forward( class EuroBertForMaskedLM(EuroBertPreTrainedModel): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/eurobert/modular_eurobert.py b/src/transformers/models/eurobert/modular_eurobert.py index 0c35455fa1c5..b40c4d043707 100644 --- a/src/transformers/models/eurobert/modular_eurobert.py +++ b/src/transformers/models/eurobert/modular_eurobert.py @@ -142,6 +142,7 @@ def forward( class EuroBertForMaskedLM(EuroBertPreTrainedModel): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/exaone4/configuration_exaone4.py b/src/transformers/models/exaone4/configuration_exaone4.py index f52652f69b4f..9aed5cc3cee3 100644 --- a/src/transformers/models/exaone4/configuration_exaone4.py +++ b/src/transformers/models/exaone4/configuration_exaone4.py @@ -91,6 +91,14 @@ class Exaone4Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 570b5e6fe160..cbe34bc5fe6c 100644 --- a/src/transformers/models/exaone4/modeling_exaone4.py +++ b/src/transformers/models/exaone4/modeling_exaone4.py @@ -441,6 +441,7 @@ def forward( class Exaone4ForCausalLM(Exaone4PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/exaone_moe/configuration_exaone_moe.py b/src/transformers/models/exaone_moe/configuration_exaone_moe.py index f76ed1c97baa..2c00f099af95 100644 --- a/src/transformers/models/exaone_moe/configuration_exaone_moe.py +++ b/src/transformers/models/exaone_moe/configuration_exaone_moe.py @@ -82,6 +82,14 @@ class ExaoneMoeConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 e5865f2fba85..78787c80cecd 100644 --- a/src/transformers/models/exaone_moe/modeling_exaone_moe.py +++ b/src/transformers/models/exaone_moe/modeling_exaone_moe.py @@ -564,6 +564,7 @@ def forward( class ExaoneMoeForCausalLM(ExaoneMoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/falcon_h1/modeling_falcon_h1.py b/src/transformers/models/falcon_h1/modeling_falcon_h1.py index f86b456fa71b..4a9c63e25cda 100644 --- a/src/transformers/models/falcon_h1/modeling_falcon_h1.py +++ b/src/transformers/models/falcon_h1/modeling_falcon_h1.py @@ -1167,6 +1167,7 @@ def _update_mamba_mask(self, attention_mask, past_key_values): class FalconH1ForCausalLM(FalconH1PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/flex_olmo/configuration_flex_olmo.py b/src/transformers/models/flex_olmo/configuration_flex_olmo.py index c80b792944a0..6727c6859227 100644 --- a/src/transformers/models/flex_olmo/configuration_flex_olmo.py +++ b/src/transformers/models/flex_olmo/configuration_flex_olmo.py @@ -61,6 +61,14 @@ class FlexOlmoConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 77d5d3203277..1884e0533c89 100644 --- a/src/transformers/models/gemma/configuration_gemma.py +++ b/src/transformers/models/gemma/configuration_gemma.py @@ -76,6 +76,14 @@ class GemmaConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 79a61722c36a..16035150b149 100644 --- a/src/transformers/models/gemma/modeling_gemma.py +++ b/src/transformers/models/gemma/modeling_gemma.py @@ -451,6 +451,7 @@ def forward( class GemmaForCausalLM(GemmaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/gemma2/configuration_gemma2.py b/src/transformers/models/gemma2/configuration_gemma2.py index 673041502934..eecd939308e8 100644 --- a/src/transformers/models/gemma2/configuration_gemma2.py +++ b/src/transformers/models/gemma2/configuration_gemma2.py @@ -80,6 +80,14 @@ class Gemma2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 7dc49c2f9e28..da784cab3c74 100644 --- a/src/transformers/models/gemma2/modeling_gemma2.py +++ b/src/transformers/models/gemma2/modeling_gemma2.py @@ -477,6 +477,7 @@ def forward( class Gemma2ForCausalLM(Gemma2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/gemma3/configuration_gemma3.py b/src/transformers/models/gemma3/configuration_gemma3.py index 0969e23f8f5e..0c29116d4e03 100644 --- a/src/transformers/models/gemma3/configuration_gemma3.py +++ b/src/transformers/models/gemma3/configuration_gemma3.py @@ -89,6 +89,14 @@ class Gemma3TextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 diff --git a/src/transformers/models/gemma3/modeling_gemma3.py b/src/transformers/models/gemma3/modeling_gemma3.py index 6c94df23cf0f..fdb9f31d2121 100644 --- a/src/transformers/models/gemma3/modeling_gemma3.py +++ b/src/transformers/models/gemma3/modeling_gemma3.py @@ -586,6 +586,7 @@ def forward( class Gemma3ForCausalLM(Gemma3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} config: Gemma3TextConfig diff --git a/src/transformers/models/gemma3n/configuration_gemma3n.py b/src/transformers/models/gemma3n/configuration_gemma3n.py index 0a20f3dd3810..08a4747f8d1f 100644 --- a/src/transformers/models/gemma3n/configuration_gemma3n.py +++ b/src/transformers/models/gemma3n/configuration_gemma3n.py @@ -93,6 +93,14 @@ class Gemma3nTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 diff --git a/src/transformers/models/gemma3n/modeling_gemma3n.py b/src/transformers/models/gemma3n/modeling_gemma3n.py index 078c73901e84..297dc14bbd15 100644 --- a/src/transformers/models/gemma3n/modeling_gemma3n.py +++ b/src/transformers/models/gemma3n/modeling_gemma3n.py @@ -1828,6 +1828,7 @@ def project_per_layer_inputs( class Gemma3nForCausalLM(Gemma3nPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} config: Gemma3nTextConfig diff --git a/src/transformers/models/gemma4/configuration_gemma4.py b/src/transformers/models/gemma4/configuration_gemma4.py index ebb6ccab2d76..f447f9b13519 100644 --- a/src/transformers/models/gemma4/configuration_gemma4.py +++ b/src/transformers/models/gemma4/configuration_gemma4.py @@ -151,6 +151,14 @@ class Gemma4TextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 diff --git a/src/transformers/models/gemma4/modeling_gemma4.py b/src/transformers/models/gemma4/modeling_gemma4.py index 11ecb36da63e..61ee14f0f6b2 100644 --- a/src/transformers/models/gemma4/modeling_gemma4.py +++ b/src/transformers/models/gemma4/modeling_gemma4.py @@ -1794,6 +1794,7 @@ def project_per_layer_inputs( class Gemma4ForCausalLM(Gemma4PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} config: Gemma4TextConfig diff --git a/src/transformers/models/gemma4_assistant/modeling_gemma4_assistant.py b/src/transformers/models/gemma4_assistant/modeling_gemma4_assistant.py index 8b7b6b1a8b67..4c300686a0e7 100644 --- a/src/transformers/models/gemma4_assistant/modeling_gemma4_assistant.py +++ b/src/transformers/models/gemma4_assistant/modeling_gemma4_assistant.py @@ -110,6 +110,7 @@ def _init_weights(self, module): class Gemma4AssistantForCausalLM(Gemma4AssistantPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_gather_output"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: Gemma4AssistantConfig): diff --git a/src/transformers/models/glm/configuration_glm.py b/src/transformers/models/glm/configuration_glm.py index 52197a8d596f..c8cdee74019a 100644 --- a/src/transformers/models/glm/configuration_glm.py +++ b/src/transformers/models/glm/configuration_glm.py @@ -67,6 +67,14 @@ class GlmConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 e87e10841bab..2fbb4e4fb64d 100644 --- a/src/transformers/models/glm/modeling_glm.py +++ b/src/transformers/models/glm/modeling_glm.py @@ -440,6 +440,7 @@ def forward( class GlmForCausalLM(GlmPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/glm4/configuration_glm4.py b/src/transformers/models/glm4/configuration_glm4.py index d4ab4053f93d..4eb5904d9ce9 100644 --- a/src/transformers/models/glm4/configuration_glm4.py +++ b/src/transformers/models/glm4/configuration_glm4.py @@ -67,6 +67,14 @@ class Glm4Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 5c4eb666da85..baf0a0e634ef 100644 --- a/src/transformers/models/glm4/modeling_glm4.py +++ b/src/transformers/models/glm4/modeling_glm4.py @@ -445,6 +445,7 @@ def forward( class Glm4ForCausalLM(Glm4PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/glm4_moe/configuration_glm4_moe.py b/src/transformers/models/glm4_moe/configuration_glm4_moe.py index b08ba1db6d6f..7b7ea57267f0 100644 --- a/src/transformers/models/glm4_moe/configuration_glm4_moe.py +++ b/src/transformers/models/glm4_moe/configuration_glm4_moe.py @@ -71,6 +71,14 @@ class Glm4MoeConfig(PreTrainedConfig): "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 aea18b1dc715..bab7e3a6084a 100644 --- a/src/transformers/models/glm4_moe/modeling_glm4_moe.py +++ b/src/transformers/models/glm4_moe/modeling_glm4_moe.py @@ -578,6 +578,7 @@ def forward( class Glm4MoeForCausalLM(Glm4MoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} 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 3885caa01e6e..8fb2a0fbf3ec 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 @@ -65,6 +65,14 @@ class Glm4MoeLiteConfig(PreTrainedConfig): "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 5b30e546217e..ba44c049339d 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,6 +652,7 @@ def forward( class Glm4MoeLiteForCausalLM(Glm4MoeLitePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/glm4v/configuration_glm4v.py b/src/transformers/models/glm4v/configuration_glm4v.py index 98200d3e3495..e7e6bb9927f2 100644 --- a/src/transformers/models/glm4v/configuration_glm4v.py +++ b/src/transformers/models/glm4v/configuration_glm4v.py @@ -102,6 +102,14 @@ class Glm4vTextConfig(PreTrainedConfig): "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 e5cbbdf2872c..be09d149e7bf 100644 --- a/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py +++ b/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py @@ -66,6 +66,14 @@ class Glm4vMoeTextConfig(PreTrainedConfig): "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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/glm_image/configuration_glm_image.py b/src/transformers/models/glm_image/configuration_glm_image.py index 94f7c61dc4fa..8c3c985ff9f0 100644 --- a/src/transformers/models/glm_image/configuration_glm_image.py +++ b/src/transformers/models/glm_image/configuration_glm_image.py @@ -115,6 +115,14 @@ class GlmImageTextConfig(PreTrainedConfig): "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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_moe_dsa/configuration_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py index 818d82eca4f3..0726f645bb77 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 @@ -75,6 +75,14 @@ class GlmMoeDsaConfig(PreTrainedConfig): "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } attribute_map = { "num_local_experts": "n_routed_experts", "head_dim": "qk_rope_head_dim", diff --git a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py index 62058f67b264..9eb5b94a9cba 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 @@ -822,6 +822,7 @@ def forward( class GlmMoeDsaForCausalLM(GlmMoeDsaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/glm_ocr/configuration_glm_ocr.py b/src/transformers/models/glm_ocr/configuration_glm_ocr.py index dcd86079241d..1b2627926de7 100644 --- a/src/transformers/models/glm_ocr/configuration_glm_ocr.py +++ b/src/transformers/models/glm_ocr/configuration_glm_ocr.py @@ -103,6 +103,14 @@ class GlmOcrTextConfig(PreTrainedConfig): "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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/gpt_neox/configuration_gpt_neox.py b/src/transformers/models/gpt_neox/configuration_gpt_neox.py index 5b4b7908e213..c62775b283bb 100644 --- a/src/transformers/models/gpt_neox/configuration_gpt_neox.py +++ b/src/transformers/models/gpt_neox/configuration_gpt_neox.py @@ -67,6 +67,14 @@ class GPTNeoXConfig(PreTrainedConfig): "final_layer_norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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_oss/configuration_gpt_oss.py b/src/transformers/models/gpt_oss/configuration_gpt_oss.py index 3a9ca00c8e9d..c7cf246fe259 100644 --- a/src/transformers/models/gpt_oss/configuration_gpt_oss.py +++ b/src/transformers/models/gpt_oss/configuration_gpt_oss.py @@ -32,6 +32,11 @@ 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", diff --git a/src/transformers/models/granite/configuration_granite.py b/src/transformers/models/granite/configuration_granite.py index 063452fbe08d..11489cd76904 100644 --- a/src/transformers/models/granite/configuration_granite.py +++ b/src/transformers/models/granite/configuration_granite.py @@ -76,6 +76,14 @@ class GraniteConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 68c8877446d4..2b6066eb8a63 100644 --- a/src/transformers/models/granite/modeling_granite.py +++ b/src/transformers/models/granite/modeling_granite.py @@ -446,6 +446,7 @@ def forward( class GraniteForCausalLM(GranitePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/granite4_vision/configuration_granite4_vision.py b/src/transformers/models/granite4_vision/configuration_granite4_vision.py index 8d9477dad556..5986d4f512b6 100644 --- a/src/transformers/models/granite4_vision/configuration_granite4_vision.py +++ b/src/transformers/models/granite4_vision/configuration_granite4_vision.py @@ -79,6 +79,14 @@ class Granite4VisionTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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/helium/configuration_helium.py b/src/transformers/models/helium/configuration_helium.py index 15cd1b77db0f..39b39dce0562 100644 --- a/src/transformers/models/helium/configuration_helium.py +++ b/src/transformers/models/helium/configuration_helium.py @@ -70,6 +70,14 @@ class HeliumConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 4463cc02980a..4668475fdb84 100644 --- a/src/transformers/models/helium/modeling_helium.py +++ b/src/transformers/models/helium/modeling_helium.py @@ -424,6 +424,7 @@ def forward( class HeliumForCausalLM(HeliumPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} 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 38e57e716c9a..4c36b638b32d 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 @@ -88,6 +88,14 @@ class HiggsAudioV2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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/hunyuan_v1_dense/modeling_hunyuan_v1_dense.py b/src/transformers/models/hunyuan_v1_dense/modeling_hunyuan_v1_dense.py index 79b21c72f890..aaeac78158b0 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 @@ -462,6 +462,7 @@ def forward( class HunYuanDenseV1ForCausalLM(HunYuanDenseV1PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} 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 9bb57e90bca1..0bf0dcf96b11 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 @@ -551,6 +551,7 @@ def forward( class HunYuanMoEV1ForCausalLM(HunYuanMoEV1PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/hy_v3/configuration_hy_v3.py b/src/transformers/models/hy_v3/configuration_hy_v3.py index 9ed4c5cc81c6..5d20419b9d0c 100644 --- a/src/transformers/models/hy_v3/configuration_hy_v3.py +++ b/src/transformers/models/hy_v3/configuration_hy_v3.py @@ -73,6 +73,14 @@ class HYV3Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 92499a1d9609..d1970f4fe8ab 100644 --- a/src/transformers/models/hy_v3/modeling_hy_v3.py +++ b/src/transformers/models/hy_v3/modeling_hy_v3.py @@ -545,6 +545,7 @@ def forward( class HYV3ForCausalLM(HYV3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/hyperclovax/configuration_hyperclovax.py b/src/transformers/models/hyperclovax/configuration_hyperclovax.py index b0b4c64d11bd..6dcbcb7576d8 100644 --- a/src/transformers/models/hyperclovax/configuration_hyperclovax.py +++ b/src/transformers/models/hyperclovax/configuration_hyperclovax.py @@ -91,6 +91,14 @@ class HyperCLOVAXConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 f314eadb363d..d228f1cae263 100644 --- a/src/transformers/models/hyperclovax/modeling_hyperclovax.py +++ b/src/transformers/models/hyperclovax/modeling_hyperclovax.py @@ -453,6 +453,7 @@ def forward( class HyperCLOVAXForCausalLM(HyperCLOVAXPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/jais2/configuration_jais2.py b/src/transformers/models/jais2/configuration_jais2.py index 5e0fa934f7a9..b03692a2342d 100644 --- a/src/transformers/models/jais2/configuration_jais2.py +++ b/src/transformers/models/jais2/configuration_jais2.py @@ -74,6 +74,14 @@ class Jais2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 0aec15924b03..d7d4ce972b01 100644 --- a/src/transformers/models/jais2/modeling_jais2.py +++ b/src/transformers/models/jais2/modeling_jais2.py @@ -398,6 +398,7 @@ def forward( class Jais2ForCausalLM(Jais2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} 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 80f5f756b98b..87460c65e115 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 @@ -875,6 +875,7 @@ def forward( class KyutaiSpeechToTextForConditionalGeneration(KyutaiSpeechToTextPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _keep_in_fp32_modules_strict = ["codec_model"] diff --git a/src/transformers/models/laguna/configuration_laguna.py b/src/transformers/models/laguna/configuration_laguna.py index dfe403281263..2843fd449da5 100644 --- a/src/transformers/models/laguna/configuration_laguna.py +++ b/src/transformers/models/laguna/configuration_laguna.py @@ -80,6 +80,14 @@ class LagunaConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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/lfm2/modeling_lfm2.py b/src/transformers/models/lfm2/modeling_lfm2.py index c8d42b555903..227afd2a35a5 100644 --- a/src/transformers/models/lfm2/modeling_lfm2.py +++ b/src/transformers/models/lfm2/modeling_lfm2.py @@ -541,6 +541,7 @@ def forward( class Lfm2ForCausalLM(Lfm2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py b/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py index ed6df4a3e34b..6e62546d4ee3 100644 --- a/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py +++ b/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py @@ -631,6 +631,7 @@ def forward( class Lfm2MoeForCausalLM(Lfm2MoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/llama/configuration_llama.py b/src/transformers/models/llama/configuration_llama.py index 73b898d619bd..c6698cb26cc5 100644 --- a/src/transformers/models/llama/configuration_llama.py +++ b/src/transformers/models/llama/configuration_llama.py @@ -76,6 +76,14 @@ class LlamaConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 c12f5e1966b1..655e542358f5 100644 --- a/src/transformers/models/llama/modeling_llama.py +++ b/src/transformers/models/llama/modeling_llama.py @@ -431,6 +431,7 @@ class LlamaForCausalLM(LlamaPreTrainedModel, GenerationMixin): _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/llama4/modeling_llama4.py b/src/transformers/models/llama4/modeling_llama4.py index a7ec69455854..094d1bb1f7fa 100644 --- a/src/transformers/models/llama4/modeling_llama4.py +++ b/src/transformers/models/llama4/modeling_llama4.py @@ -591,6 +591,7 @@ class Llama4ForCausalLM(Llama4PreTrainedModel, GenerationMixin): base_model_prefix = "language_model" _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} 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 fcdcbb8cae28..e083bf32466a 100644 --- a/src/transformers/models/longcat_flash/configuration_longcat_flash.py +++ b/src/transformers/models/longcat_flash/configuration_longcat_flash.py @@ -68,6 +68,14 @@ class LongcatFlashConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 443a860dcaae..b5e6f889bfe5 100644 --- a/src/transformers/models/longcat_flash/modeling_longcat_flash.py +++ b/src/transformers/models/longcat_flash/modeling_longcat_flash.py @@ -651,6 +651,7 @@ def forward( class LongcatFlashForCausalLM(LongcatFlashPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _keys_to_ignore_on_load_unexpected = [r"model\.mtp.*"] diff --git a/src/transformers/models/minimax/configuration_minimax.py b/src/transformers/models/minimax/configuration_minimax.py index e53ecd12617f..249a9ae2a3e0 100644 --- a/src/transformers/models/minimax/configuration_minimax.py +++ b/src/transformers/models/minimax/configuration_minimax.py @@ -73,6 +73,14 @@ class MiniMaxConfig(PreTrainedConfig): "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 2f360c509c28..1c14670d722a 100644 --- a/src/transformers/models/minimax_m2/configuration_minimax_m2.py +++ b/src/transformers/models/minimax_m2/configuration_minimax_m2.py @@ -72,6 +72,14 @@ class MiniMaxM2Config(PreTrainedConfig): "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } attribute_map = { "num_experts": "num_local_experts", } diff --git a/src/transformers/models/ministral/configuration_ministral.py b/src/transformers/models/ministral/configuration_ministral.py index 84d69e92d7f3..8f6d17ca7a17 100644 --- a/src/transformers/models/ministral/configuration_ministral.py +++ b/src/transformers/models/ministral/configuration_ministral.py @@ -78,6 +78,14 @@ class MinistralConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 510364445c3c..f9ae6841c06b 100644 --- a/src/transformers/models/ministral/modeling_ministral.py +++ b/src/transformers/models/ministral/modeling_ministral.py @@ -430,6 +430,7 @@ def forward( class MinistralForCausalLM(MinistralPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/ministral3/configuration_ministral3.py b/src/transformers/models/ministral3/configuration_ministral3.py index 21a9a6515fba..4afb236b1504 100644 --- a/src/transformers/models/ministral3/configuration_ministral3.py +++ b/src/transformers/models/ministral3/configuration_ministral3.py @@ -83,6 +83,14 @@ class Ministral3Config(PreTrainedConfig): "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 bfc4d35df5f9..0218a8b4c096 100644 --- a/src/transformers/models/ministral3/modeling_ministral3.py +++ b/src/transformers/models/ministral3/modeling_ministral3.py @@ -413,6 +413,7 @@ def forward( class Ministral3ForCausalLM(Ministral3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/mistral/configuration_mistral.py b/src/transformers/models/mistral/configuration_mistral.py index 27d125446329..fb4f4919ac8b 100644 --- a/src/transformers/models/mistral/configuration_mistral.py +++ b/src/transformers/models/mistral/configuration_mistral.py @@ -75,6 +75,14 @@ class MistralConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 4e4b24132f66..ea1cbd5647a4 100644 --- a/src/transformers/models/mistral/modeling_mistral.py +++ b/src/transformers/models/mistral/modeling_mistral.py @@ -402,6 +402,7 @@ def forward( class MistralForCausalLM(MistralPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/mistral4/configuration_mistral4.py b/src/transformers/models/mistral4/configuration_mistral4.py index 774b61016f38..ddb619cdf693 100644 --- a/src/transformers/models/mistral4/configuration_mistral4.py +++ b/src/transformers/models/mistral4/configuration_mistral4.py @@ -60,6 +60,14 @@ class Mistral4Config(PreTrainedConfig): "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 a31527ce4b4d..1f81bd21912f 100644 --- a/src/transformers/models/mistral4/modeling_mistral4.py +++ b/src/transformers/models/mistral4/modeling_mistral4.py @@ -641,6 +641,7 @@ def forward( class Mistral4ForCausalLM(Mistral4PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/nanochat/modeling_nanochat.py b/src/transformers/models/nanochat/modeling_nanochat.py index e4ea462797e9..285578dde755 100644 --- a/src/transformers/models/nanochat/modeling_nanochat.py +++ b/src/transformers/models/nanochat/modeling_nanochat.py @@ -433,6 +433,7 @@ def forward( class NanoChatForCausalLM(NanoChatPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/nanochat/modular_nanochat.py b/src/transformers/models/nanochat/modular_nanochat.py index 6bc3ddba6457..460cf321150f 100644 --- a/src/transformers/models/nanochat/modular_nanochat.py +++ b/src/transformers/models/nanochat/modular_nanochat.py @@ -199,6 +199,7 @@ def forward( @auto_docstring class NanoChatForCausalLM(Gemma2ForCausalLM): _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} def forward(self, **super_kwargs) -> CausalLMOutputWithPast: diff --git a/src/transformers/models/olmo/configuration_olmo.py b/src/transformers/models/olmo/configuration_olmo.py index 7f8fcfa2c168..32af115cfbe3 100644 --- a/src/transformers/models/olmo/configuration_olmo.py +++ b/src/transformers/models/olmo/configuration_olmo.py @@ -79,6 +79,14 @@ class OlmoConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 937f138c4c11..943e4751aa9b 100644 --- a/src/transformers/models/olmo/modeling_olmo.py +++ b/src/transformers/models/olmo/modeling_olmo.py @@ -426,6 +426,7 @@ def forward( class OlmoForCausalLM(OlmoPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/olmo2/configuration_olmo2.py b/src/transformers/models/olmo2/configuration_olmo2.py index 408e15b4bffd..8da1d8c8fee8 100644 --- a/src/transformers/models/olmo2/configuration_olmo2.py +++ b/src/transformers/models/olmo2/configuration_olmo2.py @@ -84,6 +84,14 @@ class Olmo2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 6312fdbda9a1..77d40756513c 100644 --- a/src/transformers/models/olmo2/modeling_olmo2.py +++ b/src/transformers/models/olmo2/modeling_olmo2.py @@ -430,6 +430,7 @@ def forward( class Olmo2ForCausalLM(Olmo2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/olmo3/configuration_olmo3.py b/src/transformers/models/olmo3/configuration_olmo3.py index c4597535850c..6d34f887875a 100644 --- a/src/transformers/models/olmo3/configuration_olmo3.py +++ b/src/transformers/models/olmo3/configuration_olmo3.py @@ -79,6 +79,14 @@ class Olmo3Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 b98e267cde37..5a1d57a7ab2e 100644 --- a/src/transformers/models/olmo3/modeling_olmo3.py +++ b/src/transformers/models/olmo3/modeling_olmo3.py @@ -434,6 +434,7 @@ def forward( class Olmo3ForCausalLM(Olmo3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py index a6eabbb7b99b..3d83f8548bb2 100644 --- a/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py +++ b/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py @@ -88,6 +88,14 @@ class OlmoHybridConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 5306a92f9955..322dd780abf0 100644 --- a/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py +++ b/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py @@ -1046,6 +1046,7 @@ def _update_linear_attn_mask(self, attention_mask, past_key_values): class OlmoHybridForCausalLM(OlmoHybridPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/olmoe/configuration_olmoe.py b/src/transformers/models/olmoe/configuration_olmoe.py index 33258ec2c05f..ef7077ae6867 100644 --- a/src/transformers/models/olmoe/configuration_olmoe.py +++ b/src/transformers/models/olmoe/configuration_olmoe.py @@ -66,6 +66,14 @@ class OlmoeConfig(PreTrainedConfig): "norm": "activation", } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 50304 hidden_size: int = 2048 intermediate_size: int = 2048 diff --git a/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py b/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py index 83991ac90f1c..b7e4a34e370b 100644 --- a/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py +++ b/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py @@ -124,6 +124,14 @@ class PaddleOCRTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 d10e30019bfd..03326dad422a 100644 --- a/src/transformers/models/phi/configuration_phi.py +++ b/src/transformers/models/phi/configuration_phi.py @@ -74,6 +74,14 @@ class PhiConfig(PreTrainedConfig): "final_layernorm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 diff --git a/src/transformers/models/phi/modeling_phi.py b/src/transformers/models/phi/modeling_phi.py index 70a0bb6d6bba..c628d40cf57f 100644 --- a/src/transformers/models/phi/modeling_phi.py +++ b/src/transformers/models/phi/modeling_phi.py @@ -407,6 +407,7 @@ def forward( class PhiForCausalLM(PhiPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/phi3/configuration_phi3.py b/src/transformers/models/phi3/configuration_phi3.py index e2d4d8934b05..c49823521c66 100644 --- a/src/transformers/models/phi3/configuration_phi3.py +++ b/src/transformers/models/phi3/configuration_phi3.py @@ -70,6 +70,14 @@ class Phi3Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 41d6634095b7..895426414d78 100644 --- a/src/transformers/models/phi3/modeling_phi3.py +++ b/src/transformers/models/phi3/modeling_phi3.py @@ -433,6 +433,7 @@ def forward( class Phi3ForCausalLM(Phi3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py b/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py index db4cbd0eab6b..f35ba25428c5 100644 --- a/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py +++ b/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py @@ -204,6 +204,14 @@ class Phi4MultimodalConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 5c723f3d4365..a8c05cd314fa 100644 --- a/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py +++ b/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py @@ -1597,6 +1597,7 @@ def forward( class Phi4MultimodalForCausalLM(Phi4MultimodalPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} 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/qwen2/configuration_qwen2.py b/src/transformers/models/qwen2/configuration_qwen2.py index 599511b903a7..9dcdbf2b259e 100644 --- a/src/transformers/models/qwen2/configuration_qwen2.py +++ b/src/transformers/models/qwen2/configuration_qwen2.py @@ -73,6 +73,14 @@ class Qwen2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 e22875b8c8a0..2767340db60e 100644 --- a/src/transformers/models/qwen2/modeling_qwen2.py +++ b/src/transformers/models/qwen2/modeling_qwen2.py @@ -417,6 +417,7 @@ def forward( class Qwen2ForCausalLM(Qwen2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} 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 b6ca7110bb17..15c3b509a812 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 @@ -167,6 +167,14 @@ class Qwen2_5OmniTextConfig(PreTrainedConfig): "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 c743f48e7cda..bf2723debf08 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 @@ -101,6 +101,14 @@ class Qwen2_5_VLTextConfig(PreTrainedConfig): "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 5585a7454314..3066335ce3f0 100644 --- a/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py +++ b/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py @@ -68,6 +68,14 @@ class Qwen2MoeConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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/modular_qwen2_moe.py b/src/transformers/models/qwen2_moe/modular_qwen2_moe.py index 8a5e9fb7751b..9c5b65f47896 100644 --- a/src/transformers/models/qwen2_moe/modular_qwen2_moe.py +++ b/src/transformers/models/qwen2_moe/modular_qwen2_moe.py @@ -231,6 +231,7 @@ def forward( class Qwen2MoeForCausalLM(MixtralForCausalLM, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py b/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py index cf9629d6beb9..139d425c50bb 100644 --- a/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py +++ b/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py @@ -78,6 +78,14 @@ class Qwen2VLTextConfig(PreTrainedConfig): "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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_5/configuration_qwen3_5.py b/src/transformers/models/qwen3_5/configuration_qwen3_5.py index 4f31c0c04a2f..957a99ba07f0 100644 --- a/src/transformers/models/qwen3_5/configuration_qwen3_5.py +++ b/src/transformers/models/qwen3_5/configuration_qwen3_5.py @@ -71,6 +71,14 @@ class Qwen3_5TextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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_moe/configuration_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/configuration_qwen3_5_moe.py index b7187b0c795d..baab11721f90 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 @@ -72,6 +72,14 @@ class Qwen3_5MoeTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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/modular_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/modular_qwen3_5_moe.py index 8000612aa275..58591c6f2648 100644 --- a/src/transformers/models/qwen3_5_moe/modular_qwen3_5_moe.py +++ b/src/transformers/models/qwen3_5_moe/modular_qwen3_5_moe.py @@ -245,6 +245,7 @@ def __init__(self, config): class Qwen3_5MoeForConditionalGeneration(Qwen3VLMoeForConditionalGeneration): _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} def forward(self, **super_kwargs): diff --git a/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py b/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py index 7009dfe8f835..87a8ac165d0c 100644 --- a/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py +++ b/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py @@ -96,6 +96,14 @@ class Qwen3MoeConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 151936 hidden_size: int = 2048 intermediate_size: int = 6144 diff --git a/src/transformers/models/qwen3_next/configuration_qwen3_next.py b/src/transformers/models/qwen3_next/configuration_qwen3_next.py index 1d1471d1db8c..25810a96a0fa 100644 --- a/src/transformers/models/qwen3_next/configuration_qwen3_next.py +++ b/src/transformers/models/qwen3_next/configuration_qwen3_next.py @@ -77,6 +77,14 @@ class Qwen3NextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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_vl_moe/configuration_qwen3_vl_moe.py b/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py index 0b4573b99942..612edcd5776e 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 @@ -94,6 +94,14 @@ class Qwen3VLMoeTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 151936 hidden_size: int = 2048 diff --git a/src/transformers/models/seed_oss/configuration_seed_oss.py b/src/transformers/models/seed_oss/configuration_seed_oss.py index d60348c19149..f1145ba1f7a9 100644 --- a/src/transformers/models/seed_oss/configuration_seed_oss.py +++ b/src/transformers/models/seed_oss/configuration_seed_oss.py @@ -74,6 +74,14 @@ class SeedOssConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 fa524a9de3b2..b8bda4449d72 100644 --- a/src/transformers/models/seed_oss/modeling_seed_oss.py +++ b/src/transformers/models/seed_oss/modeling_seed_oss.py @@ -430,6 +430,7 @@ def forward( class SeedOssForCausalLM(SeedOssPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/smollm3/configuration_smollm3.py b/src/transformers/models/smollm3/configuration_smollm3.py index 5de4e2c1cedf..637f9ad75fb8 100644 --- a/src/transformers/models/smollm3/configuration_smollm3.py +++ b/src/transformers/models/smollm3/configuration_smollm3.py @@ -84,6 +84,14 @@ class SmolLM3Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 d9bb9d7473b5..acb03a2ead2a 100644 --- a/src/transformers/models/smollm3/modeling_smollm3.py +++ b/src/transformers/models/smollm3/modeling_smollm3.py @@ -446,6 +446,7 @@ def forward( class SmolLM3ForCausalLM(SmolLM3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/solar_open/configuration_solar_open.py b/src/transformers/models/solar_open/configuration_solar_open.py index 685ea1cbe4a5..1d78ce714d70 100644 --- a/src/transformers/models/solar_open/configuration_solar_open.py +++ b/src/transformers/models/solar_open/configuration_solar_open.py @@ -92,6 +92,14 @@ class SolarOpenConfig(PreTrainedConfig): "layers.*.mlp.experts": "moe_experts_allreduce", "norm": "activation", } + + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } 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 9d36d2716d70..e5f0c9cfafe6 100644 --- a/src/transformers/models/solar_open/modeling_solar_open.py +++ b/src/transformers/models/solar_open/modeling_solar_open.py @@ -553,6 +553,7 @@ def forward( class SolarOpenForCausalLM(SolarOpenPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/starcoder2/configuration_starcoder2.py b/src/transformers/models/starcoder2/configuration_starcoder2.py index f508dcce3cbd..ea6310ca5278 100644 --- a/src/transformers/models/starcoder2/configuration_starcoder2.py +++ b/src/transformers/models/starcoder2/configuration_starcoder2.py @@ -72,6 +72,14 @@ class Starcoder2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 fff6c00cdd7d..caac32e2cebf 100644 --- a/src/transformers/models/starcoder2/modeling_starcoder2.py +++ b/src/transformers/models/starcoder2/modeling_starcoder2.py @@ -410,6 +410,7 @@ def forward( class Starcoder2ForCausalLM(Starcoder2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/t5gemma/configuration_t5gemma.py b/src/transformers/models/t5gemma/configuration_t5gemma.py index d5e97d034d73..3ad3a5271f9d 100644 --- a/src/transformers/models/t5gemma/configuration_t5gemma.py +++ b/src/transformers/models/t5gemma/configuration_t5gemma.py @@ -80,6 +80,14 @@ class T5GemmaModuleConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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/t5gemma2/configuration_t5gemma2.py b/src/transformers/models/t5gemma2/configuration_t5gemma2.py index b90174f4f9ec..2864e9daa240 100644 --- a/src/transformers/models/t5gemma2/configuration_t5gemma2.py +++ b/src/transformers/models/t5gemma2/configuration_t5gemma2.py @@ -76,6 +76,14 @@ class T5Gemma2TextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 diff --git a/src/transformers/models/vaultgemma/configuration_vaultgemma.py b/src/transformers/models/vaultgemma/configuration_vaultgemma.py index be5737835253..7c7654330b37 100644 --- a/src/transformers/models/vaultgemma/configuration_vaultgemma.py +++ b/src/transformers/models/vaultgemma/configuration_vaultgemma.py @@ -79,6 +79,14 @@ class VaultGemmaConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + 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 e392829c8fad..95cebf857c56 100644 --- a/src/transformers/models/vaultgemma/modeling_vaultgemma.py +++ b/src/transformers/models/vaultgemma/modeling_vaultgemma.py @@ -468,6 +468,7 @@ def forward( class VaultGemmaForCausalLM(VaultGemmaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/voxtral_realtime/configuration_voxtral_realtime.py b/src/transformers/models/voxtral_realtime/configuration_voxtral_realtime.py index 48d1035b7243..f47df59dd6a4 100644 --- a/src/transformers/models/voxtral_realtime/configuration_voxtral_realtime.py +++ b/src/transformers/models/voxtral_realtime/configuration_voxtral_realtime.py @@ -41,6 +41,14 @@ class VoxtralRealtimeTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 32000 hidden_size: int = 4096 intermediate_size: int = 14336 diff --git a/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py b/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py index bff9844eb1d3..7f8965a6de70 100644 --- a/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py +++ b/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py @@ -831,6 +831,7 @@ def forward( class VoxtralRealtimeTextForCausalLM(VoxtralRealtimeTextPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/youtu/configuration_youtu.py b/src/transformers/models/youtu/configuration_youtu.py index 6da28284a577..726df1cf3c94 100644 --- a/src/transformers/models/youtu/configuration_youtu.py +++ b/src/transformers/models/youtu/configuration_youtu.py @@ -60,6 +60,14 @@ class YoutuConfig(PreTrainedConfig): "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks + # this dict to decide what fully_shard each module gets. + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } attribute_map = {} vocab_size: int = 128256 diff --git a/src/transformers/models/youtu/modeling_youtu.py b/src/transformers/models/youtu/modeling_youtu.py index e76855d5f518..820146bf2e0a 100644 --- a/src/transformers/models/youtu/modeling_youtu.py +++ b/src/transformers/models/youtu/modeling_youtu.py @@ -534,6 +534,7 @@ def forward( class YoutuForCausalLM(YoutuPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/tests/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py index 32008c74eefb..a9f46514f642 100644 --- a/tests/test_fsdp_mixin.py +++ b/tests/test_fsdp_mixin.py @@ -64,10 +64,26 @@ LR = 3e-4 SEED = 42 FSDP_TOP_MODEL_NAMES = { - # FSDP coverage is gated on models declaring `base_model_fsdp_plan` + class-level - # `_fsdp_plan`. Roll out to more models in follow-up PRs. + # 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", + "glm_moe_dsa", + "glm4_moe_lite", } From 41652094fed6d5a6d8a65407c1a510be9d390c24 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Sat, 16 May 2026 10:20:57 +0000 Subject: [PATCH 100/116] update fsdp mixin tests --- tests/test_fsdp_mixin.py | 77 ++++++++++++++-------------------------- 1 file changed, 26 insertions(+), 51 deletions(-) diff --git a/tests/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py index a9f46514f642..e60082a142db 100644 --- a/tests/test_fsdp_mixin.py +++ b/tests/test_fsdp_mixin.py @@ -42,16 +42,17 @@ if is_torch_available(): import torch import torch.distributed as dist - import torch.distributed.checkpoint as dcp import torch.multiprocessing as mp - from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner - from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict - from torch.distributed.tensor import DTensor from torch.nn.parallel import DistributedDataParallel as DDP from transformers.distributed import DistributedConfig from transformers.distributed.fsdp import apply_fully_shard_data_parallel, initialize_fsdp from transformers.distributed.tensor_parallel import replace_layer_number_by_wildcard + from transformers.distributed.utils import ( + gather_full_state_dict, + load_optimizer_distributed, + save_optimizer_distributed, + ) # ============================================================================= @@ -210,18 +211,11 @@ def _create_shared_tmpdir(rank): return tmpdir_list[0], tmpdir_obj -def _gather_fsdp2_state_dict(model): - """Gather FSDP2 sharded parameters into full tensors via DTensor.full_tensor().""" - state_dict = {} - for name, tensor in model.state_dict().items(): - if isinstance(tensor, DTensor): - state_dict[name] = tensor.full_tensor().clone().detach().cpu() - else: - state_dict[name] = tensor.clone().detach().cpu() - return state_dict - - 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()} @@ -297,43 +291,24 @@ def _save_init_pretrained(rank, config, dtype): def _save_training_state(model, optimizer, training_state_dir): - """Save optimizer + RNG states as distcp (for training resume only).""" - _, optim_sd = get_state_dict(model, optimizer) - training_state = { - "optim": optim_sd, - "cpu_rng_state": torch.get_rng_state(), - } - accelerator_rng_state = _get_accelerator_rng_state() - if accelerator_rng_state is not None: - training_state["accelerator_rng_state"] = accelerator_rng_state - dcp.save(training_state, checkpoint_id=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): - """Load optimizer + RNG states from distcp (model weights loaded separately via from_pretrained).""" - model_sd, optim_sd = get_state_dict(model, optimizer) - loaded_training_state = { - "optim": optim_sd, - "cpu_rng_state": torch.empty_like(torch.get_rng_state()), - } - accelerator_rng_state = _get_accelerator_rng_state() - if accelerator_rng_state is not None: - loaded_training_state["accelerator_rng_state"] = torch.empty_like(accelerator_rng_state) - # MoE models can have sparse optimizer state (experts not selected yet), so - # allow partial optimizer key restoration instead of failing hard on missing keys. - dcp.load( - loaded_training_state, - checkpoint_id=training_state_dir, - planner=DefaultLoadPlanner(allow_partial_load=True), - ) - set_state_dict( - model, - optimizer, - model_state_dict=model_sd, - optim_state_dict=loaded_training_state["optim"], + """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(loaded_training_state["cpu_rng_state"]) - _set_accelerator_rng_state(loaded_training_state.get("accelerator_rng_state")) + 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): @@ -451,7 +426,7 @@ def train_fsdp2( combined_losses = pre_ckpt_losses + post_ckpt_losses combined_grad_norms = pre_ckpt_grad_norms + post_ckpt_grad_norms - combined_state_dict = _gather_fsdp2_state_dict(resumed_model) + combined_state_dict = gather_full_state_dict(resumed_model) return combined_losses, combined_grad_norms, combined_state_dict @@ -491,7 +466,7 @@ def _test_fsdp2_save_load_impl(rank, config_class, config_dict): output.loss.backward() optimizer.step() - state_dict_before = _gather_fsdp2_state_dict(model) + state_dict_before = gather_full_state_dict(model) tmpdir, tmpdir_obj = _create_shared_tmpdir(rank) try: @@ -508,7 +483,7 @@ def _test_fsdp2_save_load_impl(rank, config_class, config_dict): if rank == 0: tmpdir_obj.cleanup() - state_dict_after = _gather_fsdp2_state_dict(new_model) + 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" From 6f5dbfb39f1c9cd2c627c443993a07839cad2635 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Sun, 17 May 2026 01:43:38 +0000 Subject: [PATCH 101/116] linting --- tests/test_fsdp_mixin.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py index e60082a142db..c1b3b84f80b6 100644 --- a/tests/test_fsdp_mixin.py +++ b/tests/test_fsdp_mixin.py @@ -303,9 +303,7 @@ def _save_training_state(model, optimizer, training_state_dir): 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 - ) + 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"]) From d075668dd78732f256c9d87750980f4e9e5d4474 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Mon, 18 May 2026 01:03:22 +0000 Subject: [PATCH 102/116] fix test fsdp --- tests/utils/test_modeling_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/utils/test_modeling_utils.py b/tests/utils/test_modeling_utils.py index 296f6662ac53..ba30b775588c 100644 --- a/tests/utils/test_modeling_utils.py +++ b/tests/utils/test_modeling_utils.py @@ -446,7 +446,7 @@ def test_model_from_pretrained_fsdp_distributes_before_loading(self): def fake_apply_fsdp(model, fsdp_mesh, fsdp_plan): call_order.append("distribute") self.assertIs(fsdp_mesh, fake_mesh) - self.assertEqual(fsdp_plan, "auto") + self.assertIsNone(fsdp_plan) return model def fake_load_pretrained_model(model, state_dict, checkpoint_files, load_config, expected_keys=None): From 9472436dcd448ff15fd712afcf3ffc4f0fe11242 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Mon, 18 May 2026 05:24:53 +0000 Subject: [PATCH 103/116] fsdp linting --- src/transformers/models/afmoe/modeling_afmoe.py | 2 +- .../models/apertus/configuration_apertus.py | 2 -- .../models/apertus/modeling_apertus.py | 2 +- .../models/apertus/modular_apertus.py | 6 ++++++ .../models/arcee/configuration_arcee.py | 2 -- src/transformers/models/arcee/modeling_arcee.py | 2 +- .../models/aria/configuration_aria.py | 2 -- src/transformers/models/aria/modeling_aria.py | 2 +- src/transformers/models/bamba/modeling_bamba.py | 2 +- src/transformers/models/bitnet/modeling_bitnet.py | 1 + .../models/cohere/configuration_cohere.py | 2 -- src/transformers/models/cohere/modeling_cohere.py | 2 +- .../models/cohere2/configuration_cohere2.py | 2 -- .../models/cohere2/modeling_cohere2.py | 2 +- .../models/cohere2/modular_cohere2.py | 6 ++++++ src/transformers/models/csm/modeling_csm.py | 1 + src/transformers/models/cwm/configuration_cwm.py | 2 -- src/transformers/models/cwm/modeling_cwm.py | 2 +- .../deepseek_v2/configuration_deepseek_v2.py | 2 -- .../models/deepseek_v2/modeling_deepseek_v2.py | 2 +- .../deepseek_v3/configuration_deepseek_v3.py | 2 -- .../models/deepseek_v3/modeling_deepseek_v3.py | 2 +- .../deepseek_v4/configuration_deepseek_v4.py | 2 -- .../models/diffllama/modeling_diffllama.py | 2 +- .../models/doge/configuration_doge.py | 2 -- src/transformers/models/doge/modular_doge.py | 6 ++++++ .../models/dots1/configuration_dots1.py | 3 +-- src/transformers/models/dots1/modular_dots1.py | 7 +++++++ src/transformers/models/emu3/modeling_emu3.py | 2 +- .../models/ernie4_5/configuration_ernie4_5.py | 2 -- .../models/ernie4_5/modeling_ernie4_5.py | 2 +- .../ernie4_5_moe/configuration_ernie4_5_moe.py | 2 -- .../configuration_ernie4_5_vl_moe.py | 2 -- .../models/eurobert/configuration_eurobert.py | 2 -- .../models/exaone4/configuration_exaone4.py | 2 -- .../models/exaone4/modeling_exaone4.py | 2 +- .../models/exaone4/modular_exaone4.py | 6 ++++++ .../models/exaone_moe/configuration_exaone_moe.py | 2 -- .../models/exaone_moe/modeling_exaone_moe.py | 2 +- .../models/exaone_moe/modular_exaone_moe.py | 6 ++++++ .../models/falcon_h1/modeling_falcon_h1.py | 2 +- .../models/flex_olmo/configuration_flex_olmo.py | 2 -- .../models/flex_olmo/modular_flex_olmo.py | 6 ++++++ .../models/gemma/configuration_gemma.py | 2 -- src/transformers/models/gemma/modeling_gemma.py | 2 +- src/transformers/models/gemma/modular_gemma.py | 6 ++++++ .../models/gemma2/configuration_gemma2.py | 2 -- src/transformers/models/gemma2/modeling_gemma2.py | 2 +- src/transformers/models/gemma2/modular_gemma2.py | 6 ++++++ .../models/gemma3/configuration_gemma3.py | 3 +-- src/transformers/models/gemma3/modeling_gemma3.py | 2 +- src/transformers/models/gemma3/modular_gemma3.py | 7 +++++++ .../models/gemma3n/configuration_gemma3n.py | 3 +-- .../models/gemma3n/modeling_gemma3n.py | 2 +- .../models/gemma3n/modular_gemma3n.py | 7 +++++++ .../models/gemma4/configuration_gemma4.py | 2 -- src/transformers/models/gemma4/modeling_gemma4.py | 2 +- src/transformers/models/glm/configuration_glm.py | 2 -- src/transformers/models/glm/modeling_glm.py | 2 +- .../models/glm4/configuration_glm4.py | 2 -- src/transformers/models/glm4/modeling_glm4.py | 2 +- .../models/glm4_moe/configuration_glm4_moe.py | 3 +-- .../models/glm4_moe/modeling_glm4_moe.py | 2 +- .../models/glm4_moe/modular_glm4_moe.py | 7 +++++++ .../glm4_moe_lite/configuration_glm4_moe_lite.py | 3 +-- .../glm4_moe_lite/modeling_glm4_moe_lite.py | 2 +- .../models/glm4_moe_lite/modular_glm4_moe_lite.py | 7 +++++++ .../models/glm4v/configuration_glm4v.py | 3 +-- src/transformers/models/glm4v/modular_glm4v.py | 7 +++++++ .../models/glm4v_moe/configuration_glm4v_moe.py | 4 ++-- .../models/glm4v_moe/modular_glm4v_moe.py | 7 +++++++ .../models/glm_image/configuration_glm_image.py | 3 +-- .../models/glm_image/modular_glm_image.py | 6 ++++++ .../glm_moe_dsa/configuration_glm_moe_dsa.py | 2 -- .../models/glm_moe_dsa/modeling_glm_moe_dsa.py | 2 +- .../models/glm_moe_dsa/modular_glm_moe_dsa.py | 6 ++++++ .../models/glm_ocr/configuration_glm_ocr.py | 3 +-- .../models/glm_ocr/modular_glm_ocr.py | 6 ++++++ .../models/gpt_neox/configuration_gpt_neox.py | 2 -- .../models/granite/configuration_granite.py | 2 -- .../models/granite/modeling_granite.py | 2 +- .../configuration_granite4_vision.py | 2 -- .../models/helium/configuration_helium.py | 2 -- src/transformers/models/helium/modeling_helium.py | 2 +- .../configuration_higgs_audio_v2.py | 2 -- .../hunyuan_v1_dense/modeling_hunyuan_v1_dense.py | 2 +- .../hunyuan_v1_moe/modeling_hunyuan_v1_moe.py | 2 +- .../models/hy_v3/configuration_hy_v3.py | 2 -- src/transformers/models/hy_v3/modeling_hy_v3.py | 2 +- src/transformers/models/hy_v3/modular_hy_v3.py | 6 ++++++ .../hyperclovax/configuration_hyperclovax.py | 2 -- .../models/hyperclovax/modeling_hyperclovax.py | 2 +- .../models/jais2/configuration_jais2.py | 2 -- src/transformers/models/jais2/modeling_jais2.py | 2 +- .../modeling_kyutai_speech_to_text.py | 2 +- .../models/laguna/configuration_laguna.py | 2 -- src/transformers/models/lfm2/modeling_lfm2.py | 2 +- .../models/lfm2_moe/modeling_lfm2_moe.py | 2 +- .../models/llama/configuration_llama.py | 2 -- .../longcat_flash/configuration_longcat_flash.py | 2 -- .../longcat_flash/modeling_longcat_flash.py | 2 +- .../models/minimax/configuration_minimax.py | 3 +-- .../models/minimax/modular_minimax.py | 7 +++++++ .../models/minimax_m2/configuration_minimax_m2.py | 3 +-- .../models/minimax_m2/modular_minimax_m2.py | 7 +++++++ .../models/ministral/configuration_ministral.py | 2 -- .../models/ministral/modeling_ministral.py | 2 +- .../models/ministral3/configuration_ministral3.py | 2 -- .../models/ministral3/modeling_ministral3.py | 2 +- .../models/mistral/configuration_mistral.py | 2 -- .../models/mistral/modeling_mistral.py | 2 +- .../models/mistral4/configuration_mistral4.py | 2 -- .../models/mistral4/modeling_mistral4.py | 2 +- .../models/nanochat/modeling_nanochat.py | 2 +- .../models/olmo/configuration_olmo.py | 2 -- src/transformers/models/olmo/modeling_olmo.py | 2 +- .../models/olmo2/configuration_olmo2.py | 2 -- src/transformers/models/olmo2/modeling_olmo2.py | 2 +- .../models/olmo3/configuration_olmo3.py | 2 -- src/transformers/models/olmo3/modeling_olmo3.py | 2 +- .../olmo_hybrid/configuration_olmo_hybrid.py | 2 -- .../models/olmo_hybrid/modeling_olmo_hybrid.py | 2 +- .../models/olmoe/configuration_olmoe.py | 2 -- .../configuration_openai_privacy_filter.py | 7 +++++++ .../modular_openai_privacy_filter.py | 7 +++++++ .../paddleocr_vl/configuration_paddleocr_vl.py | 2 -- src/transformers/models/phi/configuration_phi.py | 2 -- src/transformers/models/phi/modeling_phi.py | 2 +- .../models/phi3/configuration_phi3.py | 2 -- src/transformers/models/phi3/modeling_phi3.py | 2 +- .../configuration_phi4_multimodal.py | 2 -- .../phi4_multimodal/modeling_phi4_multimodal.py | 2 +- .../models/qwen2/configuration_qwen2.py | 2 -- src/transformers/models/qwen2/modeling_qwen2.py | 2 +- .../qwen2_5_omni/configuration_qwen2_5_omni.py | 3 +-- .../models/qwen2_5_omni/modular_qwen2_5_omni.py | 7 +++++++ .../models/qwen2_5_vl/configuration_qwen2_5_vl.py | 2 -- .../models/qwen2_moe/configuration_qwen2_moe.py | 2 -- .../models/qwen2_vl/configuration_qwen2_vl.py | 2 -- .../models/qwen3_5/configuration_qwen3_5.py | 2 -- .../qwen3_5_moe/configuration_qwen3_5_moe.py | 2 -- .../models/qwen3_5_moe/modeling_qwen3_5_moe.py | 1 + .../models/qwen3_moe/configuration_qwen3_moe.py | 2 -- .../models/qwen3_next/configuration_qwen3_next.py | 2 -- .../configuration_qwen3_omni_moe.py | 13 +++++++++++++ .../qwen3_omni_moe/modular_qwen3_omni_moe.py | 7 +++++++ .../qwen3_vl_moe/configuration_qwen3_vl_moe.py | 2 -- .../models/seed_oss/configuration_seed_oss.py | 2 -- .../models/seed_oss/modeling_seed_oss.py | 2 +- .../models/smollm3/configuration_smollm3.py | 2 -- .../models/smollm3/modeling_smollm3.py | 2 +- .../models/smollm3/modular_smollm3.py | 6 ++++++ .../models/solar_open/configuration_solar_open.py | 15 +++++++-------- .../models/solar_open/modeling_solar_open.py | 2 +- .../models/solar_open/modular_solar_open.py | 6 ++++++ .../models/starcoder2/configuration_starcoder2.py | 2 -- .../models/starcoder2/modeling_starcoder2.py | 2 +- .../models/t5gemma/configuration_t5gemma.py | 2 -- .../models/t5gemma/modular_t5gemma.py | 6 ++++++ .../models/t5gemma2/configuration_t5gemma2.py | 10 ++++++++-- .../models/t5gemma2/modular_t5gemma2.py | 7 +++++++ .../models/vaultgemma/configuration_vaultgemma.py | 2 -- .../models/vaultgemma/modeling_vaultgemma.py | 2 +- .../models/vaultgemma/modular_vaultgemma.py | 6 ++++++ .../configuration_voxtral_realtime.py | 2 -- .../voxtral_realtime/modeling_voxtral_realtime.py | 2 +- .../models/youtu/configuration_youtu.py | 2 -- src/transformers/models/youtu/modeling_youtu.py | 2 +- 168 files changed, 295 insertions(+), 217 deletions(-) diff --git a/src/transformers/models/afmoe/modeling_afmoe.py b/src/transformers/models/afmoe/modeling_afmoe.py index 138df483aa15..1fb30e564e1c 100644 --- a/src/transformers/models/afmoe/modeling_afmoe.py +++ b/src/transformers/models/afmoe/modeling_afmoe.py @@ -616,9 +616,9 @@ def forward( class AfmoeForCausalLM(AfmoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/apertus/configuration_apertus.py b/src/transformers/models/apertus/configuration_apertus.py index fb707e203dc6..c7fcc4159c06 100644 --- a/src/transformers/models/apertus/configuration_apertus.py +++ b/src/transformers/models/apertus/configuration_apertus.py @@ -76,8 +76,6 @@ class ApertusConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/apertus/modeling_apertus.py b/src/transformers/models/apertus/modeling_apertus.py index 9ae78dfc3459..8c5757377b03 100644 --- a/src/transformers/models/apertus/modeling_apertus.py +++ b/src/transformers/models/apertus/modeling_apertus.py @@ -423,9 +423,9 @@ def forward( class ApertusForCausalLM(ApertusPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 1850fc0cb43e..075a6de4776b 100644 --- a/src/transformers/models/apertus/modular_apertus.py +++ b/src/transformers/models/apertus/modular_apertus.py @@ -94,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 d1773d528676..e1d957bcac43 100644 --- a/src/transformers/models/arcee/configuration_arcee.py +++ b/src/transformers/models/arcee/configuration_arcee.py @@ -73,8 +73,6 @@ class ArceeConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/arcee/modeling_arcee.py b/src/transformers/models/arcee/modeling_arcee.py index c0aea6228193..05ebab40a879 100644 --- a/src/transformers/models/arcee/modeling_arcee.py +++ b/src/transformers/models/arcee/modeling_arcee.py @@ -425,9 +425,9 @@ def forward( class ArceeForCausalLM(ArceePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/aria/configuration_aria.py b/src/transformers/models/aria/configuration_aria.py index 986574d47c73..dfc1b8e74258 100644 --- a/src/transformers/models/aria/configuration_aria.py +++ b/src/transformers/models/aria/configuration_aria.py @@ -70,8 +70,6 @@ class AriaTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/aria/modeling_aria.py b/src/transformers/models/aria/modeling_aria.py index 69f67e5fff9c..5e90529f6195 100644 --- a/src/transformers/models/aria/modeling_aria.py +++ b/src/transformers/models/aria/modeling_aria.py @@ -759,9 +759,9 @@ def forward( class AriaTextForCausalLM(AriaTextPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config: AriaTextConfig): super().__init__(config) diff --git a/src/transformers/models/bamba/modeling_bamba.py b/src/transformers/models/bamba/modeling_bamba.py index 15bcef1ad696..c3b3a8a24bf1 100644 --- a/src/transformers/models/bamba/modeling_bamba.py +++ b/src/transformers/models/bamba/modeling_bamba.py @@ -1072,9 +1072,9 @@ def _update_mamba_mask(self, attention_mask, past_key_values): class BambaForCausalLM(BambaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 27c908d204d1..4994407d74c1 100644 --- a/src/transformers/models/bitnet/modeling_bitnet.py +++ b/src/transformers/models/bitnet/modeling_bitnet.py @@ -425,6 +425,7 @@ class BitNetForCausalLM(BitNetPreTrainedModel, GenerationMixin): _tp_plan = None _sp_plan = None _pp_plan = None + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/cohere/configuration_cohere.py b/src/transformers/models/cohere/configuration_cohere.py index da2139b85c60..4c05dce75e47 100644 --- a/src/transformers/models/cohere/configuration_cohere.py +++ b/src/transformers/models/cohere/configuration_cohere.py @@ -81,8 +81,6 @@ class CohereConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/cohere/modeling_cohere.py b/src/transformers/models/cohere/modeling_cohere.py index d543e709e8f1..309ddc142dc3 100644 --- a/src/transformers/models/cohere/modeling_cohere.py +++ b/src/transformers/models/cohere/modeling_cohere.py @@ -455,9 +455,9 @@ def forward( class CohereForCausalLM(CoherePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 548938b379e5..19197e26204d 100644 --- a/src/transformers/models/cohere2/configuration_cohere2.py +++ b/src/transformers/models/cohere2/configuration_cohere2.py @@ -78,8 +78,6 @@ class Cohere2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/cohere2/modeling_cohere2.py b/src/transformers/models/cohere2/modeling_cohere2.py index 24b2be0ca578..6a584c687c79 100644 --- a/src/transformers/models/cohere2/modeling_cohere2.py +++ b/src/transformers/models/cohere2/modeling_cohere2.py @@ -434,9 +434,9 @@ def forward( class Cohere2ForCausalLM(Cohere2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 574004647d0d..2720840a931e 100644 --- a/src/transformers/models/cohere2/modular_cohere2.py +++ b/src/transformers/models/cohere2/modular_cohere2.py @@ -99,6 +99,12 @@ class Cohere2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 256000 hidden_size: int = 8192 intermediate_size: int = 22528 diff --git a/src/transformers/models/csm/modeling_csm.py b/src/transformers/models/csm/modeling_csm.py index f2988078a40b..99f643e93b44 100644 --- a/src/transformers/models/csm/modeling_csm.py +++ b/src/transformers/models/csm/modeling_csm.py @@ -551,6 +551,7 @@ class CsmDepthDecoderForCausalLM(CsmPreTrainedModel, GenerationMixin): _tp_plan = None _sp_plan = None _pp_plan = None + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/cwm/configuration_cwm.py b/src/transformers/models/cwm/configuration_cwm.py index a5db8611589e..6eb7175ad020 100644 --- a/src/transformers/models/cwm/configuration_cwm.py +++ b/src/transformers/models/cwm/configuration_cwm.py @@ -75,8 +75,6 @@ class CwmConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/cwm/modeling_cwm.py b/src/transformers/models/cwm/modeling_cwm.py index 1abeff6707a5..b20fd9b7436f 100644 --- a/src/transformers/models/cwm/modeling_cwm.py +++ b/src/transformers/models/cwm/modeling_cwm.py @@ -427,9 +427,9 @@ def forward( class CwmForCausalLM(CwmPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py b/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py index e874ad881f96..49b3cb9c6387 100644 --- a/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py +++ b/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py @@ -94,8 +94,6 @@ class DeepseekV2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/deepseek_v2/modeling_deepseek_v2.py b/src/transformers/models/deepseek_v2/modeling_deepseek_v2.py index bc45461bbb90..1f9ba33ec6a5 100644 --- a/src/transformers/models/deepseek_v2/modeling_deepseek_v2.py +++ b/src/transformers/models/deepseek_v2/modeling_deepseek_v2.py @@ -542,9 +542,9 @@ def forward( class DeepseekV2ForCausalLM(DeepseekV2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py index be557639b37f..32e67ab4d90a 100644 --- a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py @@ -63,8 +63,6 @@ class DeepseekV3Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index d3ee5ae3efac..8d828320b19b 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -635,9 +635,9 @@ def forward( class DeepseekV3ForCausalLM(DeepseekV3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/deepseek_v4/configuration_deepseek_v4.py b/src/transformers/models/deepseek_v4/configuration_deepseek_v4.py index 02cd4feed68d..9475e8d6f1d0 100644 --- a/src/transformers/models/deepseek_v4/configuration_deepseek_v4.py +++ b/src/transformers/models/deepseek_v4/configuration_deepseek_v4.py @@ -128,8 +128,6 @@ class DeepseekV4Config(PreTrainedConfig): "layers.*.mlp.experts": "moe_experts_allreduce", } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/diffllama/modeling_diffllama.py b/src/transformers/models/diffllama/modeling_diffllama.py index fe8130e84116..b822da50e2af 100644 --- a/src/transformers/models/diffllama/modeling_diffllama.py +++ b/src/transformers/models/diffllama/modeling_diffllama.py @@ -661,9 +661,9 @@ def forward( class DiffLlamaForCausalLM(DiffLlamaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/doge/configuration_doge.py b/src/transformers/models/doge/configuration_doge.py index 375a1c8738f7..feac227b5a97 100644 --- a/src/transformers/models/doge/configuration_doge.py +++ b/src/transformers/models/doge/configuration_doge.py @@ -70,8 +70,6 @@ class DogeConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/doge/modular_doge.py b/src/transformers/models/doge/modular_doge.py index afe2469fc9b2..bc7e22089ea3 100644 --- a/src/transformers/models/doge/modular_doge.py +++ b/src/transformers/models/doge/modular_doge.py @@ -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 b998a8e0bc54..9846f961250f 100644 --- a/src/transformers/models/dots1/configuration_dots1.py +++ b/src/transformers/models/dots1/configuration_dots1.py @@ -67,13 +67,12 @@ class Dots1Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. 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/modular_dots1.py b/src/transformers/models/dots1/modular_dots1.py index cf52593c63b5..94bd074871c0 100644 --- a/src/transformers/models/dots1/modular_dots1.py +++ b/src/transformers/models/dots1/modular_dots1.py @@ -80,6 +80,13 @@ class Dots1Config(PreTrainedConfig): "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + attribute_map = { "num_local_experts": "n_routed_experts", } diff --git a/src/transformers/models/emu3/modeling_emu3.py b/src/transformers/models/emu3/modeling_emu3.py index ba27551a74f4..bd3f3011a539 100644 --- a/src/transformers/models/emu3/modeling_emu3.py +++ b/src/transformers/models/emu3/modeling_emu3.py @@ -1276,9 +1276,9 @@ def forward( class Emu3ForCausalLM(Emu3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 b1d0f8c47dd8..5d805e1114e6 100644 --- a/src/transformers/models/ernie4_5/configuration_ernie4_5.py +++ b/src/transformers/models/ernie4_5/configuration_ernie4_5.py @@ -76,8 +76,6 @@ class Ernie4_5Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/ernie4_5/modeling_ernie4_5.py b/src/transformers/models/ernie4_5/modeling_ernie4_5.py index c98444bd38ce..73f15f713be4 100644 --- a/src/transformers/models/ernie4_5/modeling_ernie4_5.py +++ b/src/transformers/models/ernie4_5/modeling_ernie4_5.py @@ -423,9 +423,9 @@ def forward( class Ernie4_5ForCausalLM(Ernie4_5PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 f0ea848bff2d..035daabd966d 100644 --- a/src/transformers/models/ernie4_5_moe/configuration_ernie4_5_moe.py +++ b/src/transformers/models/ernie4_5_moe/configuration_ernie4_5_moe.py @@ -81,8 +81,6 @@ class Ernie4_5_MoeConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", 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 3cf9dfbe063e..34caaa02dbeb 100644 --- a/src/transformers/models/ernie4_5_vl_moe/configuration_ernie4_5_vl_moe.py +++ b/src/transformers/models/ernie4_5_vl_moe/configuration_ernie4_5_vl_moe.py @@ -100,8 +100,6 @@ class Ernie4_5_VLMoeTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/eurobert/configuration_eurobert.py b/src/transformers/models/eurobert/configuration_eurobert.py index ab818c3ba80d..ab2c33a28828 100644 --- a/src/transformers/models/eurobert/configuration_eurobert.py +++ b/src/transformers/models/eurobert/configuration_eurobert.py @@ -80,8 +80,6 @@ class EuroBertConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/exaone4/configuration_exaone4.py b/src/transformers/models/exaone4/configuration_exaone4.py index 9aed5cc3cee3..fd3bade9f662 100644 --- a/src/transformers/models/exaone4/configuration_exaone4.py +++ b/src/transformers/models/exaone4/configuration_exaone4.py @@ -91,8 +91,6 @@ class Exaone4Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/exaone4/modeling_exaone4.py b/src/transformers/models/exaone4/modeling_exaone4.py index cbe34bc5fe6c..d6aa12baedad 100644 --- a/src/transformers/models/exaone4/modeling_exaone4.py +++ b/src/transformers/models/exaone4/modeling_exaone4.py @@ -441,9 +441,9 @@ def forward( class Exaone4ForCausalLM(Exaone4PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 ccbd0926290c..7a88b77629f1 100644 --- a/src/transformers/models/exaone4/modular_exaone4.py +++ b/src/transformers/models/exaone4/modular_exaone4.py @@ -120,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 2c00f099af95..83a873314700 100644 --- a/src/transformers/models/exaone_moe/configuration_exaone_moe.py +++ b/src/transformers/models/exaone_moe/configuration_exaone_moe.py @@ -82,8 +82,6 @@ class ExaoneMoeConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/exaone_moe/modeling_exaone_moe.py b/src/transformers/models/exaone_moe/modeling_exaone_moe.py index 78787c80cecd..d2d8543f1417 100644 --- a/src/transformers/models/exaone_moe/modeling_exaone_moe.py +++ b/src/transformers/models/exaone_moe/modeling_exaone_moe.py @@ -564,9 +564,9 @@ def forward( class ExaoneMoeForCausalLM(ExaoneMoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 88b6a60e0e16..8d50c3db2295 100644 --- a/src/transformers/models/exaone_moe/modular_exaone_moe.py +++ b/src/transformers/models/exaone_moe/modular_exaone_moe.py @@ -82,6 +82,12 @@ class ExaoneMoeConfig(Exaone4Config): base_model_sp_plan = None + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 102400 hidden_size: int = 4096 intermediate_size: int = 16384 diff --git a/src/transformers/models/falcon_h1/modeling_falcon_h1.py b/src/transformers/models/falcon_h1/modeling_falcon_h1.py index 4a9c63e25cda..2ea3c884437e 100644 --- a/src/transformers/models/falcon_h1/modeling_falcon_h1.py +++ b/src/transformers/models/falcon_h1/modeling_falcon_h1.py @@ -1167,9 +1167,9 @@ def _update_mamba_mask(self, attention_mask, past_key_values): class FalconH1ForCausalLM(FalconH1PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 6727c6859227..a89d9f0452e8 100644 --- a/src/transformers/models/flex_olmo/configuration_flex_olmo.py +++ b/src/transformers/models/flex_olmo/configuration_flex_olmo.py @@ -61,8 +61,6 @@ class FlexOlmoConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/flex_olmo/modular_flex_olmo.py b/src/transformers/models/flex_olmo/modular_flex_olmo.py index e3dbd850f3ce..c7b73dea8f15 100644 --- a/src/transformers/models/flex_olmo/modular_flex_olmo.py +++ b/src/transformers/models/flex_olmo/modular_flex_olmo.py @@ -71,6 +71,12 @@ class FlexOlmoConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 100352 hidden_size: int = 4096 intermediate_size: int = 11008 diff --git a/src/transformers/models/gemma/configuration_gemma.py b/src/transformers/models/gemma/configuration_gemma.py index 1884e0533c89..7b80a3fda63e 100644 --- a/src/transformers/models/gemma/configuration_gemma.py +++ b/src/transformers/models/gemma/configuration_gemma.py @@ -76,8 +76,6 @@ class GemmaConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/gemma/modeling_gemma.py b/src/transformers/models/gemma/modeling_gemma.py index 16035150b149..f64702b55f19 100644 --- a/src/transformers/models/gemma/modeling_gemma.py +++ b/src/transformers/models/gemma/modeling_gemma.py @@ -451,9 +451,9 @@ def forward( class GemmaForCausalLM(GemmaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 06168ff81cbb..8a24552eaa36 100644 --- a/src/transformers/models/gemma/modular_gemma.py +++ b/src/transformers/models/gemma/modular_gemma.py @@ -95,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 eecd939308e8..8ff4fc88629a 100644 --- a/src/transformers/models/gemma2/configuration_gemma2.py +++ b/src/transformers/models/gemma2/configuration_gemma2.py @@ -80,8 +80,6 @@ class Gemma2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/gemma2/modeling_gemma2.py b/src/transformers/models/gemma2/modeling_gemma2.py index da784cab3c74..46eb89681619 100644 --- a/src/transformers/models/gemma2/modeling_gemma2.py +++ b/src/transformers/models/gemma2/modeling_gemma2.py @@ -477,9 +477,9 @@ def forward( class Gemma2ForCausalLM(Gemma2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 6c52bff8d9dd..d400331af615 100644 --- a/src/transformers/models/gemma2/modular_gemma2.py +++ b/src/transformers/models/gemma2/modular_gemma2.py @@ -107,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 0c29116d4e03..adf71ed56bba 100644 --- a/src/transformers/models/gemma3/configuration_gemma3.py +++ b/src/transformers/models/gemma3/configuration_gemma3.py @@ -89,8 +89,6 @@ class Gemma3TextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", @@ -122,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 fdb9f31d2121..e49a73596dca 100644 --- a/src/transformers/models/gemma3/modeling_gemma3.py +++ b/src/transformers/models/gemma3/modeling_gemma3.py @@ -586,9 +586,9 @@ def forward( class Gemma3ForCausalLM(Gemma3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 1bb1e5bcf370..e83243f21c4a 100644 --- a/src/transformers/models/gemma3/modular_gemma3.py +++ b/src/transformers/models/gemma3/modular_gemma3.py @@ -109,6 +109,13 @@ class Gemma3TextConfig(Gemma2Config, PreTrainedConfig): "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 08a4747f8d1f..973ff753c28b 100644 --- a/src/transformers/models/gemma3n/configuration_gemma3n.py +++ b/src/transformers/models/gemma3n/configuration_gemma3n.py @@ -93,8 +93,6 @@ class Gemma3nTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", @@ -123,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 297dc14bbd15..b2b1a667b96a 100644 --- a/src/transformers/models/gemma3n/modeling_gemma3n.py +++ b/src/transformers/models/gemma3n/modeling_gemma3n.py @@ -1828,9 +1828,9 @@ def project_per_layer_inputs( class Gemma3nForCausalLM(Gemma3nPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 539a1b4f1259..2cbd72a61ed0 100644 --- a/src/transformers/models/gemma3n/modular_gemma3n.py +++ b/src/transformers/models/gemma3n/modular_gemma3n.py @@ -125,6 +125,13 @@ class Gemma3nTextConfig(Gemma3TextConfig): "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 f447f9b13519..0b4b24b00642 100644 --- a/src/transformers/models/gemma4/configuration_gemma4.py +++ b/src/transformers/models/gemma4/configuration_gemma4.py @@ -151,8 +151,6 @@ class Gemma4TextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/gemma4/modeling_gemma4.py b/src/transformers/models/gemma4/modeling_gemma4.py index 61ee14f0f6b2..07e07aa974e5 100644 --- a/src/transformers/models/gemma4/modeling_gemma4.py +++ b/src/transformers/models/gemma4/modeling_gemma4.py @@ -1794,9 +1794,9 @@ def project_per_layer_inputs( class Gemma4ForCausalLM(Gemma4PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} config: Gemma4TextConfig base_model_prefix = "model" diff --git a/src/transformers/models/glm/configuration_glm.py b/src/transformers/models/glm/configuration_glm.py index c8cdee74019a..67275679ab8b 100644 --- a/src/transformers/models/glm/configuration_glm.py +++ b/src/transformers/models/glm/configuration_glm.py @@ -67,8 +67,6 @@ class GlmConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/glm/modeling_glm.py b/src/transformers/models/glm/modeling_glm.py index 2fbb4e4fb64d..461703ab81cd 100644 --- a/src/transformers/models/glm/modeling_glm.py +++ b/src/transformers/models/glm/modeling_glm.py @@ -440,9 +440,9 @@ def forward( class GlmForCausalLM(GlmPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 4eb5904d9ce9..8b2755323f68 100644 --- a/src/transformers/models/glm4/configuration_glm4.py +++ b/src/transformers/models/glm4/configuration_glm4.py @@ -67,8 +67,6 @@ class Glm4Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/glm4/modeling_glm4.py b/src/transformers/models/glm4/modeling_glm4.py index baf0a0e634ef..aa3d6e62a124 100644 --- a/src/transformers/models/glm4/modeling_glm4.py +++ b/src/transformers/models/glm4/modeling_glm4.py @@ -445,9 +445,9 @@ def forward( class Glm4ForCausalLM(Glm4PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 7b7ea57267f0..bd94a492bb66 100644 --- a/src/transformers/models/glm4_moe/configuration_glm4_moe.py +++ b/src/transformers/models/glm4_moe/configuration_glm4_moe.py @@ -72,13 +72,12 @@ class Glm4MoeConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. 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 bab7e3a6084a..730bf70db8ef 100644 --- a/src/transformers/models/glm4_moe/modeling_glm4_moe.py +++ b/src/transformers/models/glm4_moe/modeling_glm4_moe.py @@ -578,9 +578,9 @@ def forward( class Glm4MoeForCausalLM(Glm4MoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 f1caa381ec3b..41921eb080e6 100644 --- a/src/transformers/models/glm4_moe/modular_glm4_moe.py +++ b/src/transformers/models/glm4_moe/modular_glm4_moe.py @@ -84,6 +84,13 @@ class Glm4MoeConfig(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/glm4_moe_lite/configuration_glm4_moe_lite.py b/src/transformers/models/glm4_moe_lite/configuration_glm4_moe_lite.py index 8fb2a0fbf3ec..cc0c1fc67ffe 100644 --- a/src/transformers/models/glm4_moe_lite/configuration_glm4_moe_lite.py +++ b/src/transformers/models/glm4_moe_lite/configuration_glm4_moe_lite.py @@ -66,13 +66,12 @@ class Glm4MoeLiteConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. 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 ba44c049339d..4bab2f0316ae 100644 --- a/src/transformers/models/glm4_moe_lite/modeling_glm4_moe_lite.py +++ b/src/transformers/models/glm4_moe_lite/modeling_glm4_moe_lite.py @@ -652,9 +652,9 @@ def forward( class Glm4MoeLiteForCausalLM(Glm4MoeLitePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 41bf6b0aff0d..ea8e5f60e491 100644 --- a/src/transformers/models/glm4_moe_lite/modular_glm4_moe_lite.py +++ b/src/transformers/models/glm4_moe_lite/modular_glm4_moe_lite.py @@ -73,6 +73,13 @@ class Glm4MoeLiteConfig(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", "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 e7e6bb9927f2..dc993ab631d6 100644 --- a/src/transformers/models/glm4v/configuration_glm4v.py +++ b/src/transformers/models/glm4v/configuration_glm4v.py @@ -103,13 +103,12 @@ class Glm4vTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. 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 85e4dadde5e1..880a772a537c 100644 --- a/src/transformers/models/glm4v/modular_glm4v.py +++ b/src/transformers/models/glm4v/modular_glm4v.py @@ -147,6 +147,13 @@ class Glm4vTextConfig(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", + } + 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 be09d149e7bf..73971460b2d4 100644 --- a/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py +++ b/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py @@ -67,13 +67,12 @@ class Glm4vMoeTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. 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", } @@ -106,6 +105,7 @@ class Glm4vMoeTextConfig(PreTrainedConfig): eos_token_id: int | list[int] | None = None pad_token_id: int | None = None base_config_key = "text_config" + ignore_keys_at_rope_validation = {"mrope_section"} router_aux_loss_coef: float = 0.0001 diff --git a/src/transformers/models/glm4v_moe/modular_glm4v_moe.py b/src/transformers/models/glm4v_moe/modular_glm4v_moe.py index 4f1d62eb84b0..289dac43e160 100644 --- a/src/transformers/models/glm4v_moe/modular_glm4v_moe.py +++ b/src/transformers/models/glm4v_moe/modular_glm4v_moe.py @@ -98,6 +98,13 @@ class Glm4vMoeTextConfig(Glm4MoeConfig): "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 8c3c985ff9f0..e022da23e861 100644 --- a/src/transformers/models/glm_image/configuration_glm_image.py +++ b/src/transformers/models/glm_image/configuration_glm_image.py @@ -116,13 +116,12 @@ class GlmImageTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", "norm": "keep_full_weight", } + ignore_keys_at_rope_validation = {"mrope_section"} vocab_size: int = 168064 diff --git a/src/transformers/models/glm_image/modular_glm_image.py b/src/transformers/models/glm_image/modular_glm_image.py index 273861a5b61e..8ce758601acd 100644 --- a/src/transformers/models/glm_image/modular_glm_image.py +++ b/src/transformers/models/glm_image/modular_glm_image.py @@ -133,6 +133,12 @@ class GlmImageTextConfig(Glm4vTextConfig): >>> configuration = model.config ```""" + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 168064 max_position_embeddings: int = 131072 vision_vocab_size: int = 16512 diff --git a/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py index 0726f645bb77..31512bde976a 100644 --- a/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py @@ -76,8 +76,6 @@ class GlmMoeDsaConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", 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 9eb5b94a9cba..1098096bd4e3 100644 --- a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py @@ -822,9 +822,9 @@ def forward( class GlmMoeDsaForCausalLM(GlmMoeDsaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 0a0d284403ff..d45602619458 100644 --- a/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py @@ -119,6 +119,12 @@ class GlmMoeDsaConfig(Glm4MoeLiteConfig): "layers.*.mlp.down_proj": "rowwise_allreduce", } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + hidden_size: int = 6144 intermediate_size: int = 12288 moe_intermediate_size: int = 2048 diff --git a/src/transformers/models/glm_ocr/configuration_glm_ocr.py b/src/transformers/models/glm_ocr/configuration_glm_ocr.py index 1b2627926de7..e1381b8bc79d 100644 --- a/src/transformers/models/glm_ocr/configuration_glm_ocr.py +++ b/src/transformers/models/glm_ocr/configuration_glm_ocr.py @@ -104,13 +104,12 @@ class GlmOcrTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", "norm": "keep_full_weight", } + ignore_keys_at_rope_validation = {"mrope_section"} vocab_size: int = 59392 diff --git a/src/transformers/models/glm_ocr/modular_glm_ocr.py b/src/transformers/models/glm_ocr/modular_glm_ocr.py index 782808bb69a9..2e9fc1c25a26 100644 --- a/src/transformers/models/glm_ocr/modular_glm_ocr.py +++ b/src/transformers/models/glm_ocr/modular_glm_ocr.py @@ -82,6 +82,12 @@ class GlmOcrTextConfig(Glm4vTextConfig): >>> configuration = model.config ```""" + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 59392 hidden_size: int = 1024 intermediate_size: int = 4096 diff --git a/src/transformers/models/gpt_neox/configuration_gpt_neox.py b/src/transformers/models/gpt_neox/configuration_gpt_neox.py index c62775b283bb..abd0b4578a72 100644 --- a/src/transformers/models/gpt_neox/configuration_gpt_neox.py +++ b/src/transformers/models/gpt_neox/configuration_gpt_neox.py @@ -67,8 +67,6 @@ class GPTNeoXConfig(PreTrainedConfig): "final_layer_norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/granite/configuration_granite.py b/src/transformers/models/granite/configuration_granite.py index 11489cd76904..1464157b962d 100644 --- a/src/transformers/models/granite/configuration_granite.py +++ b/src/transformers/models/granite/configuration_granite.py @@ -76,8 +76,6 @@ class GraniteConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/granite/modeling_granite.py b/src/transformers/models/granite/modeling_granite.py index 2b6066eb8a63..5a266b3acc93 100644 --- a/src/transformers/models/granite/modeling_granite.py +++ b/src/transformers/models/granite/modeling_granite.py @@ -446,9 +446,9 @@ def forward( class GraniteForCausalLM(GranitePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 5986d4f512b6..d3fb9129e5d2 100644 --- a/src/transformers/models/granite4_vision/configuration_granite4_vision.py +++ b/src/transformers/models/granite4_vision/configuration_granite4_vision.py @@ -79,8 +79,6 @@ class Granite4VisionTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/helium/configuration_helium.py b/src/transformers/models/helium/configuration_helium.py index 39b39dce0562..5aca27c00069 100644 --- a/src/transformers/models/helium/configuration_helium.py +++ b/src/transformers/models/helium/configuration_helium.py @@ -70,8 +70,6 @@ class HeliumConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/helium/modeling_helium.py b/src/transformers/models/helium/modeling_helium.py index 4668475fdb84..555fdc7592d3 100644 --- a/src/transformers/models/helium/modeling_helium.py +++ b/src/transformers/models/helium/modeling_helium.py @@ -424,9 +424,9 @@ def forward( class HeliumForCausalLM(HeliumPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 4c36b638b32d..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 @@ -88,8 +88,6 @@ class HiggsAudioV2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", 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 aaeac78158b0..3b2386cf8ab3 100644 --- a/src/transformers/models/hunyuan_v1_dense/modeling_hunyuan_v1_dense.py +++ b/src/transformers/models/hunyuan_v1_dense/modeling_hunyuan_v1_dense.py @@ -462,9 +462,9 @@ def forward( class HunYuanDenseV1ForCausalLM(HunYuanDenseV1PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 0bf0dcf96b11..db33809260af 100644 --- a/src/transformers/models/hunyuan_v1_moe/modeling_hunyuan_v1_moe.py +++ b/src/transformers/models/hunyuan_v1_moe/modeling_hunyuan_v1_moe.py @@ -551,9 +551,9 @@ def forward( class HunYuanMoEV1ForCausalLM(HunYuanMoEV1PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 5d20419b9d0c..dddb7269f95e 100644 --- a/src/transformers/models/hy_v3/configuration_hy_v3.py +++ b/src/transformers/models/hy_v3/configuration_hy_v3.py @@ -73,8 +73,6 @@ class HYV3Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/hy_v3/modeling_hy_v3.py b/src/transformers/models/hy_v3/modeling_hy_v3.py index d1970f4fe8ab..6ade5c2d278d 100644 --- a/src/transformers/models/hy_v3/modeling_hy_v3.py +++ b/src/transformers/models/hy_v3/modeling_hy_v3.py @@ -545,9 +545,9 @@ def forward( class HYV3ForCausalLM(HYV3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 0f63e5b32f59..5c9547a2fe0e 100644 --- a/src/transformers/models/hy_v3/modular_hy_v3.py +++ b/src/transformers/models/hy_v3/modular_hy_v3.py @@ -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 6dcbcb7576d8..1090a4e00fda 100644 --- a/src/transformers/models/hyperclovax/configuration_hyperclovax.py +++ b/src/transformers/models/hyperclovax/configuration_hyperclovax.py @@ -91,8 +91,6 @@ class HyperCLOVAXConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/hyperclovax/modeling_hyperclovax.py b/src/transformers/models/hyperclovax/modeling_hyperclovax.py index d228f1cae263..a69fb46b121a 100644 --- a/src/transformers/models/hyperclovax/modeling_hyperclovax.py +++ b/src/transformers/models/hyperclovax/modeling_hyperclovax.py @@ -453,9 +453,9 @@ def forward( class HyperCLOVAXForCausalLM(HyperCLOVAXPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 b03692a2342d..6fa9783effd3 100644 --- a/src/transformers/models/jais2/configuration_jais2.py +++ b/src/transformers/models/jais2/configuration_jais2.py @@ -74,8 +74,6 @@ class Jais2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/jais2/modeling_jais2.py b/src/transformers/models/jais2/modeling_jais2.py index d7d4ce972b01..1fb6c19d9044 100644 --- a/src/transformers/models/jais2/modeling_jais2.py +++ b/src/transformers/models/jais2/modeling_jais2.py @@ -398,9 +398,9 @@ def forward( class Jais2ForCausalLM(Jais2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): 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 87460c65e115..fbe51b44259d 100644 --- a/src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py +++ b/src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py @@ -875,9 +875,9 @@ def forward( class KyutaiSpeechToTextForConditionalGeneration(KyutaiSpeechToTextPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 2843fd449da5..50e011ef08be 100644 --- a/src/transformers/models/laguna/configuration_laguna.py +++ b/src/transformers/models/laguna/configuration_laguna.py @@ -80,8 +80,6 @@ class LagunaConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/lfm2/modeling_lfm2.py b/src/transformers/models/lfm2/modeling_lfm2.py index 227afd2a35a5..95af77e0141b 100644 --- a/src/transformers/models/lfm2/modeling_lfm2.py +++ b/src/transformers/models/lfm2/modeling_lfm2.py @@ -541,9 +541,9 @@ def forward( class Lfm2ForCausalLM(Lfm2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 6e62546d4ee3..dc97dfb11454 100644 --- a/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py +++ b/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py @@ -631,9 +631,9 @@ def forward( class Lfm2MoeForCausalLM(Lfm2MoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 c6698cb26cc5..0f3a0e75f7f6 100644 --- a/src/transformers/models/llama/configuration_llama.py +++ b/src/transformers/models/llama/configuration_llama.py @@ -76,8 +76,6 @@ class LlamaConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/longcat_flash/configuration_longcat_flash.py b/src/transformers/models/longcat_flash/configuration_longcat_flash.py index e083bf32466a..ffeb6de9d01f 100644 --- a/src/transformers/models/longcat_flash/configuration_longcat_flash.py +++ b/src/transformers/models/longcat_flash/configuration_longcat_flash.py @@ -68,8 +68,6 @@ class LongcatFlashConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/longcat_flash/modeling_longcat_flash.py b/src/transformers/models/longcat_flash/modeling_longcat_flash.py index b5e6f889bfe5..4e7b3ad287cd 100644 --- a/src/transformers/models/longcat_flash/modeling_longcat_flash.py +++ b/src/transformers/models/longcat_flash/modeling_longcat_flash.py @@ -651,9 +651,9 @@ def forward( class LongcatFlashForCausalLM(LongcatFlashPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} _keys_to_ignore_on_load_unexpected = [r"model\.mtp.*"] def __init__(self, config): diff --git a/src/transformers/models/minimax/configuration_minimax.py b/src/transformers/models/minimax/configuration_minimax.py index 249a9ae2a3e0..d60e6d0c2849 100644 --- a/src/transformers/models/minimax/configuration_minimax.py +++ b/src/transformers/models/minimax/configuration_minimax.py @@ -74,13 +74,12 @@ class MiniMaxConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. 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/modular_minimax.py b/src/transformers/models/minimax/modular_minimax.py index 485417e274dd..8442eebd53e7 100644 --- a/src/transformers/models/minimax/modular_minimax.py +++ b/src/transformers/models/minimax/modular_minimax.py @@ -99,6 +99,13 @@ class MiniMaxConfig(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_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 1c14670d722a..6193523d95b3 100644 --- a/src/transformers/models/minimax_m2/configuration_minimax_m2.py +++ b/src/transformers/models/minimax_m2/configuration_minimax_m2.py @@ -73,13 +73,12 @@ class MiniMaxM2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. 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/modular_minimax_m2.py b/src/transformers/models/minimax_m2/modular_minimax_m2.py index 127638fb38ab..ca1b29125ccf 100644 --- a/src/transformers/models/minimax_m2/modular_minimax_m2.py +++ b/src/transformers/models/minimax_m2/modular_minimax_m2.py @@ -91,6 +91,13 @@ class MiniMaxM2Config(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_experts": "num_local_experts", } diff --git a/src/transformers/models/ministral/configuration_ministral.py b/src/transformers/models/ministral/configuration_ministral.py index 8f6d17ca7a17..fa0de0f0671e 100644 --- a/src/transformers/models/ministral/configuration_ministral.py +++ b/src/transformers/models/ministral/configuration_ministral.py @@ -78,8 +78,6 @@ class MinistralConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/ministral/modeling_ministral.py b/src/transformers/models/ministral/modeling_ministral.py index f9ae6841c06b..d2faea85b8cc 100644 --- a/src/transformers/models/ministral/modeling_ministral.py +++ b/src/transformers/models/ministral/modeling_ministral.py @@ -430,9 +430,9 @@ def forward( class MinistralForCausalLM(MinistralPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 4afb236b1504..d486d2efe796 100644 --- a/src/transformers/models/ministral3/configuration_ministral3.py +++ b/src/transformers/models/ministral3/configuration_ministral3.py @@ -84,8 +84,6 @@ class Ministral3Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/ministral3/modeling_ministral3.py b/src/transformers/models/ministral3/modeling_ministral3.py index 0218a8b4c096..066a71997913 100644 --- a/src/transformers/models/ministral3/modeling_ministral3.py +++ b/src/transformers/models/ministral3/modeling_ministral3.py @@ -413,9 +413,9 @@ def forward( class Ministral3ForCausalLM(Ministral3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 fb4f4919ac8b..0451d2f958a3 100644 --- a/src/transformers/models/mistral/configuration_mistral.py +++ b/src/transformers/models/mistral/configuration_mistral.py @@ -75,8 +75,6 @@ class MistralConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/mistral/modeling_mistral.py b/src/transformers/models/mistral/modeling_mistral.py index ea1cbd5647a4..d94ff3d6a312 100644 --- a/src/transformers/models/mistral/modeling_mistral.py +++ b/src/transformers/models/mistral/modeling_mistral.py @@ -402,9 +402,9 @@ def forward( class MistralForCausalLM(MistralPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 ddb619cdf693..c8218ffa5c57 100644 --- a/src/transformers/models/mistral4/configuration_mistral4.py +++ b/src/transformers/models/mistral4/configuration_mistral4.py @@ -61,8 +61,6 @@ class Mistral4Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/mistral4/modeling_mistral4.py b/src/transformers/models/mistral4/modeling_mistral4.py index 1f81bd21912f..4473107418ee 100644 --- a/src/transformers/models/mistral4/modeling_mistral4.py +++ b/src/transformers/models/mistral4/modeling_mistral4.py @@ -641,9 +641,9 @@ def forward( class Mistral4ForCausalLM(Mistral4PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/nanochat/modeling_nanochat.py b/src/transformers/models/nanochat/modeling_nanochat.py index 285578dde755..a191eb2d5425 100644 --- a/src/transformers/models/nanochat/modeling_nanochat.py +++ b/src/transformers/models/nanochat/modeling_nanochat.py @@ -433,9 +433,9 @@ def forward( class NanoChatForCausalLM(NanoChatPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/olmo/configuration_olmo.py b/src/transformers/models/olmo/configuration_olmo.py index 32af115cfbe3..f9f32c1ed15d 100644 --- a/src/transformers/models/olmo/configuration_olmo.py +++ b/src/transformers/models/olmo/configuration_olmo.py @@ -79,8 +79,6 @@ class OlmoConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/olmo/modeling_olmo.py b/src/transformers/models/olmo/modeling_olmo.py index 943e4751aa9b..dad24db7ce06 100644 --- a/src/transformers/models/olmo/modeling_olmo.py +++ b/src/transformers/models/olmo/modeling_olmo.py @@ -426,9 +426,9 @@ def forward( class OlmoForCausalLM(OlmoPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 8da1d8c8fee8..719db72c42e1 100644 --- a/src/transformers/models/olmo2/configuration_olmo2.py +++ b/src/transformers/models/olmo2/configuration_olmo2.py @@ -84,8 +84,6 @@ class Olmo2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/olmo2/modeling_olmo2.py b/src/transformers/models/olmo2/modeling_olmo2.py index 77d40756513c..12d31e551870 100644 --- a/src/transformers/models/olmo2/modeling_olmo2.py +++ b/src/transformers/models/olmo2/modeling_olmo2.py @@ -430,9 +430,9 @@ def forward( class Olmo2ForCausalLM(Olmo2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/olmo3/configuration_olmo3.py b/src/transformers/models/olmo3/configuration_olmo3.py index 6d34f887875a..cf81a4b3e332 100644 --- a/src/transformers/models/olmo3/configuration_olmo3.py +++ b/src/transformers/models/olmo3/configuration_olmo3.py @@ -79,8 +79,6 @@ class Olmo3Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/olmo3/modeling_olmo3.py b/src/transformers/models/olmo3/modeling_olmo3.py index 5a1d57a7ab2e..8325b651312a 100644 --- a/src/transformers/models/olmo3/modeling_olmo3.py +++ b/src/transformers/models/olmo3/modeling_olmo3.py @@ -434,9 +434,9 @@ def forward( class Olmo3ForCausalLM(Olmo3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py index 3d83f8548bb2..106fb2010bfc 100644 --- a/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py +++ b/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py @@ -88,8 +88,6 @@ class OlmoHybridConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py index 322dd780abf0..081149371067 100644 --- a/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py +++ b/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py @@ -1046,9 +1046,9 @@ def _update_linear_attn_mask(self, attention_mask, past_key_values): class OlmoHybridForCausalLM(OlmoHybridPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/olmoe/configuration_olmoe.py b/src/transformers/models/olmoe/configuration_olmoe.py index ef7077ae6867..feedb02bfda4 100644 --- a/src/transformers/models/olmoe/configuration_olmoe.py +++ b/src/transformers/models/olmoe/configuration_olmoe.py @@ -66,8 +66,6 @@ class OlmoeConfig(PreTrainedConfig): "norm": "activation", } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", 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 ae833c2bf514..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,6 +56,12 @@ 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", @@ -66,6 +72,7 @@ class OpenAIPrivacyFilterConfig(PreTrainedConfig): } 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 b7e4a34e370b..b9bef13987ac 100644 --- a/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py +++ b/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py @@ -124,8 +124,6 @@ class PaddleOCRTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/phi/configuration_phi.py b/src/transformers/models/phi/configuration_phi.py index 03326dad422a..bee22ad6a0b7 100644 --- a/src/transformers/models/phi/configuration_phi.py +++ b/src/transformers/models/phi/configuration_phi.py @@ -74,8 +74,6 @@ class PhiConfig(PreTrainedConfig): "final_layernorm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/phi/modeling_phi.py b/src/transformers/models/phi/modeling_phi.py index c628d40cf57f..0acf6a8060aa 100644 --- a/src/transformers/models/phi/modeling_phi.py +++ b/src/transformers/models/phi/modeling_phi.py @@ -407,9 +407,9 @@ def forward( class PhiForCausalLM(PhiPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 c49823521c66..d1288d411863 100644 --- a/src/transformers/models/phi3/configuration_phi3.py +++ b/src/transformers/models/phi3/configuration_phi3.py @@ -70,8 +70,6 @@ class Phi3Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/phi3/modeling_phi3.py b/src/transformers/models/phi3/modeling_phi3.py index 895426414d78..f0d6726b73d9 100644 --- a/src/transformers/models/phi3/modeling_phi3.py +++ b/src/transformers/models/phi3/modeling_phi3.py @@ -433,9 +433,9 @@ def forward( class Phi3ForCausalLM(Phi3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 f35ba25428c5..264119a82bec 100644 --- a/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py +++ b/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py @@ -204,8 +204,6 @@ class Phi4MultimodalConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py b/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py index a8c05cd314fa..383445a29f8b 100644 --- a/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py +++ b/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py @@ -1597,9 +1597,9 @@ def forward( class Phi4MultimodalForCausalLM(Phi4MultimodalPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/qwen2/configuration_qwen2.py b/src/transformers/models/qwen2/configuration_qwen2.py index 9dcdbf2b259e..9106216fb13d 100644 --- a/src/transformers/models/qwen2/configuration_qwen2.py +++ b/src/transformers/models/qwen2/configuration_qwen2.py @@ -73,8 +73,6 @@ class Qwen2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/qwen2/modeling_qwen2.py b/src/transformers/models/qwen2/modeling_qwen2.py index 2767340db60e..7dcc274a2cf2 100644 --- a/src/transformers/models/qwen2/modeling_qwen2.py +++ b/src/transformers/models/qwen2/modeling_qwen2.py @@ -417,9 +417,9 @@ def forward( class Qwen2ForCausalLM(Qwen2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 15c3b509a812..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 @@ -168,13 +168,12 @@ class Qwen2_5OmniTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. 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 9f8fc53ad1ef..e9eeefc9b9f8 100644 --- a/src/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py +++ b/src/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py @@ -269,6 +269,13 @@ class Qwen2_5OmniTextConfig(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", + } + 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 bf2723debf08..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 @@ -102,8 +102,6 @@ class Qwen2_5_VLTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py b/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py index 3066335ce3f0..329a6ef2a058 100644 --- a/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py +++ b/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py @@ -68,8 +68,6 @@ class Qwen2MoeConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py b/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py index 139d425c50bb..bb6871f46a05 100644 --- a/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py +++ b/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py @@ -79,8 +79,6 @@ class Qwen2VLTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/qwen3_5/configuration_qwen3_5.py b/src/transformers/models/qwen3_5/configuration_qwen3_5.py index 957a99ba07f0..04877a3bbdae 100644 --- a/src/transformers/models/qwen3_5/configuration_qwen3_5.py +++ b/src/transformers/models/qwen3_5/configuration_qwen3_5.py @@ -71,8 +71,6 @@ class Qwen3_5TextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", 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 baab11721f90..72c2dacfd2c0 100644 --- a/src/transformers/models/qwen3_5_moe/configuration_qwen3_5_moe.py +++ b/src/transformers/models/qwen3_5_moe/configuration_qwen3_5_moe.py @@ -72,8 +72,6 @@ class Qwen3_5MoeTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", 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 f64ff70a1f70..2a94d8559779 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 @@ -1945,6 +1945,7 @@ class Qwen3_5MoeForConditionalGeneration(Qwen3_5MoePreTrainedModel, GenerationMi accepts_loss_kwargs = False config: Qwen3_5MoeConfig _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} def __init__(self, config): diff --git a/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py b/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py index 87a8ac165d0c..4395b7c622bd 100644 --- a/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py +++ b/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py @@ -96,8 +96,6 @@ class Qwen3MoeConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/qwen3_next/configuration_qwen3_next.py b/src/transformers/models/qwen3_next/configuration_qwen3_next.py index 25810a96a0fa..b499291d70e2 100644 --- a/src/transformers/models/qwen3_next/configuration_qwen3_next.py +++ b/src/transformers/models/qwen3_next/configuration_qwen3_next.py @@ -77,8 +77,6 @@ class Qwen3NextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", 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 223c8932fab5..3c65badeafb4 100644 --- a/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py +++ b/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py @@ -148,6 +148,13 @@ class Qwen3OmniMoeTextConfig(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", + } + ignore_keys_at_rope_validation = {"mrope_section", "interleaved", "mrope_interleaved"} vocab_size: int = 3584 @@ -424,6 +431,12 @@ class Qwen3OmniMoeTalkerTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 3072 hidden_size: int = 1024 intermediate_size: int = 2048 diff --git a/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py b/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py index 2864f1f325ba..dd98da222cf3 100644 --- a/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py +++ b/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py @@ -290,6 +290,13 @@ class Qwen3OmniMoeTextConfig(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", + } + ignore_keys_at_rope_validation = {"mrope_section", "interleaved", "mrope_interleaved"} vocab_size: int = 3584 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 612edcd5776e..f49d389a593f 100644 --- a/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py +++ b/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py @@ -94,8 +94,6 @@ class Qwen3VLMoeTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/seed_oss/configuration_seed_oss.py b/src/transformers/models/seed_oss/configuration_seed_oss.py index f1145ba1f7a9..363b393dbc4a 100644 --- a/src/transformers/models/seed_oss/configuration_seed_oss.py +++ b/src/transformers/models/seed_oss/configuration_seed_oss.py @@ -74,8 +74,6 @@ class SeedOssConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/seed_oss/modeling_seed_oss.py b/src/transformers/models/seed_oss/modeling_seed_oss.py index b8bda4449d72..841765a7c928 100644 --- a/src/transformers/models/seed_oss/modeling_seed_oss.py +++ b/src/transformers/models/seed_oss/modeling_seed_oss.py @@ -430,9 +430,9 @@ def forward( class SeedOssForCausalLM(SeedOssPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/smollm3/configuration_smollm3.py b/src/transformers/models/smollm3/configuration_smollm3.py index 637f9ad75fb8..9a14e834db9e 100644 --- a/src/transformers/models/smollm3/configuration_smollm3.py +++ b/src/transformers/models/smollm3/configuration_smollm3.py @@ -84,8 +84,6 @@ class SmolLM3Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/smollm3/modeling_smollm3.py b/src/transformers/models/smollm3/modeling_smollm3.py index acb03a2ead2a..d82882382bb9 100644 --- a/src/transformers/models/smollm3/modeling_smollm3.py +++ b/src/transformers/models/smollm3/modeling_smollm3.py @@ -446,9 +446,9 @@ def forward( class SmolLM3ForCausalLM(SmolLM3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 2fb76822c447..69c9d4b51444 100644 --- a/src/transformers/models/smollm3/modular_smollm3.py +++ b/src/transformers/models/smollm3/modular_smollm3.py @@ -100,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 1d78ce714d70..5d9aec3ac7a0 100644 --- a/src/transformers/models/solar_open/configuration_solar_open.py +++ b/src/transformers/models/solar_open/configuration_solar_open.py @@ -49,6 +49,13 @@ class SolarOpenConfig(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", } @@ -92,14 +99,6 @@ class SolarOpenConfig(PreTrainedConfig): "layers.*.mlp.experts": "moe_experts_allreduce", "norm": "activation", } - - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. - base_model_fsdp_plan = { - "embed_tokens": "free_full_weight", - "layers.*": "free_full_weight", - "norm": "keep_full_weight", - } 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 e5f0c9cfafe6..53abac7c7d07 100644 --- a/src/transformers/models/solar_open/modeling_solar_open.py +++ b/src/transformers/models/solar_open/modeling_solar_open.py @@ -553,9 +553,9 @@ def forward( class SolarOpenForCausalLM(SolarOpenPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 92b212fcfbdd..48bba9439e08 100644 --- a/src/transformers/models/solar_open/modular_solar_open.py +++ b/src/transformers/models/solar_open/modular_solar_open.py @@ -64,6 +64,12 @@ class SolarOpenConfig(Glm4MoeConfig): "norm": "activation", } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 196608 moe_intermediate_size: int = 1280 num_hidden_layers: int = 48 diff --git a/src/transformers/models/starcoder2/configuration_starcoder2.py b/src/transformers/models/starcoder2/configuration_starcoder2.py index ea6310ca5278..88137e8f362f 100644 --- a/src/transformers/models/starcoder2/configuration_starcoder2.py +++ b/src/transformers/models/starcoder2/configuration_starcoder2.py @@ -72,8 +72,6 @@ class Starcoder2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/starcoder2/modeling_starcoder2.py b/src/transformers/models/starcoder2/modeling_starcoder2.py index caac32e2cebf..baabf385da59 100644 --- a/src/transformers/models/starcoder2/modeling_starcoder2.py +++ b/src/transformers/models/starcoder2/modeling_starcoder2.py @@ -410,9 +410,9 @@ def forward( class Starcoder2ForCausalLM(Starcoder2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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 3ad3a5271f9d..8c313bae875e 100644 --- a/src/transformers/models/t5gemma/configuration_t5gemma.py +++ b/src/transformers/models/t5gemma/configuration_t5gemma.py @@ -80,8 +80,6 @@ class T5GemmaModuleConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/t5gemma/modular_t5gemma.py b/src/transformers/models/t5gemma/modular_t5gemma.py index 7aa4599200db..f35a06beab33 100644 --- a/src/transformers/models/t5gemma/modular_t5gemma.py +++ b/src/transformers/models/t5gemma/modular_t5gemma.py @@ -84,6 +84,12 @@ class T5GemmaModuleConfig(Gemma2Config): >>> configuration = model.config ```""" + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + is_decoder: bool = False use_bidirectional_attention = AttributeError() diff --git a/src/transformers/models/t5gemma2/configuration_t5gemma2.py b/src/transformers/models/t5gemma2/configuration_t5gemma2.py index 2864e9daa240..b3c2ada26f35 100644 --- a/src/transformers/models/t5gemma2/configuration_t5gemma2.py +++ b/src/transformers/models/t5gemma2/configuration_t5gemma2.py @@ -76,8 +76,6 @@ class T5Gemma2TextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", @@ -108,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): @@ -274,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 @@ -298,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/modular_t5gemma2.py b/src/transformers/models/t5gemma2/modular_t5gemma2.py index bd61326efaa9..89b926f9d4a7 100644 --- a/src/transformers/models/t5gemma2/modular_t5gemma2.py +++ b/src/transformers/models/t5gemma2/modular_t5gemma2.py @@ -86,6 +86,13 @@ class T5Gemma2TextConfig(Gemma3TextConfig, PreTrainedConfig): """ model_type = "t5gemma2_text" + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + use_bidirectional_attention = AttributeError() def __post_init__(self, **kwargs): diff --git a/src/transformers/models/vaultgemma/configuration_vaultgemma.py b/src/transformers/models/vaultgemma/configuration_vaultgemma.py index 7c7654330b37..7ff8c838d362 100644 --- a/src/transformers/models/vaultgemma/configuration_vaultgemma.py +++ b/src/transformers/models/vaultgemma/configuration_vaultgemma.py @@ -79,8 +79,6 @@ class VaultGemmaConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/vaultgemma/modeling_vaultgemma.py b/src/transformers/models/vaultgemma/modeling_vaultgemma.py index 95cebf857c56..1ddba4af2210 100644 --- a/src/transformers/models/vaultgemma/modeling_vaultgemma.py +++ b/src/transformers/models/vaultgemma/modeling_vaultgemma.py @@ -468,9 +468,9 @@ def forward( class VaultGemmaForCausalLM(VaultGemmaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _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/voxtral_realtime/configuration_voxtral_realtime.py b/src/transformers/models/voxtral_realtime/configuration_voxtral_realtime.py index f47df59dd6a4..568c6f8748b9 100644 --- a/src/transformers/models/voxtral_realtime/configuration_voxtral_realtime.py +++ b/src/transformers/models/voxtral_realtime/configuration_voxtral_realtime.py @@ -41,8 +41,6 @@ class VoxtralRealtimeTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py b/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py index 7f8965a6de70..1e9f2dbcad26 100644 --- a/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py +++ b/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py @@ -831,9 +831,9 @@ def forward( class VoxtralRealtimeTextForCausalLM(VoxtralRealtimeTextPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/youtu/configuration_youtu.py b/src/transformers/models/youtu/configuration_youtu.py index 726df1cf3c94..6210f3c5f42b 100644 --- a/src/transformers/models/youtu/configuration_youtu.py +++ b/src/transformers/models/youtu/configuration_youtu.py @@ -61,8 +61,6 @@ class YoutuConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Bundled with other parallel plans for consistency; the applier walks - # this dict to decide what fully_shard each module gets. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/youtu/modeling_youtu.py b/src/transformers/models/youtu/modeling_youtu.py index 820146bf2e0a..8751f452fd28 100644 --- a/src/transformers/models/youtu/modeling_youtu.py +++ b/src/transformers/models/youtu/modeling_youtu.py @@ -534,9 +534,9 @@ def forward( class YoutuForCausalLM(YoutuPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) From 9a835ebab92678a758d1af7273c0b5a468811b36 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 19 May 2026 04:47:15 +0000 Subject: [PATCH 104/116] revert gitignore --- .gitignore | 5 ----- 1 file changed, 5 deletions(-) diff --git a/.gitignore b/.gitignore index fe847905730f..903efc854eef 100644 --- a/.gitignore +++ b/.gitignore @@ -8,11 +8,6 @@ __pycache__/ # C extensions *.so -verify_ckpt* -checkpoints* -*result* -debug* - # tests and logs tests/fixtures/cached_*_text.txt logs/ From 86dc7b4ac7a010868755b51281f3045c2a4775ef Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 19 May 2026 06:44:19 +0000 Subject: [PATCH 105/116] _apply within for loop --- .../distributed/tensor_parallel.py | 28 +++++-------------- 1 file changed, 7 insertions(+), 21 deletions(-) diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index 202c10af6c6e..126b404594f2 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -32,7 +32,6 @@ PrepareModuleInput, RowwiseParallel, SequenceParallel, - parallelize_module, ) from torch.distributed.tensor.parallel.style import ParallelStyle from torch.distributed.tensor.placement_types import _StridedShard @@ -462,11 +461,12 @@ class ParallelInterface(GeneralInterface): def apply_tensor_parallel(model, tp_mesh, tp_plan): - """Apply tensor parallelism using PyTorch's parallelize_module. + """Apply tensor parallelism by calling each style's ``_apply`` on the + matching submodules. - Converts the wildcard tp_plan from model config into a concrete plan - for ``parallelize_module``. Plan values are string names looked up in - ``ALL_PARALLEL_STYLES``. + Walks ``model.named_modules()``, resolves each name against the wildcard + ``tp_plan`` from the model config, and applies the corresponding style + from ``ALL_PARALLEL_STYLES`` (looked up by string name) directly. """ distributed_config = getattr(model.config, "distributed_config", None) sp_requested = getattr(distributed_config, "enable_sequence_parallel", False) @@ -488,25 +488,11 @@ def apply_tensor_parallel(model, tp_mesh, tp_plan): if not tied_source_in_plan: tp_plan.pop("lm_head", None) - parallelize_plan = {} - - for name, _ in model.named_modules(): + for name, submodule in model.named_modules(): style_value = _get_parameter_tp_plan(parameter_name=name, tp_plan=tp_plan, is_weight=False) if style_value is None: continue - - if not isinstance(style_value, str): - raise TypeError( - f"Unsupported plan value for '{name}': {style_value!r} (type {type(style_value).__name__}). " - f"TP plan values must be strings looked up in ALL_PARALLEL_STYLES." - ) - if style_value not in ALL_PARALLEL_STYLES: - raise ValueError( - f"Unknown TP style {style_value!r} for module {name!r}. Valid styles: {sorted(ALL_PARALLEL_STYLES)}" - ) - parallelize_plan[name] = ALL_PARALLEL_STYLES[style_value] - - parallelize_module(model, tp_mesh, parallelize_plan) + ALL_PARALLEL_STYLES[style_value]._apply(submodule, tp_mesh) # Under SP, inputs_embeds is sequence-sharded after embed_tokens, so # auto-generated position_ids would use the wrong (local) seq_len. From 9378ccb2b031390da0cdccc02d26b77b2949846d Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 19 May 2026 06:46:59 +0000 Subject: [PATCH 106/116] rename --- src/transformers/distributed/utils.py | 2 +- src/transformers/modeling_utils.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/transformers/distributed/utils.py b/src/transformers/distributed/utils.py index 71d409d239b1..50cf2fa4bb28 100644 --- a/src/transformers/distributed/utils.py +++ b/src/transformers/distributed/utils.py @@ -196,7 +196,7 @@ def gather_full_state_dict(model) -> dict[str, torch.Tensor]: return result -def save_model_checkpoint(model, checkpoint_dir: str) -> None: +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. diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index e75706ec4002..f3358c0ae659 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -64,7 +64,7 @@ distribute_model, gather_full_state_dict, init_device_mesh, - save_model_checkpoint, + save_model_checkpoint_distributed, ) from .dynamic_module_utils import custom_object_save from .generation import CompileConfig, GenerationConfig @@ -3466,7 +3466,7 @@ def save_pretrained( "save_pretrained(distributed_checkpoint=True) requires the model to have been " "initialized with a distributed_config (device_mesh is None)." ) - save_model_checkpoint(self, save_directory) + save_model_checkpoint_distributed(self, save_directory) return # Get the model state_dict (handles FSDP unshard + TP gather in one call) From e8785942f9627d9cb2211ce80a8b9bee5dd3b89b Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 19 May 2026 06:47:56 +0000 Subject: [PATCH 107/116] doc sp plan --- src/transformers/configuration_utils.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/transformers/configuration_utils.py b/src/transformers/configuration_utils.py index 862139d6ae21..eb22824db8b0 100755 --- a/src/transformers/configuration_utils.py +++ b/src/transformers/configuration_utils.py @@ -147,6 +147,10 @@ class PreTrainedConfig(PushToHubMixin, RotaryEmbeddingConfigMixin): naming of attributes. - **base_model_tp_plan** (`dict[str, Any]`) -- A dict that maps sub-modules FQNs of a base model to a tensor parallel plan applied to the sub-module when `model.tensor_parallel` is called. + - **base_model_sp_plan** (`dict[str, Any]`) -- A dict that maps sub-modules FQNs of a base model to a sequence + parallel plan, used in place of `base_model_tp_plan` when `distributed_config.enable_sequence_parallel` is set. + Same key/value shape as the TP plan; values are style names registered in `ALL_PARALLEL_STYLES` + (e.g. `"vocab_reduce_scatter"`, `"rowwise_reduce_scatter"`, `"activation"`, `"module_allgather"`). - **base_model_fsdp_plan** (`dict[Any, str]`) -- A dict that maps sub-modules of a base model to an FSDP2 sharding strategy (e.g. `"free_full_weight"` / `"keep_full_weight"`). Keys can be wildcard module paths (e.g. `"layers.*"`) or tuples of paths (grouped into a single `fully_shard` call). From 8caa870c65a5f910e705a238c590e206b36a53e5 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 19 May 2026 06:50:37 +0000 Subject: [PATCH 108/116] fix --- src/transformers/modeling_utils.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index f3358c0ae659..4742918a4acf 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -1380,13 +1380,11 @@ def post_init(self): cls_tp_plan = getattr(self, "_tp_plan", None) or {} cls_sp_plan = getattr(self, "_sp_plan", None) or {} cls_fsdp_plan = getattr(self, "_fsdp_plan", None) or {} - self._tp_plan, self._sp_plan, self._ep_plan, self._pp_plan, self._fsdp_plan = ( - dict(cls_tp_plan), - dict(cls_sp_plan), - {}, - {}, - dict(cls_fsdp_plan), - ) + self._tp_plan = dict(cls_tp_plan) + self._sp_plan = dict(cls_sp_plan) + self._ep_plan = {} + self._pp_plan = {} + self._fsdp_plan = dict(cls_fsdp_plan) # If current model is a base model, attach `base_model_*_plan` from config if self.base_model is self: self._pp_plan = self.config.base_model_pp_plan.copy() if self.config.base_model_pp_plan is not None else {} From e0e787bc1d11535a0de12aef069647de2c10d97c Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 19 May 2026 06:52:29 +0000 Subject: [PATCH 109/116] unified settattr + torch no grad + _local_tensor --- src/transformers/core_model_loading.py | 20 ++++++-------------- 1 file changed, 6 insertions(+), 14 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 6ce8c43d0d56..de1fac7cf7b5 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -1062,7 +1062,7 @@ def _format_op_name(curr_op: list[ConversionOps] | ConversionOps | None) -> str # Raise a specific Exception that we can catch easily raise SkipParameters() - +@torch.no_grad() def set_param_for_module( model: PreTrainedModel, target_name: str, @@ -1089,7 +1089,7 @@ def set_param_for_module( # Remove from missing keys (it's either mismatched, or all good) loading_info.missing_keys.discard(target_name) - expected_shape = ref.to_local().shape if isinstance(ref, DTensor) else 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)) @@ -1097,18 +1097,10 @@ def set_param_for_module( 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) - with torch.no_grad(): - if ref.is_meta: - torch.utils.swap_tensors( - ref, torch.nn.Parameter(dtensor_param, requires_grad=ref.requires_grad) - ) - else: - ref.copy_(dtensor_param) - ref._is_hf_initialized = True - else: - # super important otherwise _init_weight will re-init the param - param_value._is_hf_initialized = True - setattr(module_obj, param_name, param_value) + 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) def offload_and_maybe_resave_param( From 9158d99f3365c49f763fe3ced5f2c1dc664415cc Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 19 May 2026 06:54:28 +0000 Subject: [PATCH 110/116] revert --- src/transformers/modeling_utils.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 4742918a4acf..17dcda104c17 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -3779,7 +3779,10 @@ def get_init_context( elif is_quantized: init_contexts.extend([torch.device("meta"), set_quantized_state()]) else: - init_contexts.append(torch.device("meta")) + # meta_device_safe_creation_ops patches torch.linspace to default to CPU + # so that custom models calling .item() during __init__ (e.g. drop-path + # schedules) don't crash on meta tensors. + init_contexts.extend([torch.device("meta"), init.meta_device_safe_creation_ops()]) return init_contexts From 1c8203a0c9aee7e7a4dd4aa23ea31243d6e2f0ab Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 19 May 2026 07:03:29 +0000 Subject: [PATCH 111/116] linting --- src/transformers/core_model_loading.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index de1fac7cf7b5..e15c2314ac31 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -1062,6 +1062,7 @@ def _format_op_name(curr_op: list[ConversionOps] | ConversionOps | None) -> str # Raise a specific Exception that we can catch easily raise SkipParameters() + @torch.no_grad() def set_param_for_module( model: PreTrainedModel, From 7e2f68688487f91d235bc4048ad2d88ded7d82dd Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 19 May 2026 07:34:36 +0000 Subject: [PATCH 112/116] fix ruff --- tests/test_tensor_parallel_mixin.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_tensor_parallel_mixin.py b/tests/test_tensor_parallel_mixin.py index 6f2ae2256bad..d71c398c039e 100644 --- a/tests/test_tensor_parallel_mixin.py +++ b/tests/test_tensor_parallel_mixin.py @@ -180,7 +180,7 @@ def _verify_tp_sharding(rank, model_tp, model_ref): for dim in range(param_local.ndim): if param_local.size(dim) != param_full.size(dim): param_plan = _get_parameter_tp_plan(name, tp_plan, is_weight=True) - if _is_packed_colwise_plan(param_plan): + if param_plan == "packed_colwise": expected_size = param_full.size(dim) // world_size assert param_local.size(dim) == expected_size, ( f"Packed weight {name} sharding incorrect: expected {expected_size}, got {param_local.size(dim)}" @@ -265,7 +265,7 @@ def _test_tp_backward_impl(rank, model_path, model_class, atol, rtol): for dim in range(grad.ndim): if grad.size(dim) != grad_tp.size(dim): param_plan = _get_parameter_tp_plan(name, tp_plan, is_weight=True) - if _is_packed_colwise_plan(param_plan): + if param_plan == "packed_colwise": # interleaved slicing grad = get_packed_grad_shard(grad, world_size, rank, dim) else: From e6484d377f441c2457bd6ede5c88ff43ae5cd5a7 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 19 May 2026 07:48:13 +0000 Subject: [PATCH 113/116] make check-repository-consistency --- .../models/hrm_text/configuration_hrm_text.py | 21 +++++++++++++++++++ .../models/hrm_text/modeling_hrm_text.py | 4 +++- 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/src/transformers/models/hrm_text/configuration_hrm_text.py b/src/transformers/models/hrm_text/configuration_hrm_text.py index c8a328906b96..255391399960 100644 --- a/src/transformers/models/hrm_text/configuration_hrm_text.py +++ b/src/transformers/models/hrm_text/configuration_hrm_text.py @@ -65,12 +65,33 @@ class HrmTextConfig(PreTrainedConfig): **{f"{stack}.layers.*.mlp.up_proj": "colwise" for stack in ("L_module", "H_module")}, **{f"{stack}.layers.*.mlp.down_proj": "rowwise" for stack in ("L_module", "H_module")}, } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 151808 hidden_size: int = 1536 intermediate_size: int = 4096 diff --git a/src/transformers/models/hrm_text/modeling_hrm_text.py b/src/transformers/models/hrm_text/modeling_hrm_text.py index 9e10bed4997e..731f3cba1ebb 100644 --- a/src/transformers/models/hrm_text/modeling_hrm_text.py +++ b/src/transformers/models/hrm_text/modeling_hrm_text.py @@ -555,8 +555,10 @@ def forward( @auto_docstring class HrmTextForCausalLM(HrmTextPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) From 4b649ed049feae4e92ad3e7b994c62b816593cbd Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 19 May 2026 08:13:50 +0000 Subject: [PATCH 114/116] trigger fsdp mixin test in CI --- .circleci/create_circleci_config.py | 12 +++++++++++- utils/tests_fetcher.py | 1 + 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/.circleci/create_circleci_config.py b/.circleci/create_circleci_config.py index 26f254d141b7..54c69567c159 100644 --- a/.circleci/create_circleci_config.py +++ b/.circleci/create_circleci_config.py @@ -398,6 +398,15 @@ def job_name(self): parallelism=6, ) +fsdp_ci_job = CircleCIJob( + "fsdp_ci", + additional_env={"RUN_FSDP_TESTS": True}, + docker_image=[{"image": "huggingface/transformers-torch-light"}], + install_steps=["uv pip install .", "uv pip install torchao"], + marker="is_fsdp_test", + parallelism=6, +) + # We also include a `dummy.py` file in the files to be doc-tested to prevent edge case failure. Otherwise, the pytest # hangs forever during test collection while showing `collecting 0 items / 21 errors`. (To see this, we have to remove # the bash output redirection.) @@ -429,7 +438,8 @@ def job_name(self): DOC_TESTS = [doc_test_job] TRAINING_CI_TESTS = [training_ci_job] TENSOR_PARALLEL_CI_TESTS = [tensor_parallel_ci_job] -ALL_TESTS = REGULAR_TESTS + EXAMPLES_TESTS + PIPELINE_TESTS + REPO_UTIL_TESTS + DOC_TESTS + [custom_tokenizers_job] + [exotic_models_job] + TRAINING_CI_TESTS + TENSOR_PARALLEL_CI_TESTS # fmt: skip +FSDP_CI_TESTS = [fsdp_ci_job] +ALL_TESTS = REGULAR_TESTS + EXAMPLES_TESTS + PIPELINE_TESTS + REPO_UTIL_TESTS + DOC_TESTS + [custom_tokenizers_job] + [exotic_models_job] + TRAINING_CI_TESTS + TENSOR_PARALLEL_CI_TESTS + FSDP_CI_TESTS # fmt: skip def create_circleci_config(folder=None): diff --git a/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)", } From 6b835cb5b2cc3a21aafc906c848ef7b9858a4689 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 19 May 2026 08:35:27 +0000 Subject: [PATCH 115/116] fix fsdp ci --- tests/test_fsdp_mixin.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py index c1b3b84f80b6..1a7032489224 100644 --- a/tests/test_fsdp_mixin.py +++ b/tests/test_fsdp_mixin.py @@ -83,7 +83,6 @@ "qwen3_5_moe", "deepseek_v2", "gpt_oss", - "glm_moe_dsa", "glm4_moe_lite", } From d563a9c07cf228171855990d9eee9837addc1688 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 19 May 2026 07:08:15 +0000 Subject: [PATCH 116/116] Reset tests/test_modeling_common.py to main Restores legitimate improvements that were accidentally undone during a stale merge of main into fsdp-vs-ddp: - Restore test_resize_embeddings_untied_no_reinit_on_post_init - Restore clipseg / Timm / evolla / parakeet_* / pi0 / musicflamingo special-cases - Restore skip_base_model parameter on test_reverse_loading_mapping - Restore "is not None" guard on subconfig in test_initialization - Fix typo: "ot" -> "or" in test_reverse_loading_mapping assert message --- tests/test_modeling_common.py | 78 ++++++++++++++++++++++++++++++++--- 1 file changed, 73 insertions(+), 5 deletions(-) diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py index f6f2efa039e9..2e60939f9ead 100644 --- a/tests/test_modeling_common.py +++ b/tests/test_modeling_common.py @@ -446,6 +446,10 @@ def _can_output_attn(model): outputs_eager = outputs_eager["language_model_outputs"] outputs_sdpa = outputs_sdpa["language_model_outputs"] key = "hidden_states" if "hidden_states" in outputs_eager else "decoder_hidden_states" + elif "decoder_output" in outputs_eager and "clipseg" in model_class.__name__.lower(): + outputs_eager = outputs_eager["decoder_output"] + outputs_sdpa = outputs_sdpa["decoder_output"] + key = "hidden_states" if "hidden_states" in outputs_eager else "decoder_hidden_states" else: key = "hidden_states" @@ -1303,7 +1307,7 @@ def test_init_weights_can_init_buffers(self): config.scale = 0 for sub_key in config.sub_configs: subconfig = getattr(config, sub_key) - if hasattr(subconfig, "scale"): + if subconfig is not None and hasattr(subconfig, "scale"): subconfig.scale = 0 for model_class in self.all_model_classes: @@ -1752,7 +1756,10 @@ def _set_subconfig_attributes(self, config, attribute_name, value): """Helper function to recursively set a config attr to a given value""" for k in config.sub_configs: if ( - self._is_composite and attribute_name == "output_attentions" and k == "vision_config" + self._is_composite + and attribute_name == "output_attentions" + and k == "vision_config" + and "Timm" in getattr(config, k).__class__.__name__ ): # skip because it's not needed and causes errors e.g with Timm continue if getattr(config, k) is not None: @@ -2403,6 +2410,55 @@ def test_resize_embeddings_untied_with_deepspeed_multi_gpu(self): with _deepspeed_zero3(ds_config): self.test_resize_embeddings_untied() + def test_resize_embeddings_untied_no_reinit_on_post_init(self): + if not self.test_resize_embeddings: + self.skipTest(reason="test_resize_embeddings is set to `False`") + + original_config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + original_config.tie_word_embeddings = False + try: + original_config.get_text_config().tie_word_embeddings = False + except Exception as e: + model_type = getattr(original_config, "model_type", "unknown") + print(f"Could not set text config's `tie_word_embeddings` for model type `{model_type}`: {e}") + + if original_config.tie_word_embeddings: + self.skipTest(reason="Model cannot untie embeddings") + + for model_class in self.all_model_classes: + with self.subTest(model_class): + config = copy.deepcopy(original_config) + model = model_class(config).to(torch_device) + model.eval() + + # The bug only affects nn.Linear LM heads created by _get_resized_lm_head + output_embeds = model.get_output_embeddings() + if not isinstance(output_embeds, nn.Linear): + continue + + model_vocab_size = config.get_text_config().vocab_size + try: + model.resize_token_embeddings(model_vocab_size + 10) + except (NotImplementedError, AttributeError): + continue + + output_embeds = model.get_output_embeddings() + weights_before = output_embeds.weight.data.clone() + bias_before = output_embeds.bias.data.clone() if output_embeds.bias is not None else None + + model.post_init() + + output_embeds_after = model.get_output_embeddings() + self.assertTrue( + torch.equal(weights_before, output_embeds_after.weight.data), + "Output embedding weights were reinitialized by post_init() after resize_token_embeddings()", + ) + if bias_before is not None: + self.assertTrue( + torch.equal(bias_before, output_embeds_after.bias.data), + "Output embedding bias was reinitialized by post_init() after resize_token_embeddings()", + ) + def test_model_get_set_embeddings(self): config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() @@ -3619,6 +3675,7 @@ def test_sdpa_can_dispatch_on_flash(self): "PaliGemma-like models currently (transformers==4.41.0) requires an attention_mask input" ) if config.model_type in [ + "evolla", "modernbert", "gemma3", "t5gemma", @@ -3630,6 +3687,9 @@ def test_sdpa_can_dispatch_on_flash(self): "kosmos-2", "mllama", "lighton_ocr", + "parakeet_encoder", + "parakeet_ctc", + "pi0", "pixtral", "sam", "sam_hq", @@ -4690,7 +4750,7 @@ def test_tp_plan_matches_params(self): len(unused_entries) == 0, f"The following entries of the TP-plan are not valid: {unused_entries}" ) - def test_reverse_loading_mapping(self, check_keys_were_modified=True): + def test_reverse_loading_mapping(self, check_keys_were_modified=True, skip_base_model=False): """Make sure we can load and save correctly the models having any weight renaming mapping or weight conversion mapping. Note that this test would be better if we could start from the serialized keys, and check that the model @@ -4704,6 +4764,11 @@ def test_reverse_loading_mapping(self, check_keys_were_modified=True): check_keys_were_modified (`bool`, *optional*, defaults to `True`): Whether to expect keys being modified or not. In some cases, models do not change keys but their weights, e.g. via transpose, memory alignment, etc. + skip_base_model (`bool`, *optional*, defaults to `False`): + Sometimes, mappings are only visible when applied to the model with head, and not visible on the + base model. This allows to skip the check on the base model. See e.g. `llava` mapping where this + is the case. In practice, the mappings are still coherent and a base model can still be loaded from + the head model, thanks to the `base_model_prefix` which will remove the prefix automatically. """ config, _ = self.model_tester.prepare_config_and_inputs_for_common() @@ -4716,6 +4781,8 @@ def test_reverse_loading_mapping(self, check_keys_were_modified=True): config_to_set.num_dense_layers = 1 # lfm2_moe for model_class in self.all_model_classes: + if skip_base_model and "For" not in model_class.__name__: + continue # Each individual model is a subtest with self.subTest(model_class.__name__): model = model_class(copy.deepcopy(config)) @@ -4784,7 +4851,7 @@ def test_reverse_loading_mapping(self, check_keys_were_modified=True): self.assertTrue( f"g{pattern_index}" in matched_groups, f"`{source_pattern}` in `{conversion}` did not match any of the source keys. " - "This indicates whether that the pattern is not properly written, ot that it could not be reversed correctly", + "This indicates whether that the pattern is not properly written, or that it could not be reversed correctly", ) # If everything is still good at this point, let's test that we perform the same operations both when @@ -4821,7 +4888,7 @@ def test_can_load_from_already_mapped_keys(self): # Skip if no conversions conversions = get_model_conversion_mapping(model, add_legacy=False) if len(conversions) == 0: - self.skipTest("No conversion found for this model") + self.skipTest(f"No conversion found for {model_class}") with tempfile.TemporaryDirectory() as tmpdirname: # Serialize without reverting the mapping @@ -4887,6 +4954,7 @@ def _audio_features_prepare_config_and_inputs(self): or "input_values" in key or "input_features" in key or key in ["padding_mask", "is_longer", "feature_attention_mask"] + or (config.model_type == "musicflamingo" and key == "input_ids") } return config, inputs_dict