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. | |
| """Neuron backend for tensor parallelism, dispatched from `apply_tensor_parallel` when the TP mesh is on Neuron. | |
| The difference from the generic path is a workaround for a Neuron NRT bug: consecutive `reduce_scatter` collectives for | |
| large weight tensors (≥ 5120×5120) can fail when all layers are distributed in a single `parallelize_module` call. The | |
| fix is to pre-shard each weight locally on CPU via `DTensor.from_local` *before* calling `parallelize_module`; the | |
| latter then sees already-placed DTensors, skips the collective for weights, but still registers the required | |
| input/output hooks for the forward pass. | |
| """ | |
| import torch | |
| import torch.distributed as dist | |
| import torch.nn as nn | |
| def _neuron_styles(relative_plan: dict) -> dict: | |
| """Map a `{relative_path: style}` plan to no-op-partition styles for Neuron. | |
| Weights (and biases) are pre-sharded in `_pre_shard_and_tp`, so `parallelize_module` runs only to register the | |
| forward hooks; `_partition_linear_fn` must not re-partition. Packed and plain styles share hook behavior, so both | |
| collapse onto the two no-op styles. | |
| """ | |
| from torch.distributed.tensor.parallel import ColwiseParallel, RowwiseParallel | |
| from .tensor_parallel import PackedColwiseParallel, PackedRowwiseParallel | |
| class _NeuronColwise(ColwiseParallel): | |
| def _partition_linear_fn(self, name, module, device_mesh): | |
| pass # weight already Shard(0) via DTensor.from_local; parallelize_module runs only for the hooks | |
| class _NeuronRowwise(RowwiseParallel): | |
| def _partition_linear_fn(self, name, module, device_mesh): | |
| pass # weight already Shard(1) via DTensor.from_local; parallelize_module runs only for the hooks | |
| resolved = {} | |
| for path, style in relative_plan.items(): | |
| if style == "colwise" or isinstance(style, PackedColwiseParallel): | |
| resolved[path] = _NeuronColwise() | |
| elif style == "rowwise" or isinstance(style, PackedRowwiseParallel): | |
| resolved[path] = _NeuronRowwise() | |
| else: | |
| raise ValueError( | |
| f"Unsupported tensor-parallel style '{style}' for '{path}'. " | |
| f"Expected 'colwise', 'rowwise', PackedColwiseParallel, or PackedRowwiseParallel." | |
| ) | |
| return resolved | |
| def _pre_shard_and_tp( | |
| module: nn.Module, | |
| tp_mesh: "torch.distributed.device_mesh.DeviceMesh", | |
| original_plan: dict, | |
| rank: int, | |
| tp_size: int, | |
| ) -> None: | |
| """Pre-shard Linear weights via `DTensor.from_local`, then call `parallelize_module`. | |
| Workaround for a Neuron NRT bug where consecutive `reduce_scatter` calls for large weight tensors (≥ 5120×5120) | |
| fail when all layers are distributed in a single `parallelize_module` call. Pre-sharding each weight on CPU means | |
| it is already an on-device DTensor when `parallelize_module` runs (via `_neuron_styles`), so the collective is | |
| skipped while the forward hooks are still registered. | |
| """ | |
| from torch.distributed.tensor import DTensor, Replicate, Shard | |
| from torch.distributed.tensor.parallel import parallelize_module | |
| from .tensor_parallel import PackedColwiseParallel, PackedRowwiseParallel, _blocks_to_block_sizes | |
| device = torch.neuron.current_device() | |
| for path, orig_style in original_plan.items(): | |
| # Resolve nested attribute path (e.g. "attn.to_q" or "attn.to_out.0") | |
| submod = module | |
| for part in path.split("."): | |
| submod = getattr(submod, part) | |
| if not hasattr(submod, "weight"): | |
| raise ValueError(f"`_tp_plan` entry '{path}' does not resolve to a module with a `weight` parameter.") | |
| w = submod.weight.data # CPU at this point | |
| b = submod.bias.data if submod.bias is not None else None | |
| if isinstance(orig_style, PackedColwiseParallel): | |
| blocks = orig_style.blocks if orig_style.blocks is not None else getattr(submod, "_tp_packed_col_blocks") | |
| block_sizes = _blocks_to_block_sizes(w.shape[0], blocks) | |
| parts, bias_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 | |
| sl = slice(offset + rank * chunk, offset + (rank + 1) * chunk) | |
| parts.append(w[sl, :].contiguous()) | |
| if b is not None: | |
| bias_parts.append(b[sl].contiguous()) | |
| offset += bs | |
| shard = torch.cat(parts, dim=0).to(device) | |
| submod.weight = nn.Parameter(DTensor.from_local(shard, tp_mesh, [Shard(0)])) | |
| if b is not None: | |
| bias_shard = torch.cat(bias_parts, dim=0).to(device) | |
| submod.bias = nn.Parameter(DTensor.from_local(bias_shard, tp_mesh, [Shard(0)])) | |
| elif isinstance(orig_style, PackedRowwiseParallel): | |
| blocks = orig_style.blocks if orig_style.blocks is not None else getattr(submod, "_tp_packed_row_blocks") | |
| block_sizes = _blocks_to_block_sizes(w.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(w[:, offset + rank * chunk : offset + (rank + 1) * chunk].contiguous()) | |
| offset += bs | |
| shard = torch.cat(parts, dim=1).to(device) | |
| submod.weight = nn.Parameter(DTensor.from_local(shard, tp_mesh, [Shard(1)])) | |
| if b is not None: # rowwise bias is added post-reduction → keep it replicated | |
| submod.bias = nn.Parameter(DTensor.from_local(b.to(device), tp_mesh, [Replicate()])) | |
| elif orig_style == "colwise": | |
| if w.shape[0] % tp_size != 0: | |
| raise ValueError( | |
| f"Cannot colwise-shard '{path}' weight rows ({w.shape[0]}) across {tp_size} " | |
| f"tensor-parallel ranks: not divisible by {tp_size}." | |
| ) | |
| rows = w.shape[0] // tp_size | |
| sl = slice(rank * rows, (rank + 1) * rows) | |
| submod.weight = nn.Parameter(DTensor.from_local(w[sl, :].contiguous().to(device), tp_mesh, [Shard(0)])) | |
| if b is not None: | |
| submod.bias = nn.Parameter(DTensor.from_local(b[sl].contiguous().to(device), tp_mesh, [Shard(0)])) | |
| elif orig_style == "rowwise": | |
| if w.shape[1] % tp_size != 0: | |
| raise ValueError( | |
| f"Cannot rowwise-shard '{path}' weight columns ({w.shape[1]}) across {tp_size} " | |
| f"tensor-parallel ranks: not divisible by {tp_size}." | |
| ) | |
| cols = w.shape[1] // tp_size | |
| shard = w[:, rank * cols : (rank + 1) * cols].contiguous().to(device) | |
| submod.weight = nn.Parameter(DTensor.from_local(shard, tp_mesh, [Shard(1)])) | |
| if b is not None: # rowwise bias is added post-reduction → keep it replicated | |
| submod.bias = nn.Parameter(DTensor.from_local(b.to(device), tp_mesh, [Replicate()])) | |
| # parallelize_module is now a no-op for weight distribution (already DTensors) | |
| # but still registers the input/output hooks required for the forward pass. | |
| parallelize_module(module, tp_mesh, _neuron_styles(original_plan)) | |
| def _apply_tp_neuron( | |
| model: nn.Module, | |
| tp_mesh: "torch.distributed.device_mesh.DeviceMesh", | |
| groups: list, | |
| ) -> None: | |
| """Apply tensor parallelism on Neuron from resolved `_tp_plan` groups. | |
| `groups` is produced by `diffusers.hooks.tensor_parallel._resolve_tp_plan` — the same source of truth used by the | |
| generic path, so the two backends shard identical layers. For each `(block, relative_plan)` group this pre-shards | |
| the weights via `DTensor.from_local` (Neuron NRT consecutive-reduce-scatter workaround), then calls | |
| `parallelize_module` to register the forward hooks. | |
| Model weights must be on CPU when this is called. | |
| """ | |
| rank = dist.get_rank() | |
| tp_size = tp_mesh.size() | |
| for block, relative_plan in groups: | |
| _pre_shard_and_tp(block, tp_mesh, relative_plan, rank, tp_size) | |