Spaces:
Running on Zero
Running on Zero
| # 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. | |
| import torch | |
| from ..models._modeling_parallel import TensorParallelConfig | |
| from ..utils import get_logger | |
| logger = get_logger(__name__) # pylint: disable=invalid-name | |
| _SUPPORTED_TP_DEVICES = ("cuda", "neuron") | |
| class PackedColwiseParallel: | |
| """Column-wise sharding for fused projections with heterogeneous block structure. | |
| `blocks` is a list of proportional integers whose sum divides the weight's row count. For example, `[1, 1]` for a | |
| SwiGLU gate+linear projection (two equal halves) or `[1, 1, 1, 3, 3]` for a Q+K+V+gate+linear projection with | |
| `mlp_ratio=3`. If `blocks` is `None`, the Linear module must carry a `_tp_packed_col_blocks` attribute set during | |
| model `__init__`. | |
| """ | |
| def __init__(self, blocks: "list[int] | None" = None): | |
| self.blocks = blocks | |
| class PackedRowwiseParallel: | |
| """Row-wise sharding for fused projections with heterogeneous block structure. | |
| `blocks` describes the input-column partition of the fused Linear (e.g. `[1, 3]` when the input concatenates an | |
| attention projection and an MLP projection with `mlp_ratio=3`). If `blocks` is `None`, the module must carry a | |
| `_tp_packed_row_blocks` attribute. | |
| """ | |
| def __init__(self, blocks: "list[int] | None" = None): | |
| self.blocks = blocks | |
| def _blocks_to_block_sizes(total_size: int, blocks: "list[int]") -> "list[int]": | |
| """Convert proportional block counts to absolute sizes. | |
| `blocks` is a list of positive integers interpreted as proportional weights. Their sum must divide `total_size` | |
| evenly. Returns a list of absolute sizes that sum to `total_size`. For example, `_blocks_to_block_sizes(1152, [1, | |
| 1, 1, 3, 3])` returns `[128, 128, 128, 384, 384]`. | |
| """ | |
| total = sum(blocks) | |
| if total_size % total != 0: | |
| raise ValueError( | |
| f"Cannot split {total_size} into proportional blocks {blocks}: " | |
| f"sum({blocks})={total} does not divide {total_size}." | |
| ) | |
| unit = total_size // total | |
| return [b * unit for b in blocks] | |
| def _resolve_tp_plan(model: torch.nn.Module, tp_plan: dict) -> list: | |
| """Group a flat `_tp_plan` into per-block `(submodule, {relative_path: style})` plans. | |
| Each glob is split at its single `*`; the prefix must resolve to a `ModuleList` and the suffix is the per-element | |
| key. Grouping by block lets the caller issue one `parallelize_module` call per block, which `RowwiseParallel` needs | |
| to attach its input redistribution at the block boundary. | |
| Example: when `transformer_blocks` is a `ModuleList` of length 2, the input `{"transformer_blocks.*.ff.linear_out": | |
| "rowwise"}` returns `[(transformer_blocks[0], {"ff.linear_out": "rowwise"}), (transformer_blocks[1], | |
| {"ff.linear_out": "rowwise"})]`. | |
| """ | |
| grouped: dict[int, tuple] = {} | |
| order: list[int] = [] | |
| for pattern, style in tp_plan.items(): | |
| if pattern.count("*") > 1: | |
| raise ValueError(f"Wildcard '*' can only be used once in a `_tp_plan` key, got '{pattern}'.") | |
| if "*" in pattern: | |
| prefix, _, suffix = pattern.partition("*") | |
| container = model | |
| for atom in prefix.strip(".").split("."): | |
| container = getattr(container, atom) | |
| if not isinstance(container, torch.nn.ModuleList): | |
| raise ValueError( | |
| f"`_tp_plan` wildcard '{pattern}' must expand over a `ModuleList`, but " | |
| f"'{prefix.strip('.')}' resolved to '{container.__class__.__name__}'." | |
| ) | |
| relative, blocks = suffix.strip("."), list(container) | |
| else: | |
| relative, blocks = pattern, [model] | |
| for block in blocks: | |
| key = id(block) | |
| if key not in grouped: | |
| grouped[key] = (block, {}) | |
| order.append(key) | |
| grouped[key][1][relative] = style | |
| return [grouped[key] for key in order] | |
| def _styles(relative_plan: dict) -> dict: | |
| """Map a `{relative_path: style}` plan to `parallelize_module` style instances. | |
| Values may be plain strings (`"colwise"` / `"rowwise"`) or `PackedColwiseParallel` / `PackedRowwiseParallel` marker | |
| instances. Returns `{relative_path: ColwiseParallel() | RowwiseParallel() | <packed impl>}`, each subclassed to | |
| reject a sharded dim that is not divisible by the TP degree. | |
| """ | |
| import torch.nn as nn | |
| from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor | |
| from torch.distributed.tensor.parallel import ColwiseParallel, RowwiseParallel | |
| def _make_packed_col(marker: PackedColwiseParallel) -> ColwiseParallel: | |
| _blocks = marker.blocks | |
| class _PackedColwiseImpl(ColwiseParallel): | |
| def _partition_linear_fn(self, name, module, device_mesh): | |
| blocks = _blocks if _blocks is not None else getattr(module, "_tp_packed_col_blocks") | |
| rank = device_mesh.get_local_rank() | |
| tp_size = device_mesh.size() | |
| # Both weight (`[out, in]`) and bias (`[out]`) are sharded row-wise (dim 0) with the same per-block | |
| # slicing so each rank's bias rows line up with its weight rows for the packed layout. | |
| for param_name, param in module.named_parameters(): | |
| full = distribute_tensor( | |
| param, device_mesh, [Replicate()], src_data_rank=self.src_data_rank | |
| ).to_local() | |
| block_sizes = _blocks_to_block_sizes(full.shape[0], blocks) | |
| parts, offset = [], 0 | |
| for bs in block_sizes: | |
| if bs % tp_size != 0: | |
| raise ValueError( | |
| f"Cannot shard packed block of size {bs} across {tp_size} tensor-parallel ranks: " | |
| f"{bs} is not divisible by {tp_size}." | |
| ) | |
| chunk = bs // tp_size | |
| parts.append(full[offset + rank * chunk : offset + (rank + 1) * chunk].contiguous()) | |
| offset += bs | |
| local = torch.cat(parts, dim=0) | |
| dist_param = nn.Parameter( | |
| DTensor.from_local(local, device_mesh, [Shard(0)], run_check=False), | |
| requires_grad=param.requires_grad, | |
| ) | |
| module.register_parameter(param_name, dist_param) | |
| return _PackedColwiseImpl() | |
| def _make_packed_row(marker: PackedRowwiseParallel) -> RowwiseParallel: | |
| _blocks = marker.blocks | |
| class _PackedRowwiseImpl(RowwiseParallel): | |
| def _partition_linear_fn(self, name, module, device_mesh): | |
| blocks = _blocks if _blocks is not None else getattr(module, "_tp_packed_row_blocks") | |
| rank = device_mesh.get_local_rank() | |
| tp_size = device_mesh.size() | |
| for param_name, param in module.named_parameters(): | |
| if param_name == "weight": | |
| full = distribute_tensor( | |
| param, device_mesh, [Replicate()], src_data_rank=self.src_data_rank | |
| ).to_local() | |
| block_sizes = _blocks_to_block_sizes(full.shape[1], blocks) | |
| parts, offset = [], 0 | |
| for bs in block_sizes: | |
| if bs % tp_size != 0: | |
| raise ValueError( | |
| f"Cannot shard packed block of size {bs} across {tp_size} tensor-parallel ranks: " | |
| f"{bs} is not divisible by {tp_size}." | |
| ) | |
| chunk = bs // tp_size | |
| parts.append(full[:, offset + rank * chunk : offset + (rank + 1) * chunk].contiguous()) | |
| offset += bs | |
| local = torch.cat(parts, dim=1) | |
| dist_param = nn.Parameter( | |
| DTensor.from_local(local, device_mesh, [Shard(1)], run_check=False), | |
| requires_grad=param.requires_grad, | |
| ) | |
| else: | |
| dist_param = nn.Parameter( | |
| distribute_tensor(param, device_mesh, [Replicate()], src_data_rank=self.src_data_rank), | |
| requires_grad=param.requires_grad, | |
| ) | |
| module.register_parameter(param_name, dist_param) | |
| return _PackedRowwiseImpl() | |
| # `distribute_tensor` accepts an indivisible shard dim and just gives the trailing ranks a smaller (or empty) | |
| # slice, so an uneven split does not raise here — it surfaces much later as a shape or numerics error, because | |
| # the attention head split and the paired colwise/rowwise Linear both assume equal shards. Reject it up front, | |
| # matching what the packed styles above and the Neuron pre-shard path already do. | |
| def _make_checked_col(path: str) -> ColwiseParallel: | |
| class _CheckedColwiseImpl(ColwiseParallel): | |
| def _partition_linear_fn(self, name, module, device_mesh): | |
| tp_size = device_mesh.size() | |
| out_features = module.weight.shape[0] | |
| if out_features % tp_size != 0: | |
| raise ValueError( | |
| f"Cannot colwise-shard '{path}' weight rows ({out_features}) across {tp_size} " | |
| f"tensor-parallel ranks: not divisible by {tp_size}." | |
| ) | |
| super()._partition_linear_fn(name, module, device_mesh) | |
| return _CheckedColwiseImpl() | |
| def _make_checked_row(path: str) -> RowwiseParallel: | |
| class _CheckedRowwiseImpl(RowwiseParallel): | |
| def _partition_linear_fn(self, name, module, device_mesh): | |
| tp_size = device_mesh.size() | |
| in_features = module.weight.shape[1] | |
| if in_features % tp_size != 0: | |
| raise ValueError( | |
| f"Cannot rowwise-shard '{path}' weight columns ({in_features}) across {tp_size} " | |
| f"tensor-parallel ranks: not divisible by {tp_size}." | |
| ) | |
| super()._partition_linear_fn(name, module, device_mesh) | |
| return _CheckedRowwiseImpl() | |
| resolved = {} | |
| for path, style in relative_plan.items(): | |
| if style == "colwise": | |
| resolved[path] = _make_checked_col(path) | |
| elif style == "rowwise": | |
| resolved[path] = _make_checked_row(path) | |
| elif isinstance(style, PackedColwiseParallel): | |
| resolved[path] = _make_packed_col(style) | |
| elif isinstance(style, PackedRowwiseParallel): | |
| resolved[path] = _make_packed_row(style) | |
| else: | |
| raise ValueError( | |
| f"Unsupported tensor-parallel style '{style}' for '{path}'. " | |
| f"Expected 'colwise', 'rowwise', PackedColwiseParallel, or PackedRowwiseParallel." | |
| ) | |
| return resolved | |
| def apply_tensor_parallel( | |
| model: torch.nn.Module, | |
| config: TensorParallelConfig, | |
| tp_plan: dict, | |
| ) -> None: | |
| """Apply tensor parallel on a model from its flat `_tp_plan`.""" | |
| tp_mesh = config._mesh | |
| if tp_mesh is None: | |
| raise ValueError("`config._mesh` is None. Call `config.setup(rank, world_size, device)` before applying TP.") | |
| if tp_mesh.device_type not in _SUPPORTED_TP_DEVICES: | |
| raise ValueError( | |
| f"Tensor parallelism is not supported on device type '{tp_mesh.device_type}'. Supported device types are " | |
| f"{list(_SUPPORTED_TP_DEVICES)}. The device type comes from the `mesh` passed to `TensorParallelConfig`, " | |
| f"or from the active accelerator when the mesh is built from `tp_degree`." | |
| ) | |
| backend = "neuron" if tp_mesh.device_type == "neuron" else "default" | |
| groups = _resolve_tp_plan(model, tp_plan) | |
| logger.debug(f"Applying tensor parallel (backend={backend}) over {len(groups)} module group(s) on mesh {tp_mesh}.") | |
| if backend == "neuron": | |
| from .tensor_parallel_neuron import _apply_tp_neuron | |
| _apply_tp_neuron(model, tp_mesh, groups) | |
| return | |
| from torch.distributed.tensor.parallel import parallelize_module | |
| for submodule, relative_plan in groups: | |
| parallelize_module(submodule, tp_mesh, _styles(relative_plan)) | |