diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index ec0a4c41a0ba..4c4733d14ad1 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -33,14 +33,12 @@ from .distributed.sharding_utils import DtensorShardOperation, _dtensor_from_local_like from .distributed.utils import is_dtensor from .integrations.accelerate import get_device, offload_weight -from .integrations.tensor_parallel import ALL_PARALLEL_STYLES from .utils import is_env_variable_true from .utils.loading_report import LoadStateDictInfo from .utils.logging import get_logger, tqdm if TYPE_CHECKING: - from .integrations.tensor_parallel import TensorParallelLayer from .modeling_utils import LoadStateDictConfig, PreTrainedModel from .quantizers import HfQuantizer @@ -1232,9 +1230,9 @@ def spawn_materialize( ) -> Future | Callable: """Materialize (and optionally shard) a tensor, asynchronously if a thread pool is provided. - When ``sharding_op`` is given the tensor is sharded (DTensor placement or legacy TP plan); - otherwise it is simply copied to *device*/*dtype*. Without a thread pool a deferred - callable is returned instead of a Future. + When ``sharding_op`` is given the tensor is sharded according to its DTensor placements; + otherwise it is simply copied to *device*/*dtype*. Without a thread pool a deferred callable + is returned instead of a Future. """ def _job(): @@ -1319,7 +1317,6 @@ def set_param_for_module( target_name: str, param_value: torch.Tensor, loading_info: LoadStateDictInfo, - distributed_operation: TensorParallelLayer | None, hf_quantizer: HfQuantizer, ): module_path, _, param_name = target_name.rpartition(".") @@ -1341,13 +1338,8 @@ def set_param_for_module( # Remove from missing keys (it's either mismatched, or all good) loading_info.missing_keys.discard(target_name) - # Determine expected shape: for TP/Dtensor, use sharded shape; otherwise, use full shape - if distributed_operation is not None: - expected_shape = torch.Size(distributed_operation.get_expected_sharded_shape(ref.shape)) - elif is_dtensor(ref): - expected_shape = ref._local_tensor.shape - else: - expected_shape = ref.shape + # For DTensor parameters, compare against the local shard loaded on this rank. + expected_shape = ref._local_tensor.shape if is_dtensor(ref) 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)) @@ -1355,12 +1347,12 @@ def set_param_for_module( if is_dtensor(ref): local_param = param_value.detach() if isinstance(param_value, torch.nn.Parameter) else param_value dtensor_param = _dtensor_from_local_like(local_param, ref) - param_value = torch.nn.Parameter(dtensor_param, requires_grad=ref.requires_grad) + param_value = torch.nn.Parameter( + dtensor_param, requires_grad=ref.requires_grad and dtensor_param.is_floating_point() + ) # 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( @@ -1466,7 +1458,6 @@ def convert_and_load_state_dict_in_model( model: PreTrainedModel, state_dict: dict[str, Any], load_config: LoadStateDictConfig, - tp_plan: dict[str, str] | None, disk_offload_index: dict | None = None, ): r""" @@ -1556,11 +1547,9 @@ def convert_and_load_state_dict_in_model( """ base_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 disk_offload_folder = load_config.disk_offload_folder offload_buffers = load_config.offload_buffers dtype_plan = load_config.dtype_plan or {} @@ -1595,10 +1584,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())) @@ -1673,23 +1658,13 @@ def convert_and_load_state_dict_in_model( else None ) - # 4. Handle TP/Dtensor sharding or device_map placement + # 4. Handle DTensor sharding or device_map placement param_device = get_device(device_map, renamed_key, valid_torch_device=True) sharding_op = None materialize_device = param_device if is_dtensor(empty_param): sharding_op = DtensorShardOperation(empty_param) - elif 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() - ) - sharding_op = mapping.distributed_operation - materialize_device = device_map[""] future_or_tensor = spawn_materialize( thread_pool, @@ -1732,7 +1707,6 @@ def convert_and_load_state_dict_in_model( target_name, param, loading_info, - mapping.distributed_operation, hf_quantizer, ) diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py index 0710076f350e..3ea0e294bf04 100644 --- a/src/transformers/distributed/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -17,9 +17,9 @@ import os from typing import TYPE_CHECKING, Any -from ..integrations.tensor_parallel import replace_layer_number_by_wildcard from ..utils import is_torch_available, is_torch_distributed_available, is_torch_greater_or_equal, logging, strtobool from ..utils.quantization_config import QuantizationMethod +from .tensor_parallel import replace_layer_number_by_wildcard from .utils import _is_torch_distributed_initialized diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py index 2ae34bcab24d..c73287844767 100644 --- a/src/transformers/distributed/mixin.py +++ b/src/transformers/distributed/mixin.py @@ -18,16 +18,15 @@ import warnings from typing import TYPE_CHECKING -from ..integrations.tensor_parallel import ( - ALL_PARALLEL_STYLES, - apply_tensor_parallelism, - gather_state_dict_for_save, - initialize_tensor_parallelism, -) from ..utils import is_torch_greater_or_equal, logging from ..utils.hub import create_and_tag_model_card from .configuration_utils import DistributedConfig from .fsdp import apply_fully_sharded_data_parallelism, is_fsdp_managed_module +from .tensor_parallel import ( + ALL_PARALLEL_STYLES, + apply_tensor_parallelism, + gather_state_dict_for_save, +) from .utils import ( _distributed_barrier, _ensure_torch_distributed, @@ -36,6 +35,7 @@ _is_torch_distributed_initialized, gather_full_state_dict, initialize_fully_sharded_data_parallelism, + initialize_tensor_parallelism, save_model_checkpoint_distributed, ) @@ -48,11 +48,7 @@ class DistributedMixin: - """Distributed orchestration and save/load hooks for [`PreTrainedModel`]. - - Stateless heavy lifting stays in `transformers.distributed.*` and - `integrations.tensor_parallel`. This mixin owns orchestration and instance state. - """ + """Distributed orchestration and save/load hooks for [`PreTrainedModel`].""" _device_mesh = None _tp_plan: dict[str, str] | None = None @@ -197,15 +193,13 @@ def maybe_distribute_model( model._device_mesh = device_mesh if distributed_config.tp_size > 1: - model = apply_tensor_parallelism( - model, - distributed_config.tp_plan, - distributed_config, - device_mesh, - ) + tp_mesh = device_mesh["tp"] if device_mesh.ndim > 1 else device_mesh + model = apply_tensor_parallelism(model, tp_mesh) + elif distributed_config.fsdp_size > 1: fsdp_mesh = device_mesh["fsdp"] if device_mesh.ndim > 1 else device_mesh model = apply_fully_sharded_data_parallelism(model, fsdp_mesh) + return model def should_save_on_this_rank(self, is_main_process: bool) -> bool: @@ -267,7 +261,9 @@ def gather_sharded_state_dict_for_save( return state_dict if distributed_config.tp_size > 1: - state_dict = gather_state_dict_for_save(state_dict, self._tp_plan, self._device_mesh, self._tp_size) + state_dict = gather_state_dict_for_save( + state_dict, self._tp_plan, self._device_mesh, distributed_config.tp_size + ) if not save_on_this_rank: state_dict = {} return state_dict diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py new file mode 100644 index 000000000000..4534ec46887e --- /dev/null +++ b/src/transformers/distributed/tensor_parallel.py @@ -0,0 +1,817 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import contextlib +import re + +from ..utils import logging +from ..utils.generic import GeneralInterface +from ..utils.import_utils import is_torch_available, is_torch_distributed_available + + +logger = logging.get_logger(__name__) + +if is_torch_available(): + import torch + +if is_torch_distributed_available(): + import torch.distributed as dist + from torch.distributed.tensor import DTensor, Partial, Replicate, Shard, distribute_tensor + from torch.distributed.tensor.placement_types import _StridedShard + + +def replace_layer_number_by_wildcard(name: str) -> str: + """ + Replace the numbers in the `name` by wildcards, only if they are in-between dots (`.`) or if they are between + a dot (`.`) and the end of the string. + This matches how modules are named/numbered when using a nn.ModuleList or nn.Sequential, but will NOT match + numbers in a parameter name itself, e.g. if the param is named `"w1"` or `"w2"`. + """ + return re.sub(r"\.\d+(\.|$)", lambda m: ".*" + m.group(1), name) + + +def verify_tp_plan(expected_keys: list[str], tp_plan: dict[str, str] | None): + """ + Verify the TP plan of the model, log a warning if the layers that were not sharded and the rules that were not applied. + """ + + if tp_plan is None: + return + + generic_keys = {replace_layer_number_by_wildcard(key) for key in expected_keys} + unsharded_layers = set(generic_keys) + unused_rules = tp_plan.copy() + + for key in generic_keys: + param_name = key.rsplit(".", 1)[0] if "." in key else key + generic_param_name = re.sub(r"\d+", "*", param_name) + + if generic_param_name in tp_plan: + unused_rules.pop(generic_param_name, None) + unsharded_layers.discard(key) + elif "." in generic_param_name and (parent_param_name := generic_param_name.rsplit(".", 1)[0]) in tp_plan: + unused_rules.pop(parent_param_name, None) + unsharded_layers.discard(key) + + if len(unused_rules) > 0: + logger.warning(f"The following TP rules were not applied on any of the layers: {unused_rules}") + if len(unsharded_layers) > 0: + logger.warning(f"The following layers were not sharded: {', '.join(unsharded_layers)}") + + +def _get_parameter_tp_plan(parameter_name: str, tp_plan: dict[str, str], is_weight=True) -> str | None: + """ + Get the TP style for a parameter from the TP plan. + + The TP plan is a dictionary that maps parameter names to TP styles. + The parameter name can be a generic name with wildcards (e.g. "*.weight") or a specific name (e.g. "layer_1.weight"). + + The `is_weight` is important because for weights, we want to support `.weights` and `.bias` cases seamlessly! but + not parent classes for `post_init` calls + """ + generic_param_name = replace_layer_number_by_wildcard(parameter_name) + if generic_param_name in tp_plan: + return tp_plan[generic_param_name] + elif is_weight and "." in generic_param_name and (module_name := generic_param_name.rsplit(".", 1)[0]) in tp_plan: + return tp_plan[module_name] + return None + + +@contextlib.contextmanager +def _use_local_dtensor_params(module): + # Kernels as DeepGEMM require local tensors rather than DTensors. + # We temporarily convert the DTensors to local tensors for the duration of forward() and swap them back after. + originals = {name: param for name, param in module.named_parameters(recurse=False) if isinstance(param, DTensor)} + local_params = {name: param.to_local() for name, param in originals.items()} + module._parameters.update(local_params) + try: + yield + finally: + for name, original in originals.items(): + current = module._parameters[name] + + if current is local_params[name]: + module._parameters[name] = original + + # Some kernels such as Megamoe performs changing on FP4 weights and scale factors during its forward pass. + # That implies creating a new Parameter so we should not restore the original DTensor. + elif current is not None and not isinstance(current, DTensor): + replacement = DTensor.from_local( + current, + original.device_mesh, + original.placements, + run_check=False, + ) + module._parameters[name] = torch.nn.Parameter( + replacement, + requires_grad=current.requires_grad, + ) + + +class TensorParallelLayer: + def should_use_local_tensors(self, module): + """Whether this module's forward requires local inputs and parameters.""" + return False + + def validate_param(self, module, param, mesh, parameter_name=None): + """Validate a parameter before applying this TP style.""" + pass + + def shard_param(self, module, param, mesh): + """Wrap ONE parameter as a DTensor placeholder. Default: no-op.""" + pass + + def transform_inputs_pre_forward(self, module, args, kwargs, mesh): + return args, kwargs + + def context_around_forward(self, module, mesh): + if self.should_use_local_tensors(module): + return _use_local_dtensor_params(module) + return contextlib.nullcontext() + + def transform_output_post_forward(self, module, output, mesh): + return output + + def install_forward(self, module, mesh): + """Install pre / around / post transforms by replacing module.forward.""" + original_forward = module.forward + + def tp_forward(*args, **kwargs): + args, kwargs = self.transform_inputs_pre_forward(module, args, kwargs, mesh) + with self.context_around_forward(module, mesh): + output = original_forward(*args, **kwargs) + return self.transform_output_post_forward(module, output, mesh) + + module.forward = tp_forward + return module + + +class ColwiseParallel(TensorParallelLayer): + """Column-wise: weight & bias → Shard(0) (Embedding: Shard(1)); input replicated, output Shard(-1).""" + + def __init__(self, *, input_layouts=None, output_layouts=None, use_local_output: bool = True): + self.input_layouts = input_layouts or Replicate() + self.output_layouts = output_layouts if output_layouts is not None else Shard(-1) + self.use_local_output = use_local_output + + def should_use_local_tensors(self, module): + use_local_quantized_path = getattr(module, "_hf_quantized_needs_local_tp", False) + uses_local_inference_kernel = isinstance(module, torch.nn.Linear) and not torch.is_grad_enabled() + return use_local_quantized_path or uses_local_inference_kernel + + def validate_param(self, module, param, mesh, parameter_name=None): + meta = module._parameters.get(param) + gathers_output = isinstance(self.output_layouts, Replicate) + if meta is None or not gathers_output: + return + + shard_dim = 1 if isinstance(module, torch.nn.Embedding) else meta.ndim - 2 + output_size = meta.shape[shard_dim] + tp_size = mesh.size() + if output_size % tp_size != 0: + parameter_name = parameter_name or param + layer_name = parameter_name.rsplit(".", 1)[0] + raise ValueError( + f"The output size of `{layer_name}` ({output_size}) must be divisible by the tensor parallel size " + f"({tp_size}) when gathering a colwise output." + ) + + def shard_param(self, module, param, mesh): + meta = module._parameters.get(param) + if meta is None: + return + placement = Shard(1) if isinstance(module, torch.nn.Embedding) else Shard(meta.ndim - 2) + module._parameters[param] = torch.nn.Parameter( + distribute_tensor(meta, mesh, [placement], src_data_rank=None), + requires_grad=meta.requires_grad, + ) + + def transform_inputs_pre_forward(self, module, args, kwargs, mesh): + x = args[0] + is_local_input = not isinstance(x, DTensor) + + # regular Linear inference fast path: the input is already replicated, the + # local weight owns a shard of the output features, and no backward is needed. + if is_local_input and isinstance(module, torch.nn.Linear) and not torch.is_grad_enabled(): + return args, kwargs + + # Plain-local quantized fast path. The input is already replicated, so avoid + # wrapping it as a DTensor. If its gradient is needed, synchronize the partial + # input-gradient contributions explicitly. + if is_local_input and getattr(module, "_hf_quantized_needs_local_tp", False): + if torch.is_grad_enabled() and x.requires_grad: + process_group = mesh.get_group() if mesh.ndim == 1 else mesh.get_group("tp") + x = _AllReduceBackward.apply(x, process_group) + return (x,) + args[1:], kwargs + + # DTensor path: handle regular training and layout redistribution. + if is_local_input: + x = DTensor.from_local(x, mesh, [self.input_layouts], run_check=False) + if x.placements != (Replicate(),): + x = x.redistribute(placements=[Replicate()]) + if self.should_use_local_tensors(module): + x = x.to_local(grad_placements=[Partial()]) + return (x,) + args[1:], kwargs + + def transform_output_post_forward(self, module, output, mesh): + # The local forward produced this rank's shard of the output features (last dim). + output_is_local_shard = ( + not isinstance(output, DTensor) + and isinstance(self.output_layouts, Shard) + and self.output_layouts.dim in (-1, output.dim() - 1) + ) + if self.should_use_local_tensors(module) and self.use_local_output and output_is_local_shard: + return output + if not isinstance(output, DTensor): + output = DTensor.from_local(output, mesh, [Shard(-1)], run_check=False) + if output.placements != (self.output_layouts,): + output = output.redistribute(placements=[self.output_layouts]) + return output.to_local() if self.use_local_output else output + + +class RowwiseParallel(TensorParallelLayer): + """Row-wise: weight → Shard(1), bias → Replicate (Embedding: weight → Shard(0)). + + Linear input is sharded on the last dim; Embedding input is replicated. The module + forward produces a Partial output which the boundary redistribute reduces to + output_layouts (Replicate → allreduce, Shard(1) → reduce-scatter). + """ + + def __init__(self, *, input_layouts=None, output_layouts=None, use_local_output: bool = True): + self.input_layouts = input_layouts or Shard(-1) + self.output_layouts = output_layouts or Replicate() + self.use_local_output = use_local_output + + def should_use_local_tensors(self, module): + use_local_quantized_path = getattr(module, "_hf_quantized_needs_local_tp", False) + uses_local_inference_kernel = isinstance(module, torch.nn.Linear) and not torch.is_grad_enabled() + return use_local_quantized_path or uses_local_inference_kernel + + def shard_param(self, module, param, mesh): + meta = module._parameters.get(param) + if meta is None: + return + if isinstance(module, torch.nn.Embedding): + placement = Shard(0) + else: + # bias is replicated (added after the row-reduce); weight shards on input dim (-1) + placement = Replicate() if param == "bias" else Shard(-1) + module._parameters[param] = torch.nn.Parameter( + distribute_tensor(meta, mesh, [placement], src_data_rank=None), + requires_grad=meta.requires_grad, + ) + + def transform_inputs_pre_forward(self, module, args, kwargs, mesh): + # Embedding runtime sharding needs a replicated input; Linear needs Shard(-1). + desired = Replicate() if isinstance(module, torch.nn.Embedding) else Shard(-1) + x = args[0] + # A local kernel can use a plain input when the previous layer already split it. + # avoid a redundant Tensor -> DTensor -> Tensor round trip. + input_has_desired_layout = self.input_layouts == desired + if self.should_use_local_tensors(module) and input_has_desired_layout and not isinstance(x, DTensor): + return args, kwargs + if not isinstance(x, DTensor): + x = DTensor.from_local(x, mesh, [self.input_layouts], run_check=False) + if x.placements != (desired,): + x = x.redistribute(placements=[desired]) + if self.should_use_local_tensors(module): + x = x.to_local() + return (x,) + args[1:], kwargs + + @contextlib.contextmanager + def context_around_forward(self, module, mesh): + if not self.should_use_local_tensors(module): + yield + else: + # A rowwise local forward must produce only its partial matmul. If we don't hide + # the bias, we will be adding the bias x world_size times which is not correct. + # We should add it once after the all_reduce (redistribute). + bias = module._parameters.get("bias") + if bias is not None: + module._parameters["bias"] = None + try: + with _use_local_dtensor_params(module): + yield + finally: + if bias is not None: + module._parameters["bias"] = bias + + def transform_output_post_forward(self, module, output, mesh): + use_local_inference_path = ( + isinstance(module, torch.nn.Linear) + and not isinstance(output, DTensor) + and isinstance(self.output_layouts, Replicate) + and not output.requires_grad + ) + if use_local_inference_path: + process_group = mesh.get_group() if mesh.ndim == 1 else mesh.get_group("tp") + dist.all_reduce(output, group=process_group) + if (bias := module._parameters.get("bias")) is not None: + output = output + (bias.to_local() if isinstance(bias, DTensor) else bias) + else: + # Dtensor tracks whether the result is partial, sharded, or replicated. + if not isinstance(output, DTensor): + output = DTensor.from_local(output, mesh, [Partial()], run_check=False) + if output.placements != (self.output_layouts,): + output = output.redistribute(placements=[self.output_layouts]) + if self.should_use_local_tensors(module) and (bias := module._parameters.get("bias")) is not None: + output = output + bias + if self.use_local_output: + output = output.to_local() + + return output + + +class ReplicatedWithGradAllReduce(TensorParallelLayer): + """Replicated parameter whose gradient is partial. + + For norms that sit between a colwise and a rowwise layer and normalize along a sharded + axis — e.g. Qwen3's per-head ``q_norm``/``k_norm``, which only see this rank's heads. The + forward needs no collective (the param is replicated and the activation is already local), + but each rank's parameter gradient only covers its own heads, so the gradients have to be + summed across the mesh. + """ + + def install_forward(self, module, mesh): + # A module hook rather than `param.register_hook`: params are replaced during weight + # loading, which happens after TP is applied, and would drop a param-level hook. + def _all_reduce_grads(mod, grad_input, grad_output): + for param in mod.parameters(recurse=False): + if param.grad is not None: + dist.all_reduce(param.grad, group=mesh.get_group()) + + module.register_full_backward_hook(_all_reduce_grads) + return module + + +class AllReduceParallel(TensorParallelLayer): + """All-reduce a module's partial forward output across the TP mesh.""" + + def transform_output_post_forward(self, module, output, mesh): + if output is None: + return None + if not isinstance(output, DTensor): + output = DTensor.from_local(output, mesh, [Partial()], run_check=False) + if output.placements != (Replicate(),): + output = output.redistribute(placements=[Replicate()]) + return output.to_local() + + +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 transform_output_post_forward(self, module, output, mesh): + rope_dim = module.config.qk_rope_head_dim + pass_output, rope_output = output.split([output.shape[-1] - rope_dim, rope_dim], dim=-1) + rope_output = _AllReduceBackward.apply(rope_output, mesh.get_group()) + return torch.cat([pass_output, rope_output], dim=-1) + + +class SequenceParallel(TensorParallelLayer): + def __init__(self, *, sequence_dim: int = 1, use_local_output: bool = True): + self.sequence_dim = sequence_dim + self.use_local_output = use_local_output + + def install_forward(self, module, mesh): + # Replicate the module's params (LayerNorm/RMSNorm ones-init → from_local is safe). + for p_name, p in list(module.named_parameters(recurse=False)): + module.register_parameter( + p_name, torch.nn.Parameter(DTensor.from_local(p, mesh, [Replicate()], run_check=False)) + ) + return super().install_forward(module, mesh) + + def transform_inputs_pre_forward(self, module, args, kwargs, mesh): + seq = Shard(self.sequence_dim) + x = args[0] + if not isinstance(x, DTensor): + x = DTensor.from_local(x, mesh, [seq], run_check=False) + elif x.placements != (seq,): + x = x.redistribute(placements=[seq]) + return (x,) + args[1:], kwargs + + def transform_output_post_forward(self, module, output, mesh): + if isinstance(output, DTensor): + return output.to_local() if self.use_local_output else output + return output + + +class PackedColwiseParallel(TensorParallelLayer): + """Column-wise parallel style for fused linear weights packed along the output dimension.""" + + def __init__( + self, + *, + use_local_output: bool = True, + split_factor: int = 2, + ): + self.input_layouts = (Replicate(),) + self.use_local_output = use_local_output + self.split_factor = split_factor + # Same as ColwiseParallel: replicated input, output-dim sharded weight, so the input + # gradient is partial. + self.input_grad_placements = [Partial()] + + def should_use_local_tensors(self, module): + return True + + def _packed_output_shard_dim(self, param_ndim: int) -> int: + """Dimension holding packed gate/up features: dim 0 for 2D Linear, dim 1 for 3D MoE experts.""" + if param_ndim == 1: + return -1 + return param_ndim - 2 + + def shard_param(self, module, param, mesh): + meta = module._parameters.get(param) + if meta is None: + return + shard_dim = self._packed_output_shard_dim(meta.ndim) + # Wrap as a DTensor placeholder. Runs on meta — distribute_tensor builds metadata only. + if meta.ndim == 1: + placement = Shard(shard_dim) + else: + placement = _StridedShard(dim=shard_dim, split_factor=self.split_factor) + module._parameters[param] = torch.nn.Parameter( + distribute_tensor(meta, mesh, [placement], src_data_rank=None), + requires_grad=meta.requires_grad, + ) + + def transform_inputs_pre_forward(self, module, args, kwargs, mesh): + input_tensor = args[0] + # Ensure the input is a Replicate DTensor on the TP mesh. + 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) + + # The packed kernels runs on local tensors, so Dtensor cannot infer the layout of the + # gradient produced by the kernel. The kernel sees replicated input + partial weights + # which means input gradient will be partial as well + input_tensor = input_tensor.to_local(grad_placements=[Partial()]) + return (input_tensor,) + args[1:], kwargs + + 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 + ) + + +class PackedRowwiseParallel(TensorParallelLayer): + """Parameter style for fused weights packed along the final dimension.""" + + def __init__(self, *, split_factor: int = 2): + self.split_factor = split_factor + + def shard_param(self, module, param, mesh): + meta = module._parameters.get(param) + if meta is None: + return + placement = Replicate() if meta.ndim == 1 else _StridedShard(dim=-1, split_factor=self.split_factor) + module._parameters[param] = torch.nn.Parameter( + distribute_tensor(meta, mesh, [placement], src_data_rank=None), + requires_grad=meta.requires_grad, + ) + + +class MoEParamShard(TensorParallelLayer): + """Param-only EP style for MoE expert weights (``grouped_gemm``). + + Shards dim 0 and updates module.num_experts to the per-rank local count so the + experts forward and ep_router sentinel agree. + """ + + def __init__(self, placement, *, shards_expert_dim: bool = False): + self.placement = placement + self.shards_expert_dim = shards_expert_dim + + def shard_param(self, module, param, mesh): + meta = module._parameters.get(param) + if meta is None: + return + if self.shards_expert_dim and hasattr(module, "num_experts"): + module.num_experts = meta.shape[0] // mesh.size() + module._parameters[param] = torch.nn.Parameter( + distribute_tensor(meta, mesh, [self.placement], src_data_rank=None), + requires_grad=meta.requires_grad, + ) + + +if is_torch_distributed_available(): + + class _AllReduceForward(torch.autograd.Function): + """Allreduce-sum forward, identity backward.""" + + @staticmethod + def forward(ctx, x, process_group): + if dist.get_world_size(process_group) > 1: + dist.all_reduce(x, group=process_group) + return x + + @staticmethod + def backward(ctx, grad): + return grad, None + + 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): + grad = grad.contiguous() + dist.all_reduce(grad, group=ctx.process_group) + return grad, None + + +class MoeExpertsParallel(TensorParallelLayer): + def should_use_local_tensors(self, module): + return True + + def transform_inputs_pre_forward(self, module, args, kwargs, mesh, *, is_expert_parallel=False): + hidden_states, *routing_args = args + tp_group = mesh.get_group() if mesh.ndim == 1 else mesh.get_group("tp") + if isinstance(hidden_states, DTensor): + hidden_states = hidden_states.to_local() + hidden_states = _AllReduceBackward.apply(hidden_states, tp_group) + + if len(routing_args) >= 2: + top_k_index, top_k_weights, *extra_args = routing_args + if isinstance(top_k_weights, DTensor): + top_k_weights = top_k_weights.to_local() + if not is_expert_parallel: + top_k_weights = _AllReduceBackward.apply(top_k_weights, tp_group) + routing_args = [top_k_index, top_k_weights, *extra_args] + + return (hidden_states, *routing_args), kwargs + + def install_forward(self, module, mesh, *, is_expert_parallel=False): + """Install the transforms but pass `is_expert_parallel` in the forward call.""" + original_forward = module.forward + output_source = ( + Partial() + if any( + isinstance(param, DTensor) and any(not placement.is_replicate() for placement in param.placements) + for param in module.parameters() + ) + else Replicate() + ) + + def tp_forward(*args, **kwargs): + args, kwargs = self.transform_inputs_pre_forward( + module, args, kwargs, mesh, is_expert_parallel=is_expert_parallel + ) + with self.context_around_forward(module, mesh): + output = original_forward(*args, **kwargs) + return self.transform_output_post_forward(module, output, mesh, source=output_source) + + module.forward = tp_forward + return module + + def transform_output_post_forward(self, module, output, mesh, source=None): + if output is None: + return None + + has_sharded_parameters = any( + isinstance(param, DTensor) and any(not placement.is_replicate() for placement in param.placements) + for param in module.parameters() + ) + if not has_sharded_parameters: + return output + + process_group = mesh.get_group() if mesh.ndim == 1 else mesh.get_group("tp") + return _AllReduceForward.apply(output, process_group) + + +class MoeIdentityParallel(TensorParallelLayer): + """ + Used in longcat_flash zero_experts (nn.Identity) which return the same value on every GPU, but the + moe_tp_experts will sums across GPUs. Therefore we pre-divide the identity input by tp_size to cancel the extra scaling. + """ + + def transform_inputs_pre_forward(self, module, args, kwargs, mesh): + input_tensor = args[0] + return (input_tensor / mesh.size(), *args[1:]), kwargs + + +class EpRouterParallel(TensorParallelLayer): + """Expert-parallel router: forward-only slicing of router outputs to local experts. + + Expects the router to return `(router_logits, router_scores, router_indices, *extra)`. + `router_logits` and any trailing `extra` outputs are passed through unchanged. + + The gate runs replicated on every rank and emits global expert IDs and scores. + Under EP each rank owns `num_experts // ep_size` experts, so this post-forward hook: + + - zeroes scores for non-local experts + - remaps surviving global indices to local indices (`fmod` after masking non-local slots) + - sets dropped slots to sentinel `num_local_experts` (skipped by grouped_gemm experts forward) + + Downstream `moe_tp_experts` allreduce-sums partial per-rank expert outputs. + + Example: 4 tokens, top_k=4, 128 experts, EP=8 → num_local_experts=16. + + Router output (identical on all ranks): + router_indices: + [ 52, 42, 119, 67], + [102, 89, 61, 40], + [ 82, 103, 4, 34], + [ 93, 23, 109, 11] + + Owning rank (`index // 16`): + [ 3, 2, 7, 4], + [ 6, 5, 3, 2], + [ 5, 6, 0, 2], + [ 5, 1, 6, 0], + + After slicing on rank 0 (owns experts 0-15): + router_indices: router_scores (illustrative): + [16, 16, 16, 16], [0.0, 0.0, 0.0, 0.0], + [16, 16, 16, 16], [0.0, 0.0, 0.0, 0.0], + [16, 16, 4, 16], [0.0, 0.0, 0.3, 0.0], + [16, 16, 16, 11], [0.0, 0.0, 0.0, 0.1], + + On rank 1 (owns experts 16-31), global expert 23 remaps to local index 7 via `23 % 16`. + + Scores and indices stay paired element-wise in `(seq, top_k)` shape. + """ + + def transform_output_post_forward(self, module, output, mesh): + ep_rank, ep_size = mesh.get_local_rank(), mesh.size() + num_experts = getattr(module, "num_experts", None) + if num_experts is None: + num_experts = getattr(getattr(module, "config", None), "num_experts", None) + if num_experts is None: + raise AttributeError( + f"Router module {type(module).__name__} is missing `num_experts` and `config.num_experts`" + ) + if num_experts % ep_size != 0: + raise ValueError(f"num_experts must be divisible by ep_size: {num_experts} % {ep_size} != 0") + num_local_experts = num_experts // ep_size + + router_logits, router_scores, router_indices, *extra_outputs = output + non_local_mask = (router_indices // num_local_experts) != ep_rank + router_scores = router_scores.masked_fill(non_local_mask, 0.0) + router_indices = router_indices.masked_fill(non_local_mask, -1) + 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, *extra_outputs + + +class RouterParallelMegaMoe(EpRouterParallel): + """Router TP plan used with DeepGEMM Mega MoE. + + Mega MoE handles EP dispatch inside the kernel and wants raw global expert ids + with unmasked routing weights, so the router doesn't pre-shard per EP rank like + ``EpRouterParallel`` does. + """ + + def transform_output_post_forward(self, module, output, mesh): + return output + + +class MoeTensorParalellMegaMoeExperts(MoeExpertsParallel): + """TP layer for DeepGEMM Mega MoE experts. + + Mega MoE is inference-only (the kernel has no backward) and handles EP dispatch + + combine + per-rank token sharding internally — so we skip the gradient-sync hooks + that ``MoeExpertsParallel`` would apply, and we forward the EP ``process_group`` + into the module so the symm-buffer rendezvous can run on first forward. + """ + + def transform_inputs_pre_forward(self, module, args, kwargs, mesh, *, is_expert_parallel=False): + hidden_states, top_k_index, top_k_weights = args[0], args[1], args[2] + return (hidden_states, top_k_index, top_k_weights, mesh.get_group()), kwargs + + def context_around_forward(self, module, mesh): + return _use_local_dtensor_params(module) + + def transform_output_post_forward(self, module, output, mesh, source=None): + return output + + +class ParallelInterface(GeneralInterface): + """Registry of named TP styles for the DTensor backend.""" + + _global_mapping = ( + { + "embedding_rowwise": RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()), + "colwise_gather_output": ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()), + "colwise": ColwiseParallel(input_layouts=Replicate(), output_layouts=Shard(-1)), + "rowwise": RowwiseParallel(input_layouts=Shard(-1), output_layouts=Replicate()), + "rowwise_split_input": RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()), + "packed_colwise": PackedColwiseParallel(), + "packed_rowwise": PackedRowwiseParallel(), + "sequence_parallel": SequenceParallel(use_local_output=True), + "grouped_gemm": MoEParamShard(Shard(0), shards_expert_dim=True), + "ep_router": EpRouterParallel(), + "megamoe_router": RouterParallelMegaMoe(), + "moe_tp_experts": MoeExpertsParallel(), + "megamoe_experts": MoeTensorParalellMegaMoeExperts(), + "moe_identity_expert": MoeIdentityParallel(), + "replicated_with_grad_allreduce": ReplicatedWithGradAllReduce(), + "mla_kv_a_proj": MlaKvAProjParallel(), + "all_reduce": AllReduceParallel(), + } + if is_torch_distributed_available() + else {} + ) + + +ALL_PARALLEL_STYLES: ParallelInterface = ParallelInterface() + + +def apply_tensor_parallelism(model, tp_mesh): + """DTensor backend: shard params as placeholders and install TP forward hooks.""" + + for name, module in model.named_modules(): + # Create DTensor placeholders so the loader knows which shard belongs to this rank. + for p_name, _ in list(module.named_parameters(recurse=False)): + full = f"{name}.{p_name}" if name else p_name + style_name = _get_parameter_tp_plan(parameter_name=full, tp_plan=model.tp_plan, is_weight=True) + if style_name is not None and style_name in ALL_PARALLEL_STYLES: + style = ALL_PARALLEL_STYLES[style_name] + style.validate_param(module, p_name, tp_mesh, parameter_name=full) + style.shard_param(module, p_name, tp_mesh) + + # Install the input/output transforms required by this module's TP style. + style_name = _get_parameter_tp_plan(parameter_name=name, tp_plan=model.tp_plan, is_weight=False) + if style_name is not None and style_name in ALL_PARALLEL_STYLES: + if style_name == "mla_kv_a_proj": + # MLA needs to know the qk_rope_head_dim to split the projection output into KV and RoPE parts. + # TODO: Store qk_rope_head_dim on MLA projection modules when the models initialize them. + module.config = model.config.get_text_config() + ALL_PARALLEL_STYLES[style_name].install_forward(module, tp_mesh) + module._is_hooked = True + + return model + + +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 TP-sharded ``DTensor`` parameters to full CPU tensors for checkpoint saving. + + Every rank must call this function so ``DTensor.full_tensor()`` collectives complete. + """ + gathered = {} + for key, tensor in state_dict.items(): + if isinstance(tensor, torch.Tensor): + if isinstance(tensor, DTensor): + tensor = tensor.full_tensor() + gathered[key] = tensor.detach().cpu().contiguous() + else: + gathered[key] = tensor + return gathered diff --git a/src/transformers/distributed/utils.py b/src/transformers/distributed/utils.py index 14f41f01d464..530c35284d94 100644 --- a/src/transformers/distributed/utils.py +++ b/src/transformers/distributed/utils.py @@ -14,12 +14,14 @@ from __future__ import annotations import os -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, TypeGuard from ..utils import is_torch_available, is_torch_distributed_available, is_torch_greater_or_equal if TYPE_CHECKING: + from torch.distributed.tensor import DTensor + from .configuration_utils import DistributedConfig @@ -33,7 +35,7 @@ def _is_torch_distributed_initialized() -> bool: return torch.distributed.is_initialized() -def is_dtensor(obj) -> bool: +def is_dtensor(obj: object) -> TypeGuard[DTensor]: if not is_torch_distributed_available(): return False from torch.distributed.tensor import DTensor @@ -113,6 +115,51 @@ def _distributed_barrier(): torch.distributed.barrier() +# TODO(3outeille): unify initialization across parallelism +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 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 {} + + 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"] + device_map = torch.device(f"{device_mesh.device_type}:{int(os.environ['LOCAL_RANK'])}") + + return device_map, device_mesh + + def initialize_fully_sharded_data_parallelism(distributed_config: DistributedConfig): # `fully_shard` itself only needs torch>=2.6, but distributed checkpoint save/load # (DCP + HuggingFaceStorageWriter) needs 2.7, so that is the effective requirement. diff --git a/src/transformers/integrations/__init__.py b/src/transformers/integrations/__init__.py index a867011ec469..1cf5af8c3359 100755 --- a/src/transformers/integrations/__init__.py +++ b/src/transformers/integrations/__init__.py @@ -140,8 +140,6 @@ "mxfp4": [ "Mxfp4GptOssExperts", "convert_moe_packed_tensors", - "dequantize", - "load_and_swizzle_mxfp4", "quantize_to_mxfp4", "replace_with_mxfp4_linear", "swizzle_mxfp4", @@ -171,9 +169,8 @@ ] _import_structure["tensor_parallel"] = [ - "shard_and_distribute_module", "ALL_PARALLEL_STYLES", - "translate_to_torch_parallel_style", + "shard_and_distribute_module", ] _import_structure["flex_attention"] = [ "make_flex_block_causal_mask", @@ -297,8 +294,6 @@ ) from .mxfp4 import ( Mxfp4GptOssExperts, - dequantize, - load_and_swizzle_mxfp4, quantize_to_mxfp4, replace_with_mxfp4_linear, swizzle_mxfp4, @@ -320,11 +315,7 @@ from .executorch import TorchExportableModuleWithStaticCache, convert_and_export_with_cache from .flex_attention import make_flex_block_causal_mask - from .tensor_parallel import ( - ALL_PARALLEL_STYLES, - shard_and_distribute_module, - translate_to_torch_parallel_style, - ) + from .tensor_parallel import ALL_PARALLEL_STYLES, shard_and_distribute_module else: import sys diff --git a/src/transformers/integrations/deepgemm.py b/src/transformers/integrations/deepgemm.py index 20a482aa6523..ab266873f821 100644 --- a/src/transformers/integrations/deepgemm.py +++ b/src/transformers/integrations/deepgemm.py @@ -44,7 +44,6 @@ resolve_internal_import, ) from .hub_kernels import _MISSING_KERNELS_MESSAGE, lazy_load_kernel -from .tensor_parallel import to_local logger = logging.get_logger(__name__) @@ -650,10 +649,10 @@ def deepgemm_bf16_experts_forward( hidden_states, top_k_index, top_k_weights, self.num_experts, deepgemm.m_alignment, is_sm100() ) - weight_up = to_local(self.gate_up_proj if self.has_gate else self.up_proj) - weight_down = to_local(self.down_proj) - up_bias = to_local(self.gate_up_proj_bias if self.has_gate else self.up_proj_bias) if self.has_bias else None - down_bias = to_local(self.down_proj_bias) if self.has_bias else None + weight_up = self.gate_up_proj if self.has_gate else self.up_proj + weight_down = self.down_proj + up_bias = (self.gate_up_proj_bias if self.has_gate else self.up_proj_bias) if self.has_bias else None + down_bias = self.down_proj_bias if self.has_bias else None # Up projection. up_out_dim = weight_up.shape[-1] if self.is_transposed else weight_up.shape[1] @@ -717,10 +716,10 @@ def deepgemm_fp8_fp4_experts_forward( num_tokens = hidden_states.size(0) hidden_dim = hidden_states.size(-1) - weight_up = to_local(self.gate_up_proj if self.has_gate else self.up_proj) - weight_scale_up = to_local(self.gate_up_proj_scale_inv if self.has_gate else self.up_proj_scale_inv) - weight_down = to_local(self.down_proj) - weight_scale_down = to_local(self.down_proj_scale_inv) + weight_up = self.gate_up_proj if self.has_gate else self.up_proj + weight_scale_up = self.gate_up_proj_scale_inv if self.has_gate else self.up_proj_scale_inv + weight_down = self.down_proj + weight_scale_down = self.down_proj_scale_inv cast_kwargs = _select_fp8_cast_kwargs(weight_up, weight_scale_up, self.block_size, is_sm100()) ( @@ -795,11 +794,11 @@ def setup_megamoe_weights(module: torch.nn.Module) -> None: side Parameters — the kernel takes raw pointers. """ deepgemm = load_deepgemm_kernel() - gate_up_sf_raw = to_local(module.gate_up_proj_scale_inv.data) - down_sf_raw = to_local(module.down_proj_scale_inv.data) + gate_up_sf_raw = module.gate_up_proj_scale_inv.data + down_sf_raw = module.down_proj_scale_inv.data # Force int8 view: the kernel's interleave reshape/empty_like/copy_ is bit-level. - gate_up_w = to_local(module.gate_up_proj.data).view(torch.int8).contiguous() - down_w = to_local(module.down_proj.data).view(torch.int8).contiguous() + gate_up_w = module.gate_up_proj.data.view(torch.int8).contiguous() + down_w = module.down_proj.data.view(torch.int8).contiguous() intermediate_hidden = module.intermediate_dim num_local_experts = module.num_experts diff --git a/src/transformers/integrations/fbgemm_fp8.py b/src/transformers/integrations/fbgemm_fp8.py index a987d362b7df..21bb5d576549 100644 --- a/src/transformers/integrations/fbgemm_fp8.py +++ b/src/transformers/integrations/fbgemm_fp8.py @@ -267,7 +267,7 @@ def get_quantize_fp8_per_row(): def replace_with_fbgemm_fp8_linear( - model, modules_to_not_convert: list[str] | None = None, quantization_config=None, pre_quantized=False, tp_plan=None + model, modules_to_not_convert: list[str] | None = None, quantization_config=None, pre_quantized=False ): """ A helper function to replace all `torch.nn.Linear` modules by `FbgemmFp8Linear` modules. @@ -297,9 +297,6 @@ def replace_with_fbgemm_fp8_linear( with init_empty_weights(include_buffers=True): if module.__class__.__name__ == "Llama4TextExperts": # TODO: make sure tp works later - # if tp_plan is not None: - # tp_key = re.sub(r"\d+", "*", f"{module_name}.down_proj_scale") - # tp_plan[tp_key] = None text_config = getattr(model.config, "text_config", model.config) new_module = FbgemmFp8Llama4TextExperts(text_config or model.config) elif isinstance(module, nn.Linear): diff --git a/src/transformers/integrations/finegrained_fp8.py b/src/transformers/integrations/finegrained_fp8.py index 8f82c9b15196..910a8bd82ac2 100644 --- a/src/transformers/integrations/finegrained_fp8.py +++ b/src/transformers/integrations/finegrained_fp8.py @@ -37,7 +37,6 @@ ) from .hub_kernels import _MISSING_KERNELS_MESSAGE, lazy_load_kernel from .moe import ExpertsInterface, use_experts_implementation -from .tensor_parallel import to_local logger = logging.get_logger(__name__) @@ -336,13 +335,10 @@ def forward(self, input: torch.Tensor) -> torch.Tensor: if self.weight.element_size() > 1: return F.linear(input, self.weight, self.bias) - weight = to_local(self.weight) - scale_inv = to_local(self.weight_scale_inv) - return fp8_linear( input, - weight, - scale_inv, + self.weight, + self.weight_scale_inv, block_size=self.block_size, activation_scale=self.activation_scale, bias=self.bias, @@ -395,8 +391,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: y.add_(self.bias.view(self.n_groups, -1)) return y - w = to_local(self.weight) - scale_inv = to_local(self.weight_scale_inv) + w = self.weight + scale_inv = self.weight_scale_inv w = w.view(self.n_groups, -1, hidden_dim) x = x.movedim(-2, 0).reshape(-1, hidden_dim) @@ -450,10 +446,10 @@ def fp8_batched_mm_experts_forward( # zeroes them before the per-token reduction so `uninit * 0 = NaN` can't poison the sum. sentinel_mask = (expert_ids >= self.num_experts).unsqueeze(-1) - weight_up = to_local(self.gate_up_proj if self.has_gate else self.up_proj) - weight_scale_up = to_local(self.gate_up_proj_scale_inv if self.has_gate else self.up_proj_scale_inv) - weight_down = to_local(self.down_proj) - weight_scale_down = to_local(self.down_proj_scale_inv) + weight_up = self.gate_up_proj if self.has_gate else self.up_proj + weight_scale_up = self.gate_up_proj_scale_inv if self.has_gate else self.up_proj_scale_inv + weight_down = self.down_proj + weight_scale_down = self.down_proj_scale_inv # --- Up projection per expert (FP8 batched) --- proj_out = finegrained_fp8.batched_matmul( @@ -538,10 +534,10 @@ def fp8_grouped_mm_experts_forward( # quantized weights are inference-only, so no bwd pre-mask is needed. sentinel_mask = (expert_ids_g >= self.num_experts).unsqueeze(-1) - weight_up = to_local(self.gate_up_proj if self.has_gate else self.up_proj) - weight_scale_up = to_local(self.gate_up_proj_scale_inv if self.has_gate else self.up_proj_scale_inv) - weight_down = to_local(self.down_proj) - weight_scale_down = to_local(self.down_proj_scale_inv) + weight_up = self.gate_up_proj if self.has_gate else self.up_proj + weight_scale_up = self.gate_up_proj_scale_inv if self.has_gate else self.up_proj_scale_inv + weight_down = self.down_proj + weight_scale_down = self.down_proj_scale_inv # --- Up projection per expert (FP8 grouped) --- proj_out = finegrained_fp8.grouped_matmul( @@ -870,6 +866,8 @@ def replace_with_fp8_linear( has_bias=module.bias is not None, ) if new_module is not None: + # TP must use local tensors because this quantization path does not support DTensor inputs or weights. + new_module._hf_quantized_needs_local_tp = True model.set_submodule(module_name, new_module) has_been_replaced = True diff --git a/src/transformers/integrations/mxfp4.py b/src/transformers/integrations/mxfp4.py index 7121d08fdf5b..c15a8b844353 100644 --- a/src/transformers/integrations/mxfp4.py +++ b/src/transformers/integrations/mxfp4.py @@ -491,115 +491,11 @@ def mlp_forward(self, hidden_states): return routed_out, router_logits -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") - - 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) - if hasattr(module, blocks_attr) and hasattr(module, scales_attr): - dequantized = convert_moe_packed_tensors(getattr(module, blocks_attr), getattr(module, scales_attr)) - setattr(module, proj, torch.nn.Parameter(dequantized.to(target_device))) - delattr(module, blocks_attr) - delattr(module, scales_attr) - - def dequantize_convertops(blocks, scales): dequantized = convert_moe_packed_tensors(blocks, scales) return torch.nn.Parameter(dequantized) -def load_and_swizzle_mxfp4(module, param_name, param_value, target_device, triton_kernels_hub, **kwargs): - """ - This transforms the weights obtained using `convert_gpt_oss.py` to load them into `Mxfp4GptOssExperts`. - """ - PrecisionConfig, FlexCtx, InFlexData = ( - triton_kernels_hub.matmul_ogs.PrecisionConfig, - 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 "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)) - blocks_attr = f"{proj}_blocks" - scales_attr = f"{proj}_scales" - blocks = getattr(module, blocks_attr) # at this point values were loaded from ckpt - scales = getattr(module, scales_attr) - # Check if both blocks and scales both not on meta device - if blocks.device.type != "meta" and scales.device.type != "meta": - local_experts = blocks.size(0) - if proj == "gate_up_proj": - blocks = blocks.reshape(local_experts, module.intermediate_size * 2, -1) - else: - blocks = blocks.reshape(local_experts, -1, module.intermediate_size // 2) - if ( - getattr(target_device, "type", target_device) == "cpu" - and hasattr(torch, "accelerator") - and torch.accelerator.current_accelerator() is not None - ): - target_device = torch.accelerator.current_accelerator().type - blocks = blocks.to(target_device).contiguous() - scales = scales.to(target_device).contiguous() - with on_device(target_device): - triton_weight_tensor, weight_scale = swizzle_mxfp4( - blocks.transpose(-2, -1), scales.transpose(-2, -1), triton_kernels_hub - ) - - # need to overwrite the shapes for the kernels - if proj == "gate_up_proj": - triton_weight_tensor.shape = torch.Size([local_experts, module.hidden_size, module.intermediate_size * 2]) - else: - triton_weight_tensor.shape = torch.Size([local_experts, module.intermediate_size, module.hidden_size]) - - # triton_weight_tensor is what needs to be passed in oai kernels. It stores the data, the shapes and any more objects. It is like a subtensor - setattr(module, proj, triton_weight_tensor) - setattr( - module, - f"{proj}_precision_config", - PrecisionConfig(weight_scale=weight_scale, flex_ctx=FlexCtx(rhs_data=InFlexData())), - ) - - # delete blocks and scales - delattr(module, scales_attr) - delattr(module, blocks_attr) - del blocks - - def swizzle_mxfp4_convertops(blocks, scales, module, proj, target_device, triton_kernels_hub): """ This transforms the weights obtained using `convert_gpt_oss.py` to load them into `Mxfp4GptOssExperts`. diff --git a/src/transformers/integrations/sonicmoe.py b/src/transformers/integrations/sonicmoe.py index 795bb820c216..923bf09caed4 100644 --- a/src/transformers/integrations/sonicmoe.py +++ b/src/transformers/integrations/sonicmoe.py @@ -30,7 +30,6 @@ from ..utils import logging from ..utils.import_utils import is_kernels_available, maybe_import_error from .hub_kernels import lazy_load_kernel -from .tensor_parallel import to_local logger = logging.get_logger(__name__) @@ -184,10 +183,10 @@ def sonicmoe_experts_forward( # already zero (RouterParallel masks them at dispatch), so the per-token reduction # contributes nothing for sentinel slots. - w1 = to_local(self.gate_up_proj) - w2 = to_local(self.down_proj) - b1 = to_local(self.gate_up_proj_bias) if self.has_bias else None - b2 = to_local(self.down_proj_bias) if self.has_bias else None + w1 = self.gate_up_proj + w2 = self.down_proj + b1 = self.gate_up_proj_bias if self.has_bias else None + b2 = self.down_proj_bias if self.has_bias else None # Map activation function act_name = getattr(self.config, "hidden_act", "silu").lower() diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index 8eca2f441d4d..9618f83ffd3c 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -1,4 +1,4 @@ -# Copyright 2024 The HuggingFace Team. All rights reserved. +# Copyright 2026 The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -11,1610 +11,73 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from __future__ import annotations - -import math -import operator -import os -import re -from functools import reduce - -from ..distributed import DistributedConfig -from ..distributed.utils import is_dtensor -from ..utils import logging -from ..utils.generic import GeneralInterface -from ..utils.import_utils import is_torch_available, is_torch_distributed_available - - -if is_torch_available(): - import torch - import torch.distributed as dist - from torch import nn - - -logger = logging.get_logger(__name__) - - -def to_local(t): - """Unwrap a `DTensor` to its local shard if needed; pass through otherwise. - - Custom kernels (CUTLASS, CuteDSL, Triton) take raw tensor pointers and don't - understand `DTensor`, so weights wrapped by FSDP2 / EP need this unwrap before - they can be fed to the kernel. ``to_local()`` is autograd-aware on the train - path: backward rewraps the gradient as a DTensor matching each parameter's - placements. - """ - if is_dtensor(t): - return t.to_local() - return t - - -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: - # 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 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 {} - - 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"] - device_map = torch.device(f"{device_mesh.device_type}:{int(os.environ['LOCAL_RANK'])}") - - return device_map, device_mesh - - -def replace_layer_number_by_wildcard(name: str) -> str: - """ - Replace the numbers in the `name` by wildcards, only if they are in-between dots (`.`) or if they are between - a dot (`.`) and the end of the string. - This matches how modules are named/numbered when using a nn.ModuleList or nn.Sequential, but will NOT match - numbers in a parameter name itself, e.g. if the param is named `"w1"` or `"w2"`. - """ - return re.sub(r"\.\d+(\.|$)", lambda m: ".*" + m.group(1), name) - - -def _get_parameter_tp_plan(parameter_name: str, tp_plan: dict[str, str], is_weight=True) -> str | None: - """ - Get the TP style for a parameter from the TP plan. - - The TP plan is a dictionary that maps parameter names to TP styles. - The parameter name can be a generic name with wildcards (e.g. "*.weight") or a specific name (e.g. "layer_1.weight"). - - The `is_weight` is important because for weights, we want to support `.weights` and `.bias` cases seamlessly! but - not parent classes for `post_init` calls - """ - generic_param_name = replace_layer_number_by_wildcard(parameter_name) - if generic_param_name in tp_plan: - return tp_plan[generic_param_name] - elif is_weight and "." in generic_param_name and (module_name := generic_param_name.rsplit(".", 1)[0]) in tp_plan: - return tp_plan[module_name] - return None - - -# ============================================================================= -# Tensor Sharding Utilities -# ============================================================================= - - -if is_torch_available(): - str_to_dtype = { - "BOOL": torch.bool, - "U8": torch.uint8, - "I8": torch.int8, - "I16": torch.int16, - "F16": torch.float16, - "BF16": torch.bfloat16, - "I32": torch.int32, - "F32": torch.float32, - "F64": torch.float64, - "I64": torch.int64, - "F8_E4M3": torch.float8_e4m3fn, - } - - -def _blocks_to_block_sizes(total_size: int, blocks: int | list[int]) -> list[int]: - """ - Convert block count or proportions to block sizes. - - This function accepts - - - The number of blocks (int), in which case the block size is - total_size//blocks; or - - A list of block sizes (list[int]). - - In the second case, if sum(blocks) < total_size, the ratios between - the block sizes will be preserved. For instance, if blocks is - [2, 1, 1] and total_size is 1024, the returned block sizes are - [512, 256, 256]. - """ - if isinstance(blocks, list): - total_blocks = sum(blocks) - assert total_size % total_blocks == 0, f"Cannot split {total_size} in proportional blocks: {blocks}" - part_size = total_size // total_blocks - return [part_size * block for block in blocks] - else: - assert total_size % blocks == 0, f"Prepacked is not divisible by {blocks}" - single_size = total_size // blocks - return [single_size] * blocks - - -def get_packed_weights(param, empty_param, device_mesh, rank, dim): - """ - When weights are packed (gate_up_proj), we need to make sure each shard gets its correct share. - So if you have: gate_proj ( 16, 5120, 8190) - and up_proj ( 16, 5120, 8190) - packed as gate_up_proj ( 16, 5120, 2 * 8190) - And you shard along the last dimension, you need to interleave the gate and up values: - - Now, if we shard along the last dimension across TP_size (Tensor Parallelism size), we must interleave the values from gate and up projections correctly. - - Let's take TP_size = 4 for an example: - - Packed tensor `gate_up_proj` - --------------------------------------------------------------- - [ G0 G1 G2 G3 | G4 G5 G6 G7 | ... | U0 U1 U2 U3 | U4 U5 U6 U7 | ... ] - ↑─────────────↑ ↑─────────────↑ ↑─────────────↑ ↑─────────────↑ - Gate Slice 0 Gate Slice 1 Up Slice 0 Up Slice 1 - - Explanation: - - The first half of the tensor (left of the center) holds the gate_proj values. - - The second half (right of the center) holds the up_proj values. - - For TP=4, we divide each half into 4 slices. In this example, we show two slices for brevity. - - Each shard receives one slice from the gate part and the corresponding slice from the up part. - - For instance: - • Shard 0 gets: [ Gate Slice 0, Up Slice 0 ] = [ G0, G1, G2, G3, U0, U1, U2, U3 ] - • Shard 1 gets: [ Gate Slice 1, Up Slice 1 ] = [ G4, G5, G6, G7, U4, U5, U6, U7 ] - • … and so on. - - This ensures that each shard receives an equal portion of both gate and up projections, maintaining consistency across tensor parallelism. - """ - slice_ = param - total_size = empty_param.shape[dim] - world_size = device_mesh.size() - block_sizes = _blocks_to_block_sizes(total_size=total_size, blocks=2) - - tensors_slices = [] - block_offset = 0 - for block_size in block_sizes: - shard_block_size = block_size // world_size - start = rank * shard_block_size - stop = (rank + 1) * shard_block_size - tensors_slices += range(block_offset + start, block_offset + stop) - block_offset += block_size - - slice_dtype = slice_.get_dtype() - # Handle F8_E4M3 dtype by converting to float16 before slicing - # Without upcasting, the slicing causes : RuntimeError: "index_cpu" not implemented for 'Float8_e4m3fn' - casted = False - if slice_dtype == "F8_E4M3" or slice_dtype == "F8_E5M2": - slice_ = slice_[...].to(torch.float16) - casted = True - - if dim == 0: - tensor = slice_[tensors_slices, ...] - elif dim == 1 or dim == -2: - tensor = slice_[:, tensors_slices, ...] - elif dim == 2 or dim == -1: - tensor = slice_[..., tensors_slices] - else: - raise ValueError(f"Unsupported dim {dim}, only dim 0, 1 or 2 are supported") - - if casted: - return tensor - else: - return tensor.to(str_to_dtype[slice_dtype]) - - -def repack_weights( - packed_parameter: torch.Tensor, - sharded_dim: int, # The dimension index in the global tensor that was sharded - world_size: int, - num_blocks: int = 2, -) -> torch.Tensor: - """ - Reorders a tensor that was reconstructed from sharded packed weights into its canonical packed format. - - For example, if a weight was packed (e.g., gate_proj and up_proj) and then sharded, - DTensor.full_tensor() might produce an interleaved layout like [G0, U0, G1, U1, ...] - along the sharded dimension. This function reorders it to [G0, G1, ..., U0, U1, ...]. - This is an inverse operation to get_packed_weights. - - Args: - reconstructed_tensor: The tensor reconstructed from DTensor (e.g., via .full_tensor().contiguous()). - sharded_dim: The dimension index in the reconstructed_tensor that was originally sharded. - world_size: The tensor parallel world size. - num_packed_projs: The number of projections that were packed together (e.g., 2 for gate_up_proj). - - Returns: - The reordered tensor in canonical packed format. - """ - - if num_blocks != 2: - raise ValueError( - "Num blocks different from 2 is not supported yet. This is most likely a bug in your implementation as we only pack gate and up projections together." - ) - - actual_sharded_dim = sharded_dim if sharded_dim >= 0 else sharded_dim + packed_parameter.ndim - total_size_on_sharded_dim = packed_parameter.shape[actual_sharded_dim] - original_block_size_on_dim = total_size_on_sharded_dim // num_blocks - shard_chunk_size = original_block_size_on_dim // world_size - - prefix_shape = packed_parameter.shape[:actual_sharded_dim] - suffix_shape = packed_parameter.shape[actual_sharded_dim + 1 :] - - tensor_view = packed_parameter.view( - *prefix_shape, - world_size, - num_blocks, - shard_chunk_size, - *suffix_shape, - ) - - # Permute to bring num_packed_projs first, then world_size, then shard_chunk_size - # This groups all chunks of G together, then all chunks of U together. - # Target order of these middle dimensions: (num_packed_projs, world_size, shard_chunk_size) - # Current order of view's middle dimensions: (world_size, num_packed_projs, shard_chunk_size) - # Absolute indices of the dimensions to be permuted (world_size, num_packed_projs) - axis_ws_abs = len(prefix_shape) - axis_npp_abs = len(prefix_shape) + 1 - - permute_order = list(range(tensor_view.ndim)) - permute_order[axis_ws_abs], permute_order[axis_npp_abs] = permute_order[axis_npp_abs], permute_order[axis_ws_abs] - - tensor_permuted = tensor_view.permute(*permute_order) - - # Reshape back to the original tensor's ndim, with the sharded dimension now correctly ordered as [G_all, U_all]. - # The final shape should be the same as reconstructed_tensor. - final_ordered_tensor = tensor_permuted.reshape_as(packed_parameter) - - return final_ordered_tensor - - -def get_tensor_shard(param, empty_param, device_mesh, rank, dim, tensor_idx: int | None = None): - """ - Generalized tensor sharding across a multi-dimensional device mesh. - Extract only the fraction of the parameter owned by the given `rank` when the parameter would have gone sharding at provided `dim`. - Extraction follows the pytorch `Shard` placement so that sharding and materializing back to full tensor follows `Shard` semantics. - `Shard` follows torch.chunk style sharding of the tensor. We demonstrate some cases below on how sharding happens including some edge cases - such as some ranks having an empty tensor as shard. Below implementation is robust 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 validate_module(self, module: nn.Module, device_mesh, layer_name: str = ""): - """Raise if the module cannot be sharded with this style on the given mesh.""" - pass - - 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 validate_module(self, module: nn.Module, device_mesh, layer_name: str = ""): - out_features = getattr(module, "out_features", None) - if self.gather_output and out_features is not None and out_features % device_mesh.size() != 0: - raise ValueError( - f"`{layer_name}` ({type(module).__name__} with out_features={out_features}) is sharded with " - f"'colwise_gather_output', which requires out_features to be divisible by the number of ranks " - f"({device_mesh.size()}) to all-gather equal-size shards. Resize the weight (e.g. " - f"`model.resize_token_embeddings` for LM heads) or override this module's entry in the tp_plan." - ) - - 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. +"""Backward-compatibility shim for the tensor parallel API. - 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 AllReduceParallel(TensorParallelLayer): - """All-reduce a module's forward output across the TP mesh. Use as a declarative - sync point at the boundary of a multi-arg module whose compute ends in a partial - sum (e.g. the lightning indexer's score sum before its top-k). - """ - - def _prepare_input_fn(self, mod, inputs, device_mesh): - return inputs - - def _prepare_output_fn(self, mod, outputs, device_mesh): - return all_reduce_forward(outputs, device_mesh) - - 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, output_fn=self._prepare_output_fn) - - -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 = input_mask.unsqueeze(-1) - outputs = outputs * (~mask).to(outputs.dtype) - del mod._input_mask - - return all_reduce_forward(outputs, device_mesh) - - def shard_tensor( - self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None - ) -> torch.Tensor: - # If only 1 dim, shard this one (usually it's a `bias`) - dim = param.dim() if isinstance(param, torch.Tensor) else len(param.get_shape()) - if dim == 1: - parameter = get_tensor_shard(param, self.empty_param, self.device_mesh, self.rank, -1) - else: - parameter = get_tensor_shard( - param, - self.empty_param, - self.device_mesh, - self.rank, - self.embedding_dim_sharding, - ) - return parameter.to(device=device, dtype=dtype) - - def get_expected_sharded_shape(self, full_shape: tuple[int, ...] | torch.Size) -> tuple[int, ...]: - world_size = self.device_mesh.size() - shape = list(full_shape) - # EmbeddingParallel shards on self.embedding_dim_sharding (default 0) - # 1D tensors (bias) shard on dim -1 - dim = -1 if len(shape) == 1 else self.embedding_dim_sharding - dim = len(shape) + dim if dim < 0 else dim - shard_size = math.ceil(shape[dim] / world_size) - start = self.rank * shard_size - end = min(start + shard_size, shape[dim]) - shape[dim] = end - start - return tuple(shape) - - def update_module_attributes(self, module: nn.Module): - if hasattr(module, "num_embeddings") and self.embedding_dim_sharding == 0: - module.num_embeddings = self.get_expected_sharded_shape((module.num_embeddings,))[0] - if hasattr(module, "embedding_dim") and self.embedding_dim_sharding == 1: - module.embedding_dim = self.get_expected_sharded_shape((module.embedding_dim,))[0] - - -class SequenceParallel(TensorParallelLayer): - """ - Sequence Parallel: input/output sharded on sequence dimension. - Weights are replicated. - """ - - def __init__(self, sequence_dim: int = 1, use_local_output: bool = False, use_dtensor=False, **kwargs): - super().__init__(**kwargs) - self.sequence_dim = sequence_dim - - def _prepare_input_fn(self, mod, inputs, device_mesh): - input_tensor = inputs[0] if inputs else inputs - # For sequence parallel, input is sharded on sequence dim - # All-gather for the layer, then reduce-scatter after - return all_gather(input_tensor, device_mesh) - - def _prepare_output_fn(self, mod, outputs, device_mesh): - return reduce_scatter(outputs, device_mesh) - - def shard_tensor( - self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None - ) -> torch.Tensor: - return param[...].to(device=device, dtype=dtype) - - -class GroupedGemmParallel(TensorParallelLayer): - """ - Applies Expert Parallelism to MoE experts by loading the correct experts on each device. - """ - - def __init__(self, **kwargs): - super().__init__(**kwargs) - - def shard_tensor( - self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None - ) -> torch.Tensor: - global_num_experts = self.empty_param.shape[0] - if global_num_experts % self.device_mesh.size() != 0: - raise ValueError( - f"Global number of experts must be divisible by number of devices: {global_num_experts} % {self.device_mesh.size()} != 0" - ) - local_num_experts = global_num_experts // self.device_mesh.size() - shard_size = local_num_experts - start = self.rank * shard_size - end = (self.rank + 1) * shard_size - # special case we don't "shard" just send this entire tensor to the correct rank. - shape = param.get_shape() if not isinstance(param, torch.Tensor) else param.shape - if tensor_idx is not None and start <= tensor_idx < end: - # this tensor does need to be materialized on this device: - return param[:].to(device=device) - elif tensor_idx is None: # a bias or a weight, but already merged - return param[start:end].to(device=device, dtype=dtype) - elif len(shape) >= 1 and tensor_idx is not None: - return None - else: # bias case - return param[:].to(device=device, dtype=dtype) - - def get_expected_sharded_shape(self, full_shape: tuple[int, ...] | torch.Size) -> tuple[int, ...]: - # GroupedGemm shards on dim 0 (experts dimension) - world_size = self.device_mesh.size() - shape = list(full_shape) - local_num_experts = shape[0] // world_size - shape[0] = local_num_experts - return tuple(shape) - - def update_module_attributes(self, module: nn.Module): - if hasattr(module, "num_experts"): - module.num_experts = self.get_expected_sharded_shape((self.empty_param.shape[0],))[0] - - -class RouterParallel(TensorParallelLayer): - """ - Allows to reshape the router scores to support running expert parallel. - """ - - def __init__(self, **kwargs): - super().__init__(**kwargs) - - def _prepare_input_fn(self, mod, inputs, device_mesh): - return inputs - - def _prepare_output_fn(self, mod, outputs, device_mesh): - """ - Remap global expert indices to local and zero out non-local scores. - - Example: 4 tokens, top_k=4, 128 experts, EP=8. num_local_experts = 128/8 = 16. - - Router produces (all ranks see the same values): - router_scores: (4, 4) — top-k routing weights - router_indices: (4, 4) — global expert IDs - [ 52, 42, 119, 67], - [102, 89, 61, 40], - [ 82, 103, 4, 34], - [ 93, 23, 109, 11], - - Each index maps to a rank: index // 16 gives the owning rank. - [3, 2, 7, 4], - [6, 5, 3, 2], - [5, 6, 0, 2], - [5, 1, 6, 0], - - For rank 0 (owns experts 0-15), we remap local indices with fmod and - fill non-local with sentinel=16 (used for one_hot masking): - router_indices (rank 0): - [ 16, 16, 16, 16], - [ 16, 16, 16, 16], - [ 16, 16, 4, 16], - [ 16, 16, 16, 11], - - Scores for non-local experts are zeroed out via masked_fill: - router_scores (rank 0): - [0.0, 0.0, 0.0, 0.0], - [0.0, 0.0, 0.0, 0.0], - [0.0, 0.0, 0.3, 0.0], ← only expert 4 (local) keeps its score - [0.0, 0.0, 0.0, 0.1], ← only expert 11 (local) keeps its score - - both router_scores and router_indices stay (seq, top_k) shape. - They are paired element-wise: scores[i] is the weight for indices[i]. - All expert forward implementations (grouped_mm, batched_mm, eager) flatten - both with reshape(-1) and rely on this pairing. Changing the shape of one - without the other breaks routing! - - Each rank believes it is alone and computes only its part of the hidden states. - The sentinel index (num_local_experts) is skipped by one_hot encoding or clamped - + masked in grouped_mm/batched_mm. After the expert forward, an all_reduce sums - partial outputs across EP ranks to produce the full result. - """ - ep_rank, ep_size = device_mesh.get_local_rank(), device_mesh.size() - num_experts = getattr(mod, "num_experts", None) - if num_experts is None: - num_experts = getattr(getattr(mod, "config", None), "num_experts", None) - if num_experts is None: - raise AttributeError(f"Router module {type(mod).__name__} is missing num_experts and config.num_experts") - - if num_experts % ep_size != 0: - raise ValueError( - f"The number of experts must be divisible by number of ep_size: {num_experts} % {ep_size} != 0" - ) - num_local_experts = num_experts // ep_size - # Some routers return extra tensors after the standard logits/scores/indices, e.g. zaya's router state. - router_logits, router_scores, router_indices, *extra_outputs = outputs - non_local_mask = (router_indices // num_local_experts) != ep_rank - router_scores = router_scores.masked_fill(non_local_mask, 0.0) - router_indices = router_indices.masked_fill(non_local_mask, -1) - # As -1 % 1 is 0, we can only use mask fill when num_local_experts is 1 - if num_local_experts > 1: - router_indices = torch.fmod(router_indices, num_local_experts) - else: - router_indices = router_indices.masked_fill(router_indices > 0, 0).masked_fill(router_indices < 0, -1) - router_indices = router_indices.masked_fill(router_indices == -1, num_local_experts) - return router_logits, router_scores, router_indices, *extra_outputs - - 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 RouterParallelMegaMoe(RouterParallel): - """Router TP plan used with DeepGEMM Mega MoE. - - Mega MoE handles EP dispatch inside the kernel and wants raw global expert ids - with unmasked routing weights, so the router doesn't pre-shard per EP rank like - `RouterParallel._prepare_output_fn` does. The quantizer's `update_tp_plan` swaps - `"ep_router"` → `"megamoe_router"` when `experts_implementation == "deepgemm_megamoe"`. - """ - - def _prepare_output_fn(self, mod, outputs, device_mesh): - return outputs - - -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 MoeTensorParalellMegaMoeExperts(TensorParallelLayer): - """TP layer for DeepGEMM Mega MoE experts. - - Mega MoE is inference-only (the kernel has no backward) and handles EP dispatch + - combine + per-rank token sharding internally — so we skip the gradient-sync hooks - that the regular `MoeTensorParalellExperts` would apply, and we forward the EP - `process_group` into the module so the symm-buffer rendezvous can run on first - forward. The quantizer's `update_tp_plan` swaps the experts plan key from - `"moe_tp_experts"` to `"megamoe_experts"` when - `from_pretrained(..., experts_implementation="deepgemm_megamoe")`. - """ - - def _prepare_input_fn(self, mod, inputs, device_mesh): - hidden_states, top_k_index, top_k_weights = inputs[0], inputs[1], inputs[2] - return hidden_states, top_k_index, top_k_weights, device_mesh.get_group() - - def _prepare_output_fn(self, mod, outputs, device_mesh): - # Kernel returned the fully-combined gathered output; no further reduction. - return outputs - - 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 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) +The canonical implementation lives in ``transformers.distributed.tensor_parallel``. +""" +from __future__ import annotations -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(), - "megamoe_router": RouterParallelMegaMoe(), - "moe_tp_experts": MoeTensorParalellExperts(), - "megamoe_experts": MoeTensorParalellMegaMoeExperts(), - "moe_identity_expert": MoeIdentityExpertParallel(), - "replicated_with_grad_allreduce": ReplicatedWithGradAllReduce(), - "mla_kv_a_proj": MlaKvAProjParallel(), - "all_reduce": AllReduceParallel(), - } - if is_torch_distributed_available() - else {} +import warnings + +from ..distributed.tensor_parallel import ( + ALL_PARALLEL_STYLES, + AllReduceParallel, + ColwiseParallel, + EpRouterParallel, + MlaKvAProjParallel, + MoeExpertsParallel, + MoeIdentityParallel, + MoEParamShard, + MoeTensorParalellMegaMoeExperts, + PackedColwiseParallel, + PackedRowwiseParallel, + ParallelInterface, + ReplicatedWithGradAllReduce, + RouterParallelMegaMoe, + RowwiseParallel, + SequenceParallel, + TensorParallelLayer, + apply_tensor_parallelism, + gather_state_dict_for_save, + replace_layer_number_by_wildcard, + verify_tp_plan, +) + + +def shard_and_distribute_module(*args, **kwargs): + """Deprecated per-parameter sharding helper from the legacy TP loading path.""" + warnings.warn( + "`shard_and_distribute_module` is deprecated and unavailable with the DTensor tensor-parallel " + "loading path. Use `transformers.distributed.tensor_parallel.apply_tensor_parallelism` with " + "`from_pretrained(..., tp_plan=...)` instead.", + FutureWarning, + stacklevel=2, ) - - # 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, - "all_reduce": 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, - "all_reduce": None, - } - - @classmethod - def register_plan_to_weight_dim(cls, key: str, value: int | None): - cls.plan_to_weight_dim[key] = value - - @classmethod - def register_plan_to_bias_dim(cls, key: str, value: int | None): - cls.plan_to_bias_dim[key] = value - - -ALL_PARALLEL_STYLES: ParallelInterface = ParallelInterface() - - -# ============================================================================= -# High-Level API Functions -# ============================================================================= - - -def gather_full_tensor( - local_tensor: torch.Tensor, shard_dim: int, device_mesh: dist.device_mesh.DeviceMesh -) -> torch.Tensor: - """ - All-gather a sharded tensor along the specified dimension to reconstruct the full tensor. - - Args: - local_tensor: The local shard of the tensor on this rank - shard_dim: The dimension along which the tensor was sharded - device_mesh: The device mesh for distributed communication - - Returns: - The full reconstructed tensor (same on all ranks) - """ - world_size = device_mesh.size() - # In case of TP+DP configuration, the TP group should be used for gathering, not the full DP group - process_group = device_mesh.get_group("tp") if "tp" in (device_mesh.mesh_dim_names or {}) else None - - # Normalize negative dimension - if shard_dim < 0: - shard_dim = local_tensor.ndim + shard_dim - - # Gather all shards - gathered_tensors = [torch.empty_like(local_tensor) for _ in range(world_size)] - dist.all_gather(gathered_tensors, local_tensor.contiguous(), group=process_group) - - # Concatenate along the shard dimension - return torch.cat(gathered_tensors, dim=shard_dim) - - -def gather_state_dict_for_save( - state_dict: dict[str, torch.Tensor], - tp_plan: dict[str, str], - device_mesh, - tp_size: int, -) -> dict[str, torch.Tensor]: - """ - Gather sharded tensors to reconstruct full tensors for saving. - - This function all-gathers each sharded tensor along its shard dimension - to reconstruct the full unsharded tensor for checkpoint saving. - - Args: - state_dict: The model state dict with local sharded tensors - tp_plan: The tensor parallel plan mapping layer patterns to shard styles - device_mesh: The device mesh for distributed communication - tp_size: The tensor parallel world size - - Returns: - State dict with full (gathered) tensors - """ - # Use the global mappings from ParallelInterface (can be extended by users) - plan_to_weight_dim = ALL_PARALLEL_STYLES.plan_to_weight_dim - plan_to_bias_dim = ALL_PARALLEL_STYLES.plan_to_bias_dim - - result = {} - for key, tensor in state_dict.items(): - # Find the matching TP plan for this parameter - param_name = key.rsplit(".", 1)[0] if "." in key else key - param_type = key.rsplit(".", 1)[1] if "." in key else None - generic_param_name = re.sub(r"\d+", "*", param_name) - # Also check the full key for nn.Parameter (e.g., MoE experts without .weight suffix) - generic_full_key = re.sub(r"\d+", "*", key) - - # Check if this parameter has a TP plan - current_plan = None - if generic_full_key in tp_plan: - # Full key match (e.g., "model.layers.*.mlp.experts.gate_up_proj" for MoE experts) - current_plan = tp_plan[generic_full_key] - elif generic_param_name in tp_plan: - current_plan = tp_plan[generic_param_name] - elif "." in generic_param_name: - parent_param_name = generic_param_name.rsplit(".", 1)[0] - if parent_param_name in tp_plan: - current_plan = tp_plan[parent_param_name] - - if current_plan is None or current_plan not in plan_to_weight_dim: - # Not sharded, keep as-is - result[key] = tensor - continue - - # Determine sharding dimension based on param type - if param_type == "bias": - shard_dim = plan_to_bias_dim.get(current_plan) - else: - shard_dim = plan_to_weight_dim.get(current_plan) - - if shard_dim is None: - # Replicated, keep as-is - result[key] = tensor - continue - - # Gather full tensor and handle packed weights repacking - full_tensor = gather_full_tensor(tensor, shard_dim, device_mesh) - if current_plan in ("packed_colwise", "packed_rowwise"): - full_tensor = repack_weights(full_tensor, shard_dim, tp_size, 2) - result[key] = full_tensor.contiguous() - - return result - - -def add_tensor_parallel_hooks_to_module( - model, - module, - current_module_plan, - layer_name, - device_mesh, -): - r""" - This function is called in `PretrainedModel.post_init()`. It is responsible of adding hooks - to the modules of the `model`, based on the `PretrainedModel._tp_plan`. - - This is the place where we add the `pre_forward` and `post_forwards` hooks. These are defined - for each `TensorParallelLayer` as `_prepare_input_fn` and `_prepare_output_fn`. - - Args: - model (`PretrainedModel`): The model containing the modules. - module (`nn.Module`): The current module to which we want to add the hooks. - current_module_plan (`str` or `None`): The tensor parallel plan for the current module, if any. - layer_name (`str`): The qualified name of the current module. - device_mesh (`dist.device_mesh.DeviceMesh`): The device mesh for distributed communication. - - """ - if current_module_plan is not None: - tp_layer = ALL_PARALLEL_STYLES[current_module_plan] - tp_layer.validate_module(module, device_mesh, layer_name) - 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 / rowwise parallelism, you just shard all the weights of the layer (weight and bias) - - packed layers: you slice the weights, then shard like above - - custom operation: - - you want to add an all-gather at the end of a local layer. - - you want to have a layer that is isolated from the rest of the world (because torch.DTensor does not work well with `.view` for instance) - - """ - param_name, param_type = parameter_name.rsplit(".", 1) if "." in parameter_name else parameter_name - tp_plan = model.tp_plan or {} - module_to_tp = model.get_submodule(param_name) - rank = int(rank) - current_shard_plan = _get_parameter_tp_plan(parameter_name, tp_plan) - - if dist.get_rank() == 0: - if current_shard_plan is None: - logger.info(f"Tensor sharding plan for {param_name} not found, using default 'replicate' plan.") - else: - logger.info(f"Tensor sharding plan for {param_name}: {current_shard_plan}") - - tp_layer = None - if current_shard_plan is not None: - try: - tp_layer = ALL_PARALLEL_STYLES[current_shard_plan] - tp_layer.empty_param = empty_param - tp_layer.device_mesh = device_mesh - tp_layer.rank = rank - param = tp_layer.shard_tensor(param, tensor_idx=None, dtype=param_casting_dtype, device=rank) - if is_contiguous: - param = param.contiguous() - except NotImplementedError as e: - print( - f"Trying to prepare {parameter_name}, but it's not supported. Corresponding module: {module_to_tp} Fix it's TP plan, current layer: {tp_layer} : {e}" - ) - else: - param = param[:].to(param_casting_dtype) - - # SUPER IMPORTANT we have to use setattr - # otherwise loading is crazy slow - if not isinstance(param, torch.nn.Parameter): - param = torch.nn.Parameter(param, requires_grad=empty_param.is_floating_point()) - setattr(module_to_tp, param_type, param) - if tp_layer is not None: - tp_layer.update_module_attributes(module_to_tp) - return param - - -def verify_tp_plan(expected_keys: list[str], tp_plan: dict[str, str] | None): - """ - Verify the TP plan of the model, log a warning if the layers that were not sharded and the rules that were not applied. - """ - - if tp_plan is None: - return - - generic_keys = {replace_layer_number_by_wildcard(key) for key in expected_keys} - unsharded_layers = set(generic_keys) - unused_rules = tp_plan.copy() - - for key in generic_keys: - param_name = key.rsplit(".", 1)[0] if "." in key else key - generic_param_name = re.sub(r"\d+", "*", param_name) - - if generic_param_name in tp_plan: - unused_rules.pop(generic_param_name, None) - unsharded_layers.discard(key) - elif "." in generic_param_name and (parent_param_name := generic_param_name.rsplit(".", 1)[0]) in tp_plan: - unused_rules.pop(parent_param_name, None) - unsharded_layers.discard(key) - - if len(unused_rules) > 0: - logger.warning(f"The following TP rules were not applied on any of the layers: {unused_rules}") - if len(unsharded_layers) > 0: - logger.warning(f"The following layers were not sharded: {', '.join(unsharded_layers)}") - - -def apply_tensor_parallelism(model, tp_plan, distributed_config, device_mesh): - """Apply tensor parallelism to a model according to the TP plan.""" - model._tp_size = distributed_config.tp_size - model._device_mesh = device_mesh - if distributed_config is not None: - if isinstance(distributed_config, dict): - distributed_config = DistributedConfig.from_dict(distributed_config) - model.config.distributed_config = distributed_config - # Set the new requested tp_plan on the model - if isinstance(tp_plan, dict): - model.tp_plan = tp_plan - model_plan = model.tp_plan - if model_plan is not None: - for v in model_plan.values(): - if v not in ALL_PARALLEL_STYLES: - raise ValueError(f"Unsupported tensor parallel style {v}. Supported styles are {ALL_PARALLEL_STYLES}") - for name, module in model.named_modules(): - if not getattr(module, "_is_hooked", False): - plan = _get_parameter_tp_plan(parameter_name=name, tp_plan=model_plan, is_weight=False) - add_tensor_parallel_hooks_to_module( - model, - module, - plan, - name, - device_mesh, - ) - module._is_hooked = True - return model + raise RuntimeError("`shard_and_distribute_module` is unavailable with the DTensor tensor-parallel loading path.") + + +__all__ = [ + "ALL_PARALLEL_STYLES", + "AllReduceParallel", + "ColwiseParallel", + "EpRouterParallel", + "MlaKvAProjParallel", + "MoeExpertsParallel", + "MoEParamShard", + "MoeIdentityParallel", + "MoeTensorParalellMegaMoeExperts", + "PackedColwiseParallel", + "PackedRowwiseParallel", + "ParallelInterface", + "ReplicatedWithGradAllReduce", + "RouterParallelMegaMoe", + "RowwiseParallel", + "SequenceParallel", + "TensorParallelLayer", + "apply_tensor_parallelism", + "gather_state_dict_for_save", + "replace_layer_number_by_wildcard", + "shard_and_distribute_module", + "verify_tp_plan", +] diff --git a/src/transformers/integrations/torchao.py b/src/transformers/integrations/torchao.py index fd6af55e456f..b274c7b940c1 100644 --- a/src/transformers/integrations/torchao.py +++ b/src/transformers/integrations/torchao.py @@ -66,6 +66,8 @@ def _quantize(self, module, config, *args, **kwargs): module.to("cpu") else: quantize_(module, config, *args, **kwargs) + # TP must use local tensors because this quantization path does not support DTensor inputs or weights. + module._hf_quantized_needs_local_tp = True def convert( self, @@ -190,6 +192,9 @@ def convert( Float8Tensor instance as the value. """ is_unsafe_serialization = list(input_dict.keys())[0] not in source_patterns + module, _ = get_module_from_name(model, full_layer_name) + # TP must use local tensors because this quantization path does not support DTensor inputs or weights. + module._hf_quantized_needs_local_tp = True param_data = {} layer_name = ".".join(full_layer_name.split(".")[:-1]) @@ -218,7 +223,6 @@ def convert( assert not leftover_state_dict # there should be no unprocessed tensors new_param = unflattened_state_dict[full_layer_name] - module, _ = get_module_from_name(model, full_layer_name) # Add repr to the module if isinstance(module, torch.nn.Linear): module.extra_repr = types.MethodType(_linear_extra_repr, module) diff --git a/src/transformers/modeling_layers.py b/src/transformers/modeling_layers.py index c072211ab48d..e34e5d6fccda 100644 --- a/src/transformers/modeling_layers.py +++ b/src/transformers/modeling_layers.py @@ -635,7 +635,6 @@ def from_pretrained(cls, main_model: PreTrainedModel, device_map=None, **kwargs) load_config=LoadStateDictConfig( weight_mapping=weight_conversions, device_map=device_map, dtype=main_model.config.dtype ), - tp_plan=None, ) # finally close all opened file pointers for k in all_pointer: diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 44fdec436bb1..7259a89f36f7 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -42,6 +42,8 @@ from torch.distributions import constraints from torch.utils.checkpoint import checkpoint +from transformers.distributed.utils import is_dtensor + from . import initialization as init from .configuration_utils import PreTrainedConfig from .conversion_mapping import get_model_conversion_mapping @@ -53,6 +55,8 @@ ) from .distributed import DistributedConfig from .distributed.mixin import DistributedMixin +from .distributed.sharding_utils import _dtensor_from_local_like +from .distributed.tensor_parallel import _get_parameter_tp_plan, verify_tp_plan from .distributed.utils import ( _get_torch_distributed_world_size, _is_torch_distributed_initialized, @@ -81,11 +85,6 @@ from .integrations.peft import maybe_load_adapters from .integrations.sdpa_attention import sdpa_attention_forward from .integrations.sdpa_paged import sdpa_attention_paged_forward -from .integrations.tensor_parallel import ( - _get_parameter_tp_plan, - shard_and_distribute_module, - verify_tp_plan, -) from .loss.loss_utils import LOSS_MAPPING from .modeling_flash_attention_utils import ( FLASH_ATTENTION_COMPATIBILITY_MATRIX, @@ -4483,7 +4482,6 @@ def _load_pretrained_model( model=model, state_dict=merged_state_dict, load_config=load_config, - tp_plan=model.tp_plan, disk_offload_index=disk_offload_index, ) @@ -4729,11 +4727,12 @@ def _move_missing_keys_from_meta_to_device( device_mesh: "DeviceMeshLike | None", hf_quantizer: HfQuantizer | None, ) -> None: - """Move the missing keys (keys that are part of the model parameters, but were NOT found in the loaded state dicts) - back from meta device to their device according to the `device_map` if any, else cpu. Takes care of sharding those - missing parameters if `device_mesh` is provided, i.e. we are using TP. - All non-persistent buffers are also moved back to the correct device (they are not part of the state_dict, but are - not missing either). + """Move missing params/buffers off meta to their target device. + + Loaded weights are handled earlier in `convert_and_load_state_dict_in_model` + via `DtensorShardOperation` and `set_param_for_module`. This only + materializes keys that were not loaded (or mismatched) so + `_initialize_missing_keys` can run proper init on them. """ is_quantized = hf_quantizer is not None # This is the only case where we do not initialize the model on meta device, so we don't have to do anything here @@ -4758,13 +4757,13 @@ def _move_missing_keys_from_meta_to_device( param_device = get_device(device_map, key, valid_torch_device=True) value = torch.empty_like(param, device=param_device) # For TP, we may need to shard the param - if device_mesh is not None: - shard_and_distribute_module( - self, value, param, key, None, False, device_mesh.get_local_rank(), device_mesh + if is_dtensor(param): + local = torch.empty(param._local_tensor.shape, dtype=param.dtype, device=param_device) + value = torch.nn.Parameter( + _dtensor_from_local_like(local, param), + requires_grad=param.requires_grad, ) - # Otherwise, just move it to device - else: - _load_parameter_into_model(self, key, value) + _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(): buffer_device = get_device(device_map, key, valid_torch_device=True) diff --git a/src/transformers/models/hunyuan_vl/modeling_hunyuan_vl.py b/src/transformers/models/hunyuan_vl/modeling_hunyuan_vl.py index 723ce7617e90..116f36c138b2 100644 --- a/src/transformers/models/hunyuan_vl/modeling_hunyuan_vl.py +++ b/src/transformers/models/hunyuan_vl/modeling_hunyuan_vl.py @@ -1127,7 +1127,7 @@ def forward( @auto_docstring class HunYuanVLForConditionalGeneration(HunYuanVLPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_rep"} + _tp_plan = {"lm_head": "colwise_gather_output"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} config: HunYuanVLConfig diff --git a/src/transformers/models/hunyuan_vl/modular_hunyuan_vl.py b/src/transformers/models/hunyuan_vl/modular_hunyuan_vl.py index 4401d05d904f..1c00aab2d0ff 100644 --- a/src/transformers/models/hunyuan_vl/modular_hunyuan_vl.py +++ b/src/transformers/models/hunyuan_vl/modular_hunyuan_vl.py @@ -1315,7 +1315,7 @@ def forward( @auto_docstring class HunYuanVLForConditionalGeneration(HunYuanVLPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_rep"} + _tp_plan = {"lm_head": "colwise_gather_output"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} config: HunYuanVLConfig diff --git a/src/transformers/models/llama4/configuration_llama4.py b/src/transformers/models/llama4/configuration_llama4.py index 79cfd063f4d4..f9c34b9a4421 100644 --- a/src/transformers/models/llama4/configuration_llama4.py +++ b/src/transformers/models/llama4/configuration_llama4.py @@ -119,6 +119,7 @@ class Llama4TextConfig(PreTrainedConfig): "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.experts": "moe_tp_experts", "layers.*.feed_forward.gate_proj": "colwise", "layers.*.feed_forward.up_proj": "colwise", "layers.*.feed_forward.down_proj": "rowwise", @@ -233,7 +234,7 @@ class Llama4Config(PreTrainedConfig): } sub_configs = {"text_config": Llama4TextConfig, "vision_config": Llama4VisionConfig} base_model_tp_plan = { - "multi_modal_projector.linear_1": "colwise_rep", + "multi_modal_projector.linear_1": "colwise_gather_output", } vision_config: dict | PreTrainedConfig | None = None diff --git a/src/transformers/quantizers/quantizer_fbgemm_fp8.py b/src/transformers/quantizers/quantizer_fbgemm_fp8.py index 3791d33efb1d..03c655fc8189 100644 --- a/src/transformers/quantizers/quantizer_fbgemm_fp8.py +++ b/src/transformers/quantizers/quantizer_fbgemm_fp8.py @@ -133,7 +133,6 @@ def _process_model_before_weight_loading( modules_to_not_convert=self.modules_to_not_convert, quantization_config=self.quantization_config, pre_quantized=self.pre_quantized, - tp_plan=model._tp_plan, ) def _process_model_after_weight_loading(self, model, **kwargs): diff --git a/src/transformers/quantizers/quantizer_finegrained_fp8.py b/src/transformers/quantizers/quantizer_finegrained_fp8.py index f9f49b8ec64b..71386deaab2a 100644 --- a/src/transformers/quantizers/quantizer_finegrained_fp8.py +++ b/src/transformers/quantizers/quantizer_finegrained_fp8.py @@ -182,20 +182,25 @@ def update_tp_plan(self, config): config.base_model_tp_plan = text_plan - # Per-impl rewrite of the experts parallel-layer kind. Applied LAST so it composes - # on top of any plan written above (e.g. the Qwen3 dense plan). Models carry the - # experts mapping under `base_model_tp_plan` and/or `base_model_ep_plan` — rewrite - # both. See `FP8Experts._impl_tp_layer_overrides`. from ..integrations.finegrained_fp8 import FP8Experts impl = getattr(config, "_experts_implementation", None) layer_overrides = FP8Experts._impl_tp_layer_overrides.get(impl) - if layer_overrides: - for plan_attr in ("base_model_tp_plan", "base_model_ep_plan"): - base_plan = getattr(config, plan_attr, None) or {} - updated_plan = {k: layer_overrides.get(v, v) for k, v in base_plan.items()} - if updated_plan != base_plan: - setattr(config, plan_attr, updated_plan) + for plan_attr in ("base_model_tp_plan", "base_model_ep_plan"): + base_plan = getattr(config, plan_attr, None) or {} + # Per-impl rewrite of the experts parallel-layer kind. Applied LAST so it composes + # on top of any plan written above (e.g. the Qwen3 dense plan). Models carry the + # experts mapping under `base_model_tp_plan` and/or `base_model_ep_plan` — rewrite + # both. See `FP8Experts._impl_tp_layer_overrides`. + updated_plan = {k: layer_overrides.get(v, v) for k, v in base_plan.items()} + + # Expert scales must be sharded along with their corresponding weights. + for key, style in list(updated_plan.items()): + if style == "grouped_gemm": + updated_plan.setdefault(f"{key}_scale_inv", style) + + if updated_plan != base_plan: + setattr(config, plan_attr, updated_plan) return config diff --git a/tests/kernels/test_deepgemm.py b/tests/kernels/test_deepgemm.py index 37f3a61d9bc8..28ff68538916 100644 --- a/tests/kernels/test_deepgemm.py +++ b/tests/kernels/test_deepgemm.py @@ -46,13 +46,13 @@ from test_utils import make_experts, make_fp8_experts import transformers.integrations.deepgemm as dg +from transformers.distributed.utils import is_dtensor from transformers.integrations.deepgemm import ( deepgemm_bf16_experts_forward, deepgemm_fp8_fp4_experts_forward, deepgemm_fp8_fp4_linear, deepgemm_fp8_fp4_megamoe_experts_forward, ) -from transformers.integrations.tensor_parallel import to_local from transformers.testing_utils import ( require_torch, require_torch_greater_or_equal, @@ -240,7 +240,7 @@ def test_to_local_is_compile_safe(self): @torch.compile(fullgraph=True) def run(x): - return to_local(x) + 1 + return x.to_local() + 1 if is_dtensor(x) else x + 1 out = run(torch.zeros(3, device=torch_device)) # a graph break / traced probe would raise here self.assertTrue(torch.equal(out, torch.ones(3, device=torch_device))) diff --git a/tests/models/doge/test_modeling_doge.py b/tests/models/doge/test_modeling_doge.py index 9325794e2d76..576a2c5c21b9 100644 --- a/tests/models/doge/test_modeling_doge.py +++ b/tests/models/doge/test_modeling_doge.py @@ -341,9 +341,9 @@ def test_tp_plan_matches_params(self): config, _ = self.model_tester.prepare_config_and_inputs_for_common() # They are valid but not always used, depending on config.is_moe flag (the modules are not the same in both cases) problematic_keys = { - "layers.*.mlp.router_gate": "colwise_rep", - "layers.*.mlp.down_embed": "rowwise_rep", - "layers.*.mlp.up_embed": "rowwise_rep", + "layers.*.mlp.router_gate": "colwise_gather_output", + "layers.*.mlp.down_embed": "rowwise_split_input", + "layers.*.mlp.up_embed": "rowwise_split_input", } if not config.is_moe: for key in problematic_keys: diff --git a/tests/quantization/finegrained_fp8/test_fp8.py b/tests/quantization/finegrained_fp8/test_fp8.py index ed08bc81a325..c99292a0c5fa 100644 --- a/tests/quantization/finegrained_fp8/test_fp8.py +++ b/tests/quantization/finegrained_fp8/test_fp8.py @@ -462,7 +462,7 @@ def test_linear_with_diff_feature_size_preserves_shape(self): class FP8DeepGEMMMultiDeviceTest(unittest.TestCase): - """`disable_deepgemm_on_multi_device` must flag FP8 modules based on the devices they actually + """`_disable_deepgemm_on_multi_device` must flag FP8 modules based on the devices they actually occupy — DeepGEMM's kernels are bound to a single CUDA context and corrupt across devices, but a model that fits on one device must keep DeepGEMM even when other GPUs are visible (no overshoot). """ @@ -475,22 +475,22 @@ def _fp8_module(device): @require_torch_multi_gpu def test_multi_device_disables_deepgemm(self): - from transformers.integrations.finegrained_fp8 import disable_deepgemm_on_multi_device + from transformers.integrations.finegrained_fp8 import _disable_deepgemm_on_multi_device model = torch.nn.Module() model.a = self._fp8_module("cuda:0") model.b = self._fp8_module("cuda:1") - disable_deepgemm_on_multi_device(model) + _disable_deepgemm_on_multi_device(model) self.assertTrue(model.a._deepgemm_disabled) self.assertTrue(model.b._deepgemm_disabled) @require_torch_gpu def test_single_device_keeps_deepgemm(self): - from transformers.integrations.finegrained_fp8 import disable_deepgemm_on_multi_device + from transformers.integrations.finegrained_fp8 import _disable_deepgemm_on_multi_device model = torch.nn.Module() model.a = self._fp8_module("cuda:0") model.b = self._fp8_module("cuda:0") - disable_deepgemm_on_multi_device(model) + _disable_deepgemm_on_multi_device(model) self.assertFalse(model.a._deepgemm_disabled) self.assertFalse(model.b._deepgemm_disabled) diff --git a/tests/tensor_parallel/test_tensor_parallel.py b/tests/tensor_parallel/test_tensor_parallel.py index 3d0e644f69cd..891911460458 100644 --- a/tests/tensor_parallel/test_tensor_parallel.py +++ b/tests/tensor_parallel/test_tensor_parallel.py @@ -11,57 +11,24 @@ # 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 +from unittest.mock import patch import torch from transformers import AutoModelForCausalLM -from transformers.integrations.tensor_parallel import ( +from transformers.distributed import tensor_parallel +from transformers.distributed.sharding_utils import DtensorShardOperation +from transformers.distributed.tensor_parallel import ( + ALL_PARALLEL_STYLES, ColwiseParallel, - EmbeddingParallel, - GroupedGemmParallel, PackedColwiseParallel, PackedRowwiseParallel, RowwiseParallel, - add_tensor_parallel_hooks_to_module, - 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): @@ -186,6 +153,7 @@ def __init__(self, world_size, rank): self.world_size = world_size self.rank = rank self.shape = (world_size,) + self.ndim = 1 def size(self): return self.world_size @@ -193,255 +161,188 @@ def size(self): def get_local_rank(self): return self.rank + def _get_parameter_placements(self, module, style, mesh=None): + placements = {} + mesh = object() if mesh is None else mesh + with patch.object( + tensor_parallel, "distribute_tensor", side_effect=lambda tensor, *args, **kwargs: tensor + ) as distribute: + for parameter_name in list(module._parameters): + style.shard_param(module, parameter_name, mesh) + placements[parameter_name] = distribute.call_args.args[2][0] + + return placements + + def _get_local_shape(self, global_shape, placement, world_size, rank): + if placement.is_replicate(): + return tuple(global_shape) + + shard_dim = placement.dim + local_size, _ = placement._local_shard_size_and_offset(global_shape[shard_dim], world_size, rank) + local_shape = list(global_shape) + local_shape[shard_dim] = local_size + return tuple(local_shape) + + def _make_dtensor_shard_op(self, mesh, placement, param_shape, local_shape): + op = object.__new__(DtensorShardOperation) + op.device_mesh = mesh + op.placements = (placement,) + op.param_ndim = len(param_shape) + op._axis0_offset = 0 + op._axis0_local_size = local_shape[0] + return op + def test_colwise_gather_output_rejects_indivisible_out_features(self): + model = torch.nn.Module() + model.lm_head = torch.nn.Linear(8, 99) + model.tp_plan = {"lm_head": "colwise_gather_output"} device_mesh = self.MockDeviceMesh(world_size=2, rank=0) with self.assertRaises(ValueError) as context: - add_tensor_parallel_hooks_to_module( - model=SimpleNamespace(config=None), - module=torch.nn.Linear(8, 99), - current_module_plan="colwise_gather_output", - layer_name="lm_head", - device_mesh=device_mesh, - ) + tensor_parallel.apply_tensor_parallelism(model, device_mesh) self.assertIn("lm_head", str(context.exception)) self.assertIn("divisible", str(context.exception)) - 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_colwise_uneven_local_shapes(self): + module = torch.nn.Module() + module.register_parameter("weight", torch.nn.Parameter(torch.empty(10, 32))) + module.register_parameter("bias", torch.nn.Parameter(torch.empty(10))) + placements = self._get_parameter_placements(module, ColwiseParallel()) + expected_local_sizes = (4, 4, 2) - 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) + for rank, expected_size in enumerate(expected_local_sizes): + weight_shape = self._get_local_shape((10, 32), placements["weight"], world_size=3, rank=rank) + bias_shape = self._get_local_shape((10,), placements["bias"], world_size=3, rank=rank) + + self.assertEqual(weight_shape, (expected_size, 32)) + self.assertEqual(bias_shape, (expected_size,)) - 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) + def test_rowwise_uneven_local_shapes(self): + module = torch.nn.Module() + module.register_parameter("weight", torch.nn.Parameter(torch.empty(32, 10))) + module.register_parameter("bias", torch.nn.Parameter(torch.empty(10))) + placements = self._get_parameter_placements(module, RowwiseParallel()) + expected_local_sizes = (4, 4, 2) - self.assertEqual( - module.__dict__, - {"random_attr": 123}, - "update_module_attributes should not modify attributes that don't exist", - ) + for rank, expected_size in enumerate(expected_local_sizes): + weight_shape = self._get_local_shape((32, 10), placements["weight"], world_size=3, rank=rank) + bias_shape = self._get_local_shape((10,), placements["bias"], world_size=3, rank=rank) + + self.assertEqual(weight_shape, (32, expected_size)) + self.assertEqual(bias_shape, (10,)) + + def test_embedding_uneven_local_shapes(self): + rowwise_embedding = torch.nn.Embedding(10, 10) + rowwise_placement = self._get_parameter_placements(rowwise_embedding, RowwiseParallel())["weight"] + + colwise_embedding = torch.nn.Embedding(10, 10) + colwise_placement = self._get_parameter_placements(colwise_embedding, ColwiseParallel())["weight"] + + expected_local_sizes = (4, 4, 2) + for rank, expected_size in enumerate(expected_local_sizes): + rowwise_shape = self._get_local_shape((10, 10), rowwise_placement, world_size=3, rank=rank) + colwise_shape = self._get_local_shape((10, 10), colwise_placement, world_size=3, rank=rank) + + self.assertEqual(rowwise_shape, (expected_size, 10)) + self.assertEqual(colwise_shape, (10, expected_size)) 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) + world_size = 4 + cases = { + "colwise": { + "module": torch.nn.Linear(32, 16), + "style": ColwiseParallel(), + "expected_shapes": {"weight": (4, 32), "bias": (4,)}, + }, + "colwise_gather_output": { + "module": torch.nn.Linear(32, 16), + "style": ALL_PARALLEL_STYLES["colwise_gather_output"], + "expected_shapes": {"weight": (4, 32), "bias": (4,)}, + }, + "rowwise": { + "module": torch.nn.Linear(32, 16), + "style": RowwiseParallel(), + "expected_shapes": {"weight": (16, 8), "bias": (16,)}, + }, + "embedding_rowwise": { + "module": torch.nn.Embedding(32, 16), + "style": ALL_PARALLEL_STYLES["embedding_rowwise"], + "expected_shapes": {"weight": (8, 16)}, + }, + "embedding_colwise": { + "module": torch.nn.Embedding(32, 16), + "style": ColwiseParallel(), + "expected_shapes": {"weight": (32, 4)}, + }, + } + + for case_name, case in cases.items(): + module = case["module"] + placements = self._get_parameter_placements(module, case["style"]) + + for parameter_name, expected_shape in case["expected_shapes"].items(): + global_shape = module._parameters[parameter_name].shape + placement = placements[parameter_name] + + for rank in range(world_size): + with self.subTest(case=case_name, parameter=parameter_name, rank=rank): + local_shape = self._get_local_shape(global_shape, placement, world_size, rank) + self.assertEqual(local_shape, expected_shape) + + def test_packed_colwise_packed_and_unpacked_shapes(self): + module = torch.nn.Module() + module.register_parameter("weight", torch.nn.Parameter(torch.empty(2, 16, 64))) + placement = self._get_parameter_placements(module, PackedColwiseParallel())["weight"] + packed = torch.randn(2, 16, 64) + unpacked_expert = torch.randn(16, 64) + + self.assertEqual(placement.dim, 1) + self.assertEqual(placement.split_factor, 2) + for rank in range(2): + mesh = self.MockDeviceMesh(world_size=2, rank=rank) + op = self._make_dtensor_shard_op(mesh, placement, param_shape=(2, 16, 64), local_shape=(2, 8, 64)) + + self.assertEqual(op.shard_tensor(packed).shape, (2, 8, 64)) + self.assertEqual(op.shard_tensor(unpacked_expert, tensor_idx=0).shape, (8, 64)) + + def test_packed_rowwise_packed_and_unpacked_shapes(self): + module = torch.nn.Module() + module.register_parameter("weight", torch.nn.Parameter(torch.empty(16, 64))) + placement = self._get_parameter_placements(module, PackedRowwiseParallel())["weight"] + packed = torch.randn(16, 64) + unpacked = torch.randn(16, 32) + + self.assertEqual(placement.dim, -1) + self.assertEqual(placement.split_factor, 2) + for rank in range(2): + mesh = self.MockDeviceMesh(world_size=2, rank=rank) + op = self._make_dtensor_shard_op(mesh, placement, param_shape=(16, 64), local_shape=(16, 32)) + + self.assertEqual(op.shard_tensor(packed).shape, (16, 32)) + self.assertEqual(op.shard_tensor(unpacked).shape, (16, 16)) + + def test_grouped_gemm_updates_local_expert_count(self): + module = torch.nn.Module() + module.num_experts = 8 + module.register_parameter("weight", torch.nn.Parameter(torch.empty(8, 16, 32))) + grouped_gemm = ALL_PARALLEL_STYLES["grouped_gemm"] + + placements = self._get_parameter_placements(module, grouped_gemm, self.MockDeviceMesh(world_size=4, rank=0)) + + self.assertEqual(placements["weight"].dim, 0) + self.assertEqual(module.num_experts, 2) + + def test_sharding_does_not_create_unrelated_module_attributes(self): + styles = (ColwiseParallel(), RowwiseParallel(), ALL_PARALLEL_STYLES["grouped_gemm"]) + + for style in styles: + with self.subTest(style=type(style).__name__): + module = torch.nn.Module() + module.random_attr = 123 + module.register_parameter("weight", torch.nn.Parameter(torch.empty(8, 16, 32))) + + self._get_parameter_placements(module, style, self.MockDeviceMesh(world_size=4, rank=0)) + + self.assertEqual(module.random_attr, 123) + self.assertFalse(hasattr(module, "num_experts")) diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py index c3d50e4433b4..74005dcc17be 100644 --- a/tests/test_modeling_common.py +++ b/tests/test_modeling_common.py @@ -135,8 +135,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 f2dbadd9338e..e1d6db0d217f 100644 --- a/tests/test_tensor_parallel_mixin.py +++ b/tests/test_tensor_parallel_mixin.py @@ -19,7 +19,7 @@ from transformers import TorchAoConfig, set_seed from transformers.distributed.configuration_utils import DistributedConfig -from transformers.integrations.tensor_parallel import _get_parameter_tp_plan +from transformers.distributed.tensor_parallel import _get_parameter_tp_plan from transformers.testing_utils import ( is_tensor_parallel_test, is_torch_available, @@ -31,10 +31,12 @@ from torchao.quantization import Float8WeightOnlyConfig +# TODO(3outeille): better guarding if is_torch_available(): import torch import torch.distributed as dist import torch.multiprocessing as mp + from torch.distributed.tensor import DTensor from torch.multiprocessing.spawn import ProcessRaisedException @@ -179,7 +181,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 in ("packed_colwise", "packed_rowwise"): 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)}" @@ -255,12 +257,17 @@ def _test_tp_backward_impl(rank, model_path, model_class, atol, rtol): grad = param.grad grad_tp = param_tp.grad + # A sharded param's grad is a DTensor: take this rank's local shard, since a DTensor + # reports the *global* shape and can't be compared against a plain tensor. + if isinstance(grad_tp, DTensor): + grad_tp = grad_tp.to_local() + # Slice reference gradient to match local shard if parameter is sharded if grad.shape != grad_tp.shape: for dim in range(grad.ndim): if grad.size(dim) != grad_tp.size(dim): param_plan = _get_parameter_tp_plan(name, model_tp._tp_plan, is_weight=True) - if param_plan in ("packed_colwise",): + if param_plan in ("packed_colwise", "packed_rowwise"): # interleaved slicing grad = get_packed_grad_shard(grad, world_size, rank, dim) else: diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py index c47aaf5e3afc..309879ab5091 100644 --- a/tests/utils/test_core_model_loading.py +++ b/tests/utils/test_core_model_loading.py @@ -386,7 +386,6 @@ def test_moe_and_qkv_conversion(self): model, state_dict, load_config, - tp_plan=None, ) self.assertEqual( @@ -503,7 +502,6 @@ def test_moe_and_qkv_conversion_reversed(self): model, state_dict, load_config, - tp_plan=None, ) self.assertTrue(len(loading_info.missing_keys) == 0) self.assertTrue(len(loading_info.unexpected_keys) == 0) @@ -603,7 +601,7 @@ def __init__(self, config): ) ] load_config = LoadStateDictConfig(weight_mapping=weight_mapping, hf_quantizer=quantizer) - loading_info, _ = convert_and_load_state_dict_in_model(model, state_dict, load_config, tp_plan=None) + loading_info, _ = convert_and_load_state_dict_in_model(model, state_dict, load_config) self.assertEqual(loading_info.missing_keys, set()) self.assertEqual(loading_info.unexpected_keys, set()) @@ -724,7 +722,6 @@ def __init__(self, config): model, checkpoint, LoadStateDictConfig(weight_mapping=[scoped_rename]), - tp_plan=None, ) # Sibling and parent keys must be unmatched. @@ -769,7 +766,6 @@ def __init__(self, config): model, checkpoint, LoadStateDictConfig(weight_mapping=[scoped_rename]), - tp_plan=None, ) self.assertEqual(loading_info.missing_keys, set()) @@ -815,7 +811,6 @@ def __init__(self, config): model, checkpoint, LoadStateDictConfig(weight_mapping=[scoped_rename]), - tp_plan=None, ) self.assertEqual(loading_info.missing_keys, set()) @@ -856,7 +851,6 @@ def __init__(self, config): model, checkpoint, LoadStateDictConfig(weight_mapping=[scoped_rename]), - tp_plan=None, ) self.assertEqual(loading_info.missing_keys, set()) self.assertEqual(loading_info.unexpected_keys, set()) @@ -928,7 +922,6 @@ def __init__(self, config): model, checkpoint, LoadStateDictConfig(weight_mapping=weight_mapping), - tp_plan=None, ) self.assertEqual(loading_info.missing_keys, set()) @@ -1004,7 +997,7 @@ def test_ernie4_5_vl_moe_conversion(self): WeightRenaming("mlp.w2.weight", "mlp.down_proj.weight"), ] loading_info, _ = convert_and_load_state_dict_in_model( - model, state_dict, LoadStateDictConfig(weight_mapping=weight_mapping), tp_plan=None + model, state_dict, LoadStateDictConfig(weight_mapping=weight_mapping) ) self.assertEqual(loading_info.missing_keys, set()) @@ -1127,7 +1120,7 @@ def test_ernie4_5_vl_moe_conversion_reversed(self): # Use the mapping to load loading_info, _ = convert_and_load_state_dict_in_model( - model, state_dict, LoadStateDictConfig(weight_mapping=weight_mapping), tp_plan=None + model, state_dict, LoadStateDictConfig(weight_mapping=weight_mapping) ) self.assertTrue(len(loading_info.missing_keys) == 0) self.assertTrue(len(loading_info.unexpected_keys) == 0) @@ -1204,7 +1197,7 @@ def __init__(self, config, fused_qkv: bool = False): ) ] load_config = LoadStateDictConfig(weight_mapping=weight_mapping) - loading_info, _ = convert_and_load_state_dict_in_model(model, state_dict_fused, load_config, tp_plan=None) + loading_info, _ = convert_and_load_state_dict_in_model(model, state_dict_fused, load_config) self.assertEqual(loading_info.missing_keys, set()) self.assertEqual(loading_info.unexpected_keys, set()) @@ -1236,9 +1229,7 @@ def __init__(self, config, fused_qkv: bool = False): ) ] load_config = LoadStateDictConfig(weight_mapping=weight_mapping) - loading_info, _ = convert_and_load_state_dict_in_model( - model_fused, state_dict_unfused, load_config, tp_plan=None - ) + loading_info, _ = convert_and_load_state_dict_in_model(model_fused, state_dict_unfused, load_config) self.assertEqual(loading_info.missing_keys, set()) self.assertEqual(loading_info.unexpected_keys, set()) @@ -1267,7 +1258,6 @@ def test_group_weight_rename(self): model, bad_serialized_checkpoints, LoadStateDictConfig(weight_mapping=copy.deepcopy(weight_mapping)), - tp_plan=None, ) # Assert we can load without issues @@ -1289,7 +1279,6 @@ def test_group_weight_rename(self): model, good_serialized_checkpoints, LoadStateDictConfig(weight_mapping=copy.deepcopy(weight_mapping)), - tp_plan=None, ) # Assert we can load without issues @@ -1372,7 +1361,6 @@ def test_can_remove_prefix(self): model, bad_serialized_checkpoints, LoadStateDictConfig(weight_mapping=copy.deepcopy(weight_mapping)), - tp_plan=None, ) # Assert we can load without issues @@ -1394,7 +1382,6 @@ def test_can_remove_prefix(self): model, good_serialized_checkpoints, LoadStateDictConfig(weight_mapping=copy.deepcopy(weight_mapping)), - tp_plan=None, ) # Assert we can load without issues @@ -1431,7 +1418,6 @@ def test_can_add_prefix(self): model, bad_serialized_checkpoints, LoadStateDictConfig(weight_mapping=copy.deepcopy(weight_mapping)), - tp_plan=None, ) # Assert we can load without issues @@ -1453,7 +1439,6 @@ def test_can_add_prefix(self): model, good_serialized_checkpoints, LoadStateDictConfig(weight_mapping=copy.deepcopy(weight_mapping)), - tp_plan=None, ) # Assert we can load without issues @@ -1491,7 +1476,6 @@ def test_can_remove_prefix_submodule(self): model, bad_serialized_checkpoints, LoadStateDictConfig(weight_mapping=copy.deepcopy(weight_mapping)), - tp_plan=None, ) # Assert we can load without issues @@ -1513,7 +1497,6 @@ def test_can_remove_prefix_submodule(self): model, good_serialized_checkpoints, LoadStateDictConfig(weight_mapping=copy.deepcopy(weight_mapping)), - tp_plan=None, ) # Assert we can load without issues @@ -1550,7 +1533,6 @@ def test_can_add_prefix_submodule(self): model, bad_serialized_checkpoints, LoadStateDictConfig(weight_mapping=copy.deepcopy(weight_mapping)), - tp_plan=None, ) # Assert we can load without issues @@ -1572,7 +1554,6 @@ def test_can_add_prefix_submodule(self): model, good_serialized_checkpoints, LoadStateDictConfig(weight_mapping=copy.deepcopy(weight_mapping)), - tp_plan=None, ) # Assert we can load without issues