diff --git a/docs/source/en/api/parallel.md b/docs/source/en/api/parallel.md index f2a6bee3910e..5f300d5dd566 100644 --- a/docs/source/en/api/parallel.md +++ b/docs/source/en/api/parallel.md @@ -22,3 +22,9 @@ Parallelism strategies help speed up diffusion transformers by distributing comp [[autodoc]] ContextParallelConfig [[autodoc]] hooks.apply_context_parallel + +## TensorParallelConfig + +[[autodoc]] TensorParallelConfig + +[[autodoc]] hooks.apply_tensor_parallel diff --git a/docs/source/en/training/distributed_inference.md b/docs/source/en/training/distributed_inference.md index a0a07b7f678a..6213fd905949 100644 --- a/docs/source/en/training/distributed_inference.md +++ b/docs/source/en/training/distributed_inference.md @@ -431,3 +431,133 @@ pipeline = DiffusionPipeline.from_pretrained( CKPT_ID, transformer=transformer, dtype=torch.bfloat16, ).to(device) ``` + +## Tensor parallelism + +[Tensor parallelism](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=tensor_parallelism) shards the weight matrices of a model across devices. Each device holds a column-wise (`"colwise"`) or row-wise (`"rowwise"`) slice of each layer, computes a partial result, and an `AllReduce`/`AllGather` at the layer boundary reconstructs the full output. Unlike context parallelism, it reduces the per-device *weight* memory, which is useful for models that do not fit on a single device. + +Pass a [`TensorParallelConfig`] to [`~ModelMixin.enable_parallelism`]. `tp_degree` is the number of devices to shard across and must divide the model's number of attention heads. The model must define a `_tp_plan` (a flat mapping of module-name globs to a `"colwise"`/`"rowwise"` style). + +```py +import torch +from torch import distributed as dist +from diffusers import DiffusionPipeline, TensorParallelConfig + +def setup_distributed(): + if not dist.is_initialized(): + dist.init_process_group(backend="nccl") + rank = dist.get_rank() + device = torch.device(f"cuda:{rank}") + torch.cuda.set_device(device) + return device + +def main(): + device = setup_distributed() + world_size = dist.get_world_size() + + pipeline = DiffusionPipeline.from_pretrained( + "black-forest-labs/FLUX.2-dev", torch_dtype=torch.bfloat16 + ) # weights stay on CPU + + # Shard the transformer first, then move only each rank's slice onto the accelerator. + pipeline.transformer.enable_parallelism(config=TensorParallelConfig(tp_degree=world_size)) + pipeline.transformer.to(device) + + # Move the remaining, non-sharded components onto the accelerator individually. + pipeline.text_encoder.to(device) + pipeline.vae.to(device) + + generator = torch.Generator().manual_seed(42) + image = pipeline(prompt="a cat holding a sign that says hello", generator=generator).images[0] + if dist.get_rank() == 0: + image.save("output.png") + if dist.is_initialized(): + dist.destroy_process_group() + +if __name__ == "__main__": + main() +``` + +```shell +torchrun --nproc-per-node 4 tensor_parallel_flux.py +``` + +`tp_degree` is taken from `world_size` above, so `--nproc-per-node 4` shards the transformer across 4 devices. + +### Writing a _tp_plan + +Tensor parallelism only works on models that define a `_tp_plan`, a flat class attribute mapping module-name globs to a sharding style. Writing one is mostly a matter of pairing each projection that *expands* the hidden dimension with the projection that *contracts* it back. + +Each key may contain **at most one `*`**, and the prefix before it must resolve to an [`nn.ModuleList`](https://pytorch.org/docs/stable/generated/torch.nn.ModuleList.html) so a single entry covers every block. A key without a `*` applies to the model itself. Paths are relative to the model. + +#### Colwise and rowwise + +| Style | Shards | Each rank | Use for | +|---|---|---|---| +| `"colwise"` | output features (`weight` dim 0) | computes a slice of the output | `to_q`, `to_k`, `to_v`, FFN in-projection | +| `"rowwise"` | input features (`weight` dim 1) | computes a partial sum | `to_out.0`, FFN out-projection | + +Always pair them in that order. A `"colwise"` projection leaves its output sharded, the following `"rowwise"` projection consumes that shard directly, and a single `AllReduce` at the block boundary reconstructs the result. Sharding the pair any other way forces a gather in the middle and communicates far more. + +For attention this means each rank owns a subset of heads, which is why `tp_degree` must divide the head count. Encoder-stream duplicates (`add_q_proj`, `to_add_out`, `ff_context`) follow the same pattern as their image-stream counterparts. + +#### Fused projections + +When one `Linear` packs several logical tensors along the dimension being sharded, plain `"colwise"`/`"rowwise"` slices straight across the concatenation and misaligns the pieces. Use `PackedColwiseParallel`/`PackedRowwiseParallel` instead, which shard each packed block independently. + +```py +# in src/diffusers/models/transformers/your_model.py +from ...hooks.tensor_parallel import PackedColwiseParallel, PackedRowwiseParallel +``` + +`blocks` is a list of proportional integers whose sum divides the packed dimension — `[1, 1]` for a SwiGLU gate+up projection of equal halves, or `[1, 1, 1, 3, 3]` for a fused Q+K+V+gate+up projection with `mlp_ratio=3`. + +```py +"transformer_blocks.*.ff.linear_in": PackedColwiseParallel([1, 1]), +``` + +When the block sizes are only known from the config, omit the argument and store the absolute sizes on the `Linear` during `__init__` instead, as `_tp_packed_col_blocks` or `_tp_packed_row_blocks`. + +```py +# in the attention module's __init__ +self.to_out._tp_packed_row_blocks = [self.inner_dim, self.mlp_hidden_dim] +``` + +```py +# in the model's _tp_plan +"single_transformer_blocks.*.attn.to_out": PackedRowwiseParallel(), +``` + +#### What to leave out + +Anything absent from the plan stays replicated on every rank, which is the right choice for normalization layers, AdaLN modulation (`img_mod`/`txt_mod`), patch and text embeddings, and the final `norm_out`/`proj_out`. These are small, so sharding them saves little memory while adding communication. + +#### Constraints and verification + +- `tp_degree` must divide `config.num_attention_heads`. This is validated in [`~ModelMixin.enable_parallelism`]. +- Every packed block must *individually* be divisible by `tp_degree`, not just their sum. + +Validate a new plan numerically rather than by eye: generate with a fixed seed on a single device, then again under tensor parallelism, and compare the outputs. A misplaced `"colwise"`/`"rowwise"` usually still runs and produces a plausible but wrong image. + +> [!TIP] +> Start from an existing plan for a similar architecture. [`QwenImageTransformer2DModel`] is fully unfused and every entry is plain `"colwise"`/`"rowwise"`, [`FluxTransformer2DModel`] adds a single packed row-wise projection, and [`Flux2Transformer2DModel`] covers both packed styles. + +## Choosing a strategy + +The strategies above solve different problems, and the useful question is not which is fastest in the abstract but what you are running out of. + +| Strategy | Splits | Reduces | Latency for one prompt | Best when | +|---|---|---|---|---| +| [Accelerate](#accelerate) / [DDP](#pytorch-distributed) | prompts across replicas | nothing — each device holds a full copy | unchanged | the model already fits and you have many prompts | +| [`device_map`](#device_map) | components across devices | weight memory | slightly worse | the model doesn't fit and the interconnect is slow | +| [Context parallelism](#context-parallelism) | the input sequence | activation memory | lower | sequences are long — high resolution or video | +| [Tensor parallelism](#tensor-parallelism) | weight matrices | weight memory | lower | one component's weights don't fit and the interconnect is fast | + +Some practical guidance: + +- **Throughput on many prompts, model already fits.** Use data parallelism. It is the only strategy here that scales throughput linearly without touching the model, and it leaves single-prompt latency alone. +- **A single component's weights don't fit.** Reach for tensor parallelism first, since it lowers both memory and latency. It communicates at every block boundary, so it wants a fast interconnect like NVLink; over PCIe that per-layer traffic can outweigh the compute it saves, and `device_map` becomes the better choice. `device_map` also handles the case where the components are individually fine but collectively too large. +- **Activations, not weights, are the problem.** This is the long-sequence regime — large images, many frames — and context parallelism is the direct answer. For picking a backend within it, see the [Ulysses/Ring benchmarks](#ulysses-attention) above; Ulysses gives the best throughput but caps at the attention head count, and unified attention lifts that cap once you have at least 4 devices. +- **Both weights and sequence are too large.** Combine tensor and context parallelism. [`TensorParallelConfig`] accepts a `mesh` argument so both can share one device mesh. + +Two constraints often decide this before performance does: tensor parallelism requires the model to define a [`_tp_plan`](#writing-a-_tp_plan), and its `tp_degree` must divide the attention head count. Context parallelism has no such per-model requirement and works with most attention backends. diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 7a8d727aefea..4cbf3366b6c9 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -337,6 +337,7 @@ "StableCascadeUNet", "T2IAdapter", "T5FilmDecoder", + "TensorParallelConfig", "Transformer2DModel", "TransformerTemporalModel", "UNet1DModel", @@ -1177,6 +1178,7 @@ StableAudioDiTModel, T2IAdapter, T5FilmDecoder, + TensorParallelConfig, Transformer2DModel, TransformerTemporalModel, UNet1DModel, diff --git a/src/diffusers/hooks/__init__.py b/src/diffusers/hooks/__init__.py index 2a9aa81608e7..d999ab32d6d7 100644 --- a/src/diffusers/hooks/__init__.py +++ b/src/diffusers/hooks/__init__.py @@ -27,4 +27,5 @@ from .pyramid_attention_broadcast import PyramidAttentionBroadcastConfig, apply_pyramid_attention_broadcast from .smoothed_energy_guidance_utils import SmoothedEnergyGuidanceConfig from .taylorseer_cache import TaylorSeerCacheConfig, apply_taylorseer_cache + from .tensor_parallel import apply_tensor_parallel from .text_kv_cache import TextKVCacheConfig, apply_text_kv_cache diff --git a/src/diffusers/hooks/tensor_parallel.py b/src/diffusers/hooks/tensor_parallel.py new file mode 100644 index 000000000000..b90a5761d043 --- /dev/null +++ b/src/diffusers/hooks/tensor_parallel.py @@ -0,0 +1,273 @@ +# 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() | }`, 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)) diff --git a/src/diffusers/hooks/tensor_parallel_neuron.py b/src/diffusers/hooks/tensor_parallel_neuron.py new file mode 100644 index 000000000000..6b8f219a17ff --- /dev/null +++ b/src/diffusers/hooks/tensor_parallel_neuron.py @@ -0,0 +1,178 @@ +# 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) diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index 167ee7a534de..a91758c1d5a3 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -24,7 +24,7 @@ _import_structure = {} if is_torch_available(): - _import_structure["_modeling_parallel"] = ["ContextParallelConfig", "ParallelConfig"] + _import_structure["_modeling_parallel"] = ["ContextParallelConfig", "ParallelConfig", "TensorParallelConfig"] _import_structure["adapter"] = ["MultiAdapter", "T2IAdapter"] _import_structure["attention_dispatch"] = ["AttentionBackendName", "attention_backend"] _import_structure["auto_model"] = ["AutoModel"] @@ -158,7 +158,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: if is_torch_available(): - from ._modeling_parallel import ContextParallelConfig, ParallelConfig + from ._modeling_parallel import ContextParallelConfig, ParallelConfig, TensorParallelConfig from .adapter import MultiAdapter, T2IAdapter from .attention_dispatch import AttentionBackendName, attention_backend from .auto_model import AutoModel diff --git a/src/diffusers/models/_modeling_parallel.py b/src/diffusers/models/_modeling_parallel.py index f5693f1033cf..86627284e078 100644 --- a/src/diffusers/models/_modeling_parallel.py +++ b/src/diffusers/models/_modeling_parallel.py @@ -35,7 +35,6 @@ # - Unified Attention # - More dispatcher attention backends # - CFG/Data Parallel -# - Tensor Parallel @dataclass @@ -154,6 +153,46 @@ def setup(self, rank: int, world_size: int, device: torch.device, mesh: torch.di self._ulysses_local_rank = self._ulysses_mesh.get_local_rank() +@dataclass +class TensorParallelConfig: + """ + Configuration for tensor parallelism. + + Tensor parallelism shards weight matrices (column-wise and row-wise) across devices. Each device computes a partial + result; an AllReduce/AllGather at layer boundaries reconstructs the full output. Uses + `torch.distributed.tensor.parallelize_module` with `ColwiseParallel` / `RowwiseParallel` sharding styles. Supported + device types are `"cuda"` and `"neuron"`. + + Args: + tp_degree (`int`, defaults to `1`): + Number of devices to shard across. Must be a divisor of the number of attention heads (and FFN hidden + dimensions) of the model being parallelised. + mesh (`torch.distributed.device_mesh.DeviceMesh`, *optional*): + A custom device mesh to use. If provided, `tp_degree` is inferred from `mesh.size()` and the argument is + ignored. Useful when combining TP with other parallelism strategies (e.g. CP) that share the same mesh. + """ + + tp_degree: int = 1 + mesh: torch.distributed.device_mesh.DeviceMesh | None = None + + _rank: int = None + _world_size: int = None + _device: torch.device = None + _mesh: torch.distributed.device_mesh.DeviceMesh = None + _tp_degree: int = None + + def __post_init__(self): + if self.tp_degree < 1: + raise ValueError("`tp_degree` must be >= 1.") + + def setup(self, rank: int, world_size: int, device: torch.device, mesh: torch.distributed.device_mesh.DeviceMesh): + self._rank = rank + self._world_size = world_size + self._device = device + self._mesh = mesh + self._tp_degree = mesh.size() + + @dataclass class ParallelConfig: """ @@ -162,15 +201,25 @@ class ParallelConfig: Args: context_parallel_config (`ContextParallelConfig`, *optional*): Configuration for context parallelism. + tensor_parallel_config (`TensorParallelConfig`, *optional*): + Configuration for tensor parallelism. """ context_parallel_config: ContextParallelConfig | None = None + tensor_parallel_config: TensorParallelConfig | None = None _rank: int = None _world_size: int = None _device: torch.device = None _mesh: torch.distributed.device_mesh.DeviceMesh = None + def __post_init__(self): + if self.context_parallel_config is not None and self.tensor_parallel_config is not None: + raise ValueError( + "Combining context parallelism and tensor parallelism in a single `ParallelConfig` is not supported. " + "Please specify only one of `context_parallel_config` or `tensor_parallel_config`." + ) + def setup( self, rank: int, @@ -185,6 +234,8 @@ def setup( self._mesh = mesh if self.context_parallel_config is not None: self.context_parallel_config.setup(rank, world_size, device, mesh) + if self.tensor_parallel_config is not None: + self.tensor_parallel_config.setup(rank, world_size, device, mesh) @dataclass(frozen=True) diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index 61dfc3133fbd..5af0ca0e6278 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -63,7 +63,12 @@ from ..utils.distributed_utils import is_torch_dist_rank_zero from ..utils.hub_utils import PushToHubMixin, load_or_create_model_card, populate_model_card from ..utils.torch_utils import empty_device_cache -from ._modeling_parallel import ContextParallelConfig, ContextParallelModelPlan, ParallelConfig +from ._modeling_parallel import ( + ContextParallelConfig, + ContextParallelModelPlan, + ParallelConfig, + TensorParallelConfig, +) from .model_loading_utils import ( _caching_allocator_warmup, _determine_device_map, @@ -254,6 +259,7 @@ class ModelMixin(torch.nn.Module, PushToHubMixin): _repeated_blocks = [] _parallel_config = None _cp_plan = None + _tp_plan = None _skip_keys = None def __init__(self): @@ -1601,7 +1607,7 @@ def compile_repeated_blocks(self, *args, **kwargs): def enable_parallelism( self, *, - config: ParallelConfig | ContextParallelConfig, + config: ParallelConfig | ContextParallelConfig | TensorParallelConfig, cp_plan: dict[str, ContextParallelModelPlan] | None = None, ): logger.warning( @@ -1620,6 +1626,8 @@ def enable_parallelism( if isinstance(config, ContextParallelConfig): config = ParallelConfig(context_parallel_config=config) + elif isinstance(config, TensorParallelConfig): + config = ParallelConfig(tensor_parallel_config=config) rank = torch.distributed.get_rank() world_size = torch.distributed.get_world_size() @@ -1665,19 +1673,30 @@ def enable_parallelism( mesh_shape=cp_config.mesh_shape, mesh_dim_names=cp_config.mesh_dim_names, ) + elif config.tensor_parallel_config is not None: + tp_config = config.tensor_parallel_config + mesh = tp_config.mesh or torch.distributed.device_mesh.init_device_mesh( + device_type=device_type, + mesh_shape=(tp_config.tp_degree,), + mesh_dim_names=("tp",), + ) + # `config.setup()` records the mesh resolved above onto the config; see `ParallelConfig.setup`. config.setup(rank, world_size, device, mesh=mesh) self._parallel_config = config - for module in self.modules(): - if not isinstance(module, attention_classes): - continue - processor = module.processor - if processor is None or not hasattr(processor, "_parallel_config"): - continue - processor._parallel_config = config - + # Only context parallelism needs the config inside attention: it replaces the attention computation itself + # (Ulysses all-to-all / ring). Tensor parallelism only shards `Linear` weights, so each rank runs the ordinary + # attention op over its own heads and the processors must stay unaware of it. if config.context_parallel_config is not None: + for module in self.modules(): + if not isinstance(module, attention_classes): + continue + processor = module.processor + if processor is None or not hasattr(processor, "_parallel_config"): + continue + processor._parallel_config = config + if cp_plan is None and self._cp_plan is None: raise ValueError( "`cp_plan` must be provided either as an argument or set in the model's `_cp_plan` attribute." @@ -1685,6 +1704,21 @@ def enable_parallelism( cp_plan = cp_plan if cp_plan is not None else self._cp_plan apply_context_parallel(self, config.context_parallel_config, cp_plan) + if config.tensor_parallel_config is not None: + if self._tp_plan is None: + raise ValueError( + "`_tp_plan` must be set on the model class to use tensor parallelism. " + f"'{self.__class__.__name__}' does not define one." + ) + tp_degree = config.tensor_parallel_config._tp_degree + num_heads = getattr(self.config, "num_attention_heads", None) + if num_heads is not None and num_heads % tp_degree != 0: + raise ValueError(f"`tp_degree` ({tp_degree}) must divide the number of attention heads ({num_heads}).") + + from ..hooks.tensor_parallel import apply_tensor_parallel + + apply_tensor_parallel(self, config.tensor_parallel_config, self._tp_plan) + @classmethod def _load_pretrained_model( cls, diff --git a/src/diffusers/models/transformers/transformer_bria_fibo.py b/src/diffusers/models/transformers/transformer_bria_fibo.py index 78545cb7da31..a02f59461a1e 100644 --- a/src/diffusers/models/transformers/transformer_bria_fibo.py +++ b/src/diffusers/models/transformers/transformer_bria_fibo.py @@ -64,7 +64,6 @@ def _get_qkv_projections(attn: "BriaFiboAttention", hidden_states, encoder_hidde return _get_projections(attn, hidden_states, encoder_hidden_states) -# Copied from diffusers.models.transformers.transformer_flux.FluxAttnProcessor with FluxAttnProcessor->BriaFiboAttnProcessor, FluxAttention->BriaFiboAttention class BriaFiboAttnProcessor: _attention_backend = None _parallel_config = None diff --git a/src/diffusers/models/transformers/transformer_flux.py b/src/diffusers/models/transformers/transformer_flux.py index 94857dffacb2..c8ba2d517988 100644 --- a/src/diffusers/models/transformers/transformer_flux.py +++ b/src/diffusers/models/transformers/transformer_flux.py @@ -21,6 +21,7 @@ import torch.nn.functional as F from ...configuration_utils import ConfigMixin, register_to_config +from ...hooks.tensor_parallel import PackedRowwiseParallel from ...loaders import FluxTransformer2DLoadersMixin, FromOriginalModelMixin, PeftAdapterMixin from ...utils import apply_lora_scale, logging from ...utils.torch_utils import maybe_adjust_dtype_for_device, maybe_allow_in_graph @@ -92,17 +93,19 @@ def __call__( attn, hidden_states, encoder_hidden_states ) - query = query.unflatten(-1, (attn.heads, -1)) - key = key.unflatten(-1, (attn.heads, -1)) - value = value.unflatten(-1, (attn.heads, -1)) + # Reshape by a fixed `head_dim` and let `-1` absorb the head count. Under tensor parallelism each rank + # holds a column-sharded slice (`attn.heads // tp_degree` heads); this keeps the processor TP-agnostic. + query = query.unflatten(-1, (-1, attn.head_dim)) + key = key.unflatten(-1, (-1, attn.head_dim)) + value = value.unflatten(-1, (-1, attn.head_dim)) query = attn.norm_q(query) key = attn.norm_k(key) if attn.added_kv_proj_dim is not None: - encoder_query = encoder_query.unflatten(-1, (attn.heads, -1)) - encoder_key = encoder_key.unflatten(-1, (attn.heads, -1)) - encoder_value = encoder_value.unflatten(-1, (attn.heads, -1)) + encoder_query = encoder_query.unflatten(-1, (-1, attn.head_dim)) + encoder_key = encoder_key.unflatten(-1, (-1, attn.head_dim)) + encoder_value = encoder_value.unflatten(-1, (-1, attn.head_dim)) encoder_query = attn.norm_added_q(encoder_query) encoder_key = attn.norm_added_k(encoder_key) @@ -362,6 +365,9 @@ def __init__(self, dim: int, num_attention_heads: int, attention_head_dim: int, self.proj_mlp = nn.Linear(dim, self.mlp_hidden_dim) self.act_mlp = nn.GELU(approximate="tanh") self.proj_out = nn.Linear(dim + self.mlp_hidden_dim, dim) + # proj_out input concatenates the attention output (`dim`) and the MLP branch (`mlp_hidden_dim`); each is + # independently column-sharded, so PackedRowwiseParallel splits its input columns along these block sizes. + self.proj_out._tp_packed_row_blocks = [dim, self.mlp_hidden_dim] self.attn = FluxAttention( query_dim=dim, @@ -573,6 +579,35 @@ class FluxTransformer2DModel( }, "proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3), } + # Tensor-parallel plan: how each block's fused/unfused Linears shard across the TP mesh. Flux1 is unfused, so + # nearly everything is plain "colwise"/"rowwise" (torch's ColwiseParallel / RowwiseParallel). The one exception is + # the single-stream `proj_out`, whose input is `cat([attn_output, mlp_hidden_states])` — two independently + # column-sharded streams — so it needs PackedRowwiseParallel with block sizes stored on the Linear at init. + # AdaLN modulation (norm1/norm1_context, the single-block norm, norm_out) and the QK-norms stay replicated + # (intentionally absent here): modulation indexes the full hidden dim, and QK-norm applies over head_dim after + # the heads are already split, so every rank needs the whole vector. + _tp_plan = { + # double-stream (MMDiT) blocks + "transformer_blocks.*.attn.to_q": "colwise", + "transformer_blocks.*.attn.to_k": "colwise", + "transformer_blocks.*.attn.to_v": "colwise", + "transformer_blocks.*.attn.to_out.0": "rowwise", + "transformer_blocks.*.attn.add_q_proj": "colwise", + "transformer_blocks.*.attn.add_k_proj": "colwise", + "transformer_blocks.*.attn.add_v_proj": "colwise", + "transformer_blocks.*.attn.to_add_out": "rowwise", + "transformer_blocks.*.ff.net.0.proj": "colwise", + "transformer_blocks.*.ff.net.2": "rowwise", + "transformer_blocks.*.ff_context.net.0.proj": "colwise", + "transformer_blocks.*.ff_context.net.2": "rowwise", + # single-stream (parallel attention + MLP) blocks + "single_transformer_blocks.*.attn.to_q": "colwise", + "single_transformer_blocks.*.attn.to_k": "colwise", + "single_transformer_blocks.*.attn.to_v": "colwise", + "single_transformer_blocks.*.proj_mlp": "colwise", + # input = cat([attn_output (inner_dim), mlp_hidden_states (mlp_hidden_dim)]); blocks set in __init__ + "single_transformer_blocks.*.proj_out": PackedRowwiseParallel(), + } @register_to_config def __init__( diff --git a/src/diffusers/models/transformers/transformer_flux2.py b/src/diffusers/models/transformers/transformer_flux2.py index 17c8bd0ffd52..be56ecb935a1 100644 --- a/src/diffusers/models/transformers/transformer_flux2.py +++ b/src/diffusers/models/transformers/transformer_flux2.py @@ -21,6 +21,7 @@ import torch.nn.functional as F from ...configuration_utils import ConfigMixin, register_to_config +from ...hooks.tensor_parallel import PackedColwiseParallel, PackedRowwiseParallel from ...loaders import FluxTransformer2DLoadersMixin, FromOriginalModelMixin, PeftAdapterMixin from ...utils import BaseOutput, apply_lora_scale, logging from ...utils.torch_utils import maybe_adjust_dtype_for_device @@ -343,17 +344,19 @@ def __call__( attn, hidden_states, encoder_hidden_states ) - query = query.unflatten(-1, (attn.heads, -1)) - key = key.unflatten(-1, (attn.heads, -1)) - value = value.unflatten(-1, (attn.heads, -1)) + # Reshape by a fixed `head_dim` and let `-1` absorb the head count. Under tensor parallelism each rank + # holds a column-sharded slice (`attn.heads // tp_degree` heads); this keeps the processor TP-agnostic. + query = query.unflatten(-1, (-1, attn.head_dim)) + key = key.unflatten(-1, (-1, attn.head_dim)) + value = value.unflatten(-1, (-1, attn.head_dim)) query = attn.norm_q(query) key = attn.norm_k(key) if attn.added_kv_proj_dim is not None: - encoder_query = encoder_query.unflatten(-1, (attn.heads, -1)) - encoder_key = encoder_key.unflatten(-1, (attn.heads, -1)) - encoder_value = encoder_value.unflatten(-1, (attn.heads, -1)) + encoder_query = encoder_query.unflatten(-1, (-1, attn.head_dim)) + encoder_key = encoder_key.unflatten(-1, (-1, attn.head_dim)) + encoder_value = encoder_value.unflatten(-1, (-1, attn.head_dim)) encoder_query = attn.norm_added_q(encoder_query) encoder_key = attn.norm_added_k(encoder_key) @@ -423,17 +426,17 @@ def __call__( attn, hidden_states, encoder_hidden_states ) - query = query.unflatten(-1, (attn.heads, -1)) - key = key.unflatten(-1, (attn.heads, -1)) - value = value.unflatten(-1, (attn.heads, -1)) + query = query.unflatten(-1, (-1, attn.head_dim)) + key = key.unflatten(-1, (-1, attn.head_dim)) + value = value.unflatten(-1, (-1, attn.head_dim)) query = attn.norm_q(query) key = attn.norm_k(key) if attn.added_kv_proj_dim is not None: - encoder_query = encoder_query.unflatten(-1, (attn.heads, -1)) - encoder_key = encoder_key.unflatten(-1, (attn.heads, -1)) - encoder_value = encoder_value.unflatten(-1, (attn.heads, -1)) + encoder_query = encoder_query.unflatten(-1, (-1, attn.head_dim)) + encoder_key = encoder_key.unflatten(-1, (-1, attn.head_dim)) + encoder_value = encoder_value.unflatten(-1, (-1, attn.head_dim)) encoder_query = attn.norm_added_q(encoder_query) encoder_key = attn.norm_added_k(encoder_key) @@ -583,16 +586,21 @@ def __call__( ) -> torch.Tensor: # Parallel in (QKV + MLP in) projection hidden_states = attn.to_qkv_mlp_proj(hidden_states) - qkv, mlp_hidden_states = torch.split( - hidden_states, [3 * attn.inner_dim, attn.mlp_hidden_dim * attn.mlp_mult_factor], dim=-1 - ) + + # Split the fused output into its QKV and MLP halves by their global ratio, so the same code path works + # whether or not `to_qkv_mlp_proj` is column-sharded (PackedColwiseParallel shards each block, so every + # rank's contiguous slice keeps the QKV and MLP chunks proportional). No tensor-parallel state is read here. + qkv_dim = 3 * attn.inner_dim + mlp_dim = attn.mlp_hidden_dim * attn.mlp_mult_factor + local_qkv = hidden_states.shape[-1] * qkv_dim // (qkv_dim + mlp_dim) + qkv, mlp_hidden_states = torch.split(hidden_states, [local_qkv, hidden_states.shape[-1] - local_qkv], dim=-1) # Handle the attention logic query, key, value = qkv.chunk(3, dim=-1) - query = query.unflatten(-1, (attn.heads, -1)) - key = key.unflatten(-1, (attn.heads, -1)) - value = value.unflatten(-1, (attn.heads, -1)) + query = query.unflatten(-1, (-1, attn.head_dim)) + key = key.unflatten(-1, (-1, attn.head_dim)) + value = value.unflatten(-1, (-1, attn.head_dim)) query = attn.norm_q(query) key = attn.norm_k(key) @@ -651,15 +659,21 @@ def __call__( ) -> torch.Tensor: # Parallel in (QKV + MLP in) projection hidden_states_proj = attn.to_qkv_mlp_proj(hidden_states) + + # Split by the global QKV:MLP ratio so the path is identical with or without column sharding (see + # `Flux2ParallelSelfAttnProcessor`). No tensor-parallel state is read here. + qkv_dim = 3 * attn.inner_dim + mlp_dim = attn.mlp_hidden_dim * attn.mlp_mult_factor + local_qkv = hidden_states_proj.shape[-1] * qkv_dim // (qkv_dim + mlp_dim) qkv, mlp_hidden_states = torch.split( - hidden_states_proj, [3 * attn.inner_dim, attn.mlp_hidden_dim * attn.mlp_mult_factor], dim=-1 + hidden_states_proj, [local_qkv, hidden_states_proj.shape[-1] - local_qkv], dim=-1 ) query, key, value = qkv.chunk(3, dim=-1) - query = query.unflatten(-1, (attn.heads, -1)) - key = key.unflatten(-1, (attn.heads, -1)) - value = value.unflatten(-1, (attn.heads, -1)) + query = query.unflatten(-1, (-1, attn.head_dim)) + key = key.unflatten(-1, (-1, attn.head_dim)) + value = value.unflatten(-1, (-1, attn.head_dim)) query = attn.norm_q(query) key = attn.norm_k(key) @@ -754,6 +768,10 @@ def __init__( self.to_qkv_mlp_proj = torch.nn.Linear( self.query_dim, self.inner_dim * 3 + self.mlp_hidden_dim * self.mlp_mult_factor, bias=bias ) + # Block structure for tensor-parallel packed sharding: Q, K, V, gate, linear + self.to_qkv_mlp_proj._tp_packed_col_blocks = [self.inner_dim] * 3 + [ + self.mlp_hidden_dim + ] * self.mlp_mult_factor self.mlp_act_fn = Flux2SwiGLU() # QK Norm @@ -762,6 +780,8 @@ def __init__( # Fused attention output projection + MLP output projection self.to_out = torch.nn.Linear(self.inner_dim + self.mlp_hidden_dim, self.out_dim, bias=out_bias) + # Block structure for tensor-parallel packed sharding: attn columns, mlp columns + self.to_out._tp_packed_row_blocks = [self.inner_dim, self.mlp_hidden_dim] if processor is None: processor = self._default_processor_cls() @@ -1090,6 +1110,34 @@ class Flux2Transformer2DModel( "proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3), } + # Tensor-parallel sharding plan: a flat mapping of module-name globs (relative to the model) to a + # parallel style. Plain strings ("colwise" / "rowwise") map to torch's ColwiseParallel / + # RowwiseParallel. Fused projections use PackedColwiseParallel / PackedRowwiseParallel, which + # take the per-block structure from either an explicit `blocks` argument or a + # `_tp_packed_col_blocks` / `_tp_packed_row_blocks` attribute stored on the Linear at init. + # AdaLN modulation (double_stream_modulation_img/_txt, single_stream_modulation, norm_out) and the QK-norms stay + # replicated (intentionally absent here): modulation indexes the full hidden dim, and QK-norm applies over + # head_dim after the heads are already split, so every rank needs the whole vector. + _tp_plan = { + # double-stream (cross-attention + FFN) blocks + "transformer_blocks.*.attn.to_q": "colwise", + "transformer_blocks.*.attn.to_k": "colwise", + "transformer_blocks.*.attn.to_v": "colwise", + "transformer_blocks.*.attn.to_out.0": "rowwise", + "transformer_blocks.*.attn.add_q_proj": "colwise", + "transformer_blocks.*.attn.add_k_proj": "colwise", + "transformer_blocks.*.attn.add_v_proj": "colwise", + "transformer_blocks.*.attn.to_add_out": "rowwise", + # SwiGLU gate+linear are equal halves → blocks=[1,1] + "transformer_blocks.*.ff.linear_in": PackedColwiseParallel([1, 1]), + "transformer_blocks.*.ff.linear_out": "rowwise", + "transformer_blocks.*.ff_context.linear_in": PackedColwiseParallel([1, 1]), + "transformer_blocks.*.ff_context.linear_out": "rowwise", + # single-stream: block structure stored on the Linear by Flux2ParallelSelfAttention.__init__ + "single_transformer_blocks.*.attn.to_qkv_mlp_proj": PackedColwiseParallel(), + "single_transformer_blocks.*.attn.to_out": PackedRowwiseParallel(), + } + @register_to_config def __init__( self, diff --git a/src/diffusers/models/transformers/transformer_nucleusmoe_image.py b/src/diffusers/models/transformers/transformer_nucleusmoe_image.py index f1c0eee949f7..cced492fd4ea 100644 --- a/src/diffusers/models/transformers/transformer_nucleusmoe_image.py +++ b/src/diffusers/models/transformers/transformer_nucleusmoe_image.py @@ -36,7 +36,6 @@ logger = logging.get_logger(__name__) -# Copied from diffusers.models.transformers.transformer_qwenimage.apply_rotary_emb_qwen with qwen->nucleus def _apply_rotary_emb_nucleus( x: torch.Tensor, freqs_cis: torch.Tensor | tuple[torch.Tensor], diff --git a/src/diffusers/models/transformers/transformer_qwenimage.py b/src/diffusers/models/transformers/transformer_qwenimage.py index 464712bd94fd..6ed2ac7a8961 100644 --- a/src/diffusers/models/transformers/transformer_qwenimage.py +++ b/src/diffusers/models/transformers/transformer_qwenimage.py @@ -93,52 +93,26 @@ def get_timestep_embedding( return emb -def apply_rotary_emb_qwen( - x: torch.Tensor, - freqs_cis: torch.Tensor | tuple[torch.Tensor], - use_real: bool = True, - use_real_unbind_dim: int = -1, -) -> tuple[torch.Tensor, torch.Tensor]: - """ - Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings - to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are - reshaped as complex numbers, and the frequency tensor is reshaped for broadcasting compatibility. The resulting - tensors contain rotary embeddings and are returned as real tensors. +def apply_rotary_emb_qwen(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + """Apply rotary position embeddings to `x` using real-valued cos/sin. + + `freqs` holds the per-position rotation *angles* with shape `[seq_len, head_dim // 2]`. The real cos/sin rotation + is numerically identical to the equivalent complex-exponential formulation. Args: - x (`torch.Tensor`): - Query or key tensor to apply rotary embeddings. [B, S, H, D] xk (torch.Tensor): Key tensor to apply - freqs_cis (`tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],) + x (`torch.Tensor`): Query or key tensor to rotate, shape `[B, S, H, D]`. + freqs (`torch.Tensor`): Rotation angles, shape `[S, D // 2]`. Returns: - tuple[torch.Tensor, torch.Tensor]: tuple of modified query tensor and key tensor with rotary embeddings. + `torch.Tensor`: `x` with rotary embeddings applied. """ - if use_real: - cos, sin = freqs_cis # [S, D] - cos = cos[None, None] - sin = sin[None, None] - cos, sin = cos.to(x.device), sin.to(x.device) - - if use_real_unbind_dim == -1: - # Used for flux, cogvideox, hunyuan-dit - x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, S, H, D//2] - x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3) - elif use_real_unbind_dim == -2: - # Used for Stable Audio, OmniGen, CogView4 and Cosmos - x_real, x_imag = x.reshape(*x.shape[:-1], 2, -1).unbind(-2) # [B, S, H, D//2] - x_rotated = torch.cat([-x_imag, x_real], dim=-1) - else: - raise ValueError(f"`use_real_unbind_dim={use_real_unbind_dim}` but should be -1 or -2.") - - out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype) - - return out - else: - x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) - freqs_cis = freqs_cis.unsqueeze(1) - x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3) - - return x_out.type_as(x) + # Adjacent feature pairs (2k, 2k+1) share angle k, so each angle is repeated twice along the last dim; unsqueeze + # the head axis so the freqs broadcast over heads (this is what keeps it tensor-parallel-agnostic). + cos = torch.cos(freqs).repeat_interleave(2, dim=-1).unsqueeze(1) # [S, 1, D] + sin = torch.sin(freqs).repeat_interleave(2, dim=-1).unsqueeze(1) # [S, 1, D] + x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, S, H, D//2] + x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3) # [B, S, H, D] + return (x.float() * cos + x_rotated.float() * sin).to(x.dtype) def compute_text_seq_len_from_mask( @@ -220,7 +194,6 @@ def __init__(self, theta: int, axes_dim: list[int], scale_rope=False): dim=1, ) - # DO NOT USING REGISTER BUFFER HERE, IT WILL CAUSE COMPLEX NUMBERS LOSE ITS IMAGINARY PART self.scale_rope = scale_rope def rope_params(self, index, dim, theta=10000): @@ -229,8 +202,8 @@ def rope_params(self, index, dim, theta=10000): index: [0, 1, 2, 3] 1D Tensor representing the position index of the token """ assert dim % 2 == 0 + # Return the raw rotation angles (real); cos/sin are taken later in apply_rotary_emb_qwen. freqs = torch.outer(index, 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float32).div(dim))) - freqs = torch.polar(torch.ones_like(freqs), freqs) return freqs @lru_cache_unless_export(maxsize=None) @@ -353,8 +326,8 @@ def rope_params(self, index, dim, theta=10000): index: [0, 1, 2, 3] 1D Tensor representing the position index of the token """ assert dim % 2 == 0 + # Real rotation angles (see QwenEmbedRope.rope_params). freqs = torch.outer(index, 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float32).div(dim))) - freqs = torch.polar(torch.ones_like(freqs), freqs) return freqs @lru_cache_unless_export(maxsize=None) @@ -521,14 +494,17 @@ def __call__( txt_key = attn.add_k_proj(encoder_hidden_states) txt_value = attn.add_v_proj(encoder_hidden_states) - # Reshape for multi-head attention - img_query = img_query.unflatten(-1, (attn.heads, -1)) - img_key = img_key.unflatten(-1, (attn.heads, -1)) - img_value = img_value.unflatten(-1, (attn.heads, -1)) + # Reshape by a fixed `head_dim` and let `-1` absorb the head count. Under tensor parallelism each rank + # holds a column-sharded slice (`attn.heads // tp_degree` heads); this keeps the processor TP-agnostic. + # The shared Attention has no `head_dim` attribute; derive it from the unsharded `inner_dim`/`heads`. + head_dim = attn.inner_dim // attn.heads + img_query = img_query.unflatten(-1, (-1, head_dim)) + img_key = img_key.unflatten(-1, (-1, head_dim)) + img_value = img_value.unflatten(-1, (-1, head_dim)) - txt_query = txt_query.unflatten(-1, (attn.heads, -1)) - txt_key = txt_key.unflatten(-1, (attn.heads, -1)) - txt_value = txt_value.unflatten(-1, (attn.heads, -1)) + txt_query = txt_query.unflatten(-1, (-1, head_dim)) + txt_key = txt_key.unflatten(-1, (-1, head_dim)) + txt_value = txt_value.unflatten(-1, (-1, head_dim)) # Apply QK normalization if attn.norm_q is not None: @@ -543,10 +519,10 @@ def __call__( # Apply RoPE if image_rotary_emb is not None: img_freqs, txt_freqs = image_rotary_emb - img_query = apply_rotary_emb_qwen(img_query, img_freqs, use_real=False) - img_key = apply_rotary_emb_qwen(img_key, img_freqs, use_real=False) - txt_query = apply_rotary_emb_qwen(txt_query, txt_freqs, use_real=False) - txt_key = apply_rotary_emb_qwen(txt_key, txt_freqs, use_real=False) + img_query = apply_rotary_emb_qwen(img_query, img_freqs) + img_key = apply_rotary_emb_qwen(img_key, img_freqs) + txt_query = apply_rotary_emb_qwen(txt_query, txt_freqs) + txt_key = apply_rotary_emb_qwen(txt_key, txt_freqs) # Concatenate for joint attention # Order: [text, image] @@ -790,6 +766,25 @@ class QwenImageTransformer2DModel( }, "proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3), } + # Tensor-parallel plan: how each block's fused/unfused Linears shard across the TP mesh. Qwen-Image is fully + # unfused (separate Q/K/V and separate img/txt MLPs), so every entry is plain "colwise"/"rowwise" (torch's + # ColwiseParallel / RowwiseParallel) and no packed sharding is needed. AdaLN modulation (img_mod/txt_mod) and + # norm_out stay replicated (intentionally absent here). + _tp_plan = { + # double-stream (MMDiT) blocks + "transformer_blocks.*.attn.to_q": "colwise", + "transformer_blocks.*.attn.to_k": "colwise", + "transformer_blocks.*.attn.to_v": "colwise", + "transformer_blocks.*.attn.to_out.0": "rowwise", + "transformer_blocks.*.attn.add_q_proj": "colwise", + "transformer_blocks.*.attn.add_k_proj": "colwise", + "transformer_blocks.*.attn.add_v_proj": "colwise", + "transformer_blocks.*.attn.to_add_out": "rowwise", + "transformer_blocks.*.img_mlp.net.0.proj": "colwise", + "transformer_blocks.*.img_mlp.net.2": "rowwise", + "transformer_blocks.*.txt_mlp.net.0.proj": "colwise", + "transformer_blocks.*.txt_mlp.net.2": "rowwise", + } @register_to_config def __init__( diff --git a/src/diffusers/pipelines/pipeline_utils.py b/src/diffusers/pipelines/pipeline_utils.py index a683973df5d7..3fd10c6bc256 100644 --- a/src/diffusers/pipelines/pipeline_utils.py +++ b/src/diffusers/pipelines/pipeline_utils.py @@ -1160,6 +1160,15 @@ def _execution_device(self): except ValueError: pass + # When text encoders are offloaded to CPU while the denoising backbone + # (transformer, unet, vae) runs on an accelerator, self.device returns CPU + # (first component). Prefer any non-CPU, non-meta component so that + # latent tensors land on the accelerator. This covers CUDA, XPU, NPU, HPU, + # and any other backend, including TP-sharded models via DTensor. + for name, model in self.components.items(): + if isinstance(model, torch.nn.Module) and model.device.type not in ("cpu", "meta"): + return model.device + for name, model in self.components.items(): if not isinstance(model, torch.nn.Module) or name in self._exclude_from_cpu_offload: continue diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 9035efb3e6e2..29c36f0d7d33 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -2085,6 +2085,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class TensorParallelConfig(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class Transformer2DModel(metaclass=DummyObject): _backends = ["torch"] diff --git a/src/diffusers/utils/torch_utils.py b/src/diffusers/utils/torch_utils.py index 981e05a7c055..9f0877eb1d13 100644 --- a/src/diffusers/utils/torch_utils.py +++ b/src/diffusers/utils/torch_utils.py @@ -200,7 +200,7 @@ def randn_tensor( layout = layout or torch.strided device = device or torch.device("cpu") - # Neuron (XLA) does not support creating random tensors directly on device; always use CPU + # Neuron does not support creating random tensors directly on device; always use CPU if device.type == "neuron": rand_device = torch.device("cpu") diff --git a/tests/conftest.py b/tests/conftest.py index 4e83806661d2..b6f8c18183c7 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -49,6 +49,7 @@ def pytest_configure(config): config.addinivalue_line("markers", "modelopt: marks tests for NVIDIA ModelOpt quantization functionality") config.addinivalue_line("markers", "sdnq: marks tests for SDNQ quantization functionality") config.addinivalue_line("markers", "context_parallel: marks tests for context parallel inference functionality") + config.addinivalue_line("markers", "tensor_parallel: marks tests for tensor parallel inference functionality") config.addinivalue_line("markers", "slow: mark test as slow") config.addinivalue_line("markers", "nightly: mark test as nightly") diff --git a/tests/models/testing_utils/__init__.py b/tests/models/testing_utils/__init__.py index 23aa871b80a2..1bd464d5677f 100644 --- a/tests/models/testing_utils/__init__.py +++ b/tests/models/testing_utils/__init__.py @@ -17,7 +17,11 @@ from .ip_adapter import IPAdapterTesterMixin from .lora import LoraHotSwappingForModelTesterMixin, LoraTesterMixin from .memory import CPUOffloadTesterMixin, GroupOffloadTesterMixin, LayerwiseCastingTesterMixin, MemoryTesterMixin -from .parallelism import ContextParallelAttentionBackendsTesterMixin, ContextParallelTesterMixin +from .parallelism import ( + ContextParallelAttentionBackendsTesterMixin, + ContextParallelTesterMixin, + TensorParallelTesterMixin, +) from .quantization import ( AutoRoundCompileTesterMixin, AutoRoundConfigMixin, @@ -62,6 +66,7 @@ "CacheTesterMixin", "ContextParallelTesterMixin", "ContextParallelAttentionBackendsTesterMixin", + "TensorParallelTesterMixin", "CPUOffloadTesterMixin", "FasterCacheConfigMixin", "FasterCacheTesterMixin", diff --git a/tests/models/testing_utils/parallelism.py b/tests/models/testing_utils/parallelism.py index dad2f51f1f2c..63575abf6b7b 100644 --- a/tests/models/testing_utils/parallelism.py +++ b/tests/models/testing_utils/parallelism.py @@ -21,13 +21,14 @@ import torch.distributed as dist import torch.multiprocessing as mp -from diffusers.models._modeling_parallel import ContextParallelConfig +from diffusers.models._modeling_parallel import ContextParallelConfig, TensorParallelConfig from diffusers.models.attention_dispatch import AttentionBackendName, _AttentionBackendRegistry from ...testing_utils import ( is_attention, is_context_parallel, is_kernels_available, + is_tensor_parallel, require_torch_multi_accelerator, torch_device, ) @@ -241,6 +242,109 @@ def _custom_mesh_worker( dist.destroy_process_group() +def _tensor_parallel_worker( + rank, world_size, master_port, model_class, init_dict, inputs_dict, return_dict, state_dict +): + """Worker function for tensor parallel inference testing. + + Each rank builds the (identical, `state_dict`-loaded) model, sets up the accelerator device, shards the model + with `enable_parallelism(config=TensorParallelConfig(tp_degree=world_size))` and runs a forward pass. Rank 0 + reports its output so the caller can compare it against a single-device reference (TP is mathematically equivalent + to the unsharded model up to floating-point reduction order). + """ + try: + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = str(master_port) + os.environ["RANK"] = str(rank) + os.environ["WORLD_SIZE"] = str(world_size) + + device_config = DEVICE_CONFIG.get(torch_device, DEVICE_CONFIG["cuda"]) + backend = device_config["backend"] + device_module = device_config["module"] + + dist.init_process_group(backend=backend, rank=rank, world_size=world_size) + + device_module.set_device(rank) + device = torch.device(f"{torch_device}:{rank}") + + model = model_class(**init_dict) + model.load_state_dict(state_dict) + model.to(device) + model.eval() + + inputs_on_device = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in inputs_dict.items()} + + # Shard the model across all ranks; the device mesh is built from `tp_degree` on the active accelerator. + model.enable_parallelism(config=TensorParallelConfig(tp_degree=world_size)) + + with torch.no_grad(): + output = model(**inputs_on_device, return_dict=False)[0] + + if rank == 0: + return_dict["status"] = "success" + return_dict["output_shape"] = list(output.shape) + # Serialise via nested list so the manager dict can transport it across processes. + return_dict["output"] = output.float().cpu().tolist() + + except Exception as e: + if rank == 0: + return_dict["status"] = "error" + return_dict["error"] = str(e) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +@is_tensor_parallel +@require_torch_multi_accelerator +class TensorParallelTesterMixin: + def test_tensor_parallel_inference(self, batch_size: int = 1): + if not torch.distributed.is_available(): + pytest.skip("torch.distributed is not available.") + + if getattr(self.model_class, "_tp_plan", None) is None: + pytest.skip("Model does not define a `_tp_plan` for tensor parallel inference.") + + world_size = 2 + init_dict = self.get_init_dict() + num_heads = init_dict.get("num_attention_heads") + if num_heads is not None and num_heads % world_size != 0: + pytest.skip(f"`num_attention_heads` ({num_heads}) is not divisible by tp_degree ({world_size}).") + + inputs_dict = self.get_dummy_inputs(batch_size=batch_size) + + # Single-device reference + model = self.model_class(**init_dict).eval().to(torch_device) + state_dict = {k: v.cpu() for k, v in model.state_dict().items()} + with torch.no_grad(): + ref_output = model(**inputs_dict, return_dict=False)[0].float().cpu() + + # Move all tensors to CPU for multiprocessing + inputs_dict = {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in inputs_dict.items()} + + master_port = _find_free_port() + manager = mp.Manager() + return_dict = manager.dict() + + mp.spawn( + _tensor_parallel_worker, + args=(world_size, master_port, self.model_class, init_dict, inputs_dict, return_dict, state_dict), + nprocs=world_size, + join=True, + ) + + assert return_dict.get("status") == "success", ( + f"Tensor parallel inference failed: {return_dict.get('error', 'Unknown error')}" + ) + + tp_output = torch.tensor(return_dict["output"]) + # Sharded matmuls + all-reduce reorder the summation, so allow a small tolerance over the reference. + torch.testing.assert_close(ref_output, tp_output, atol=1e-3, rtol=1e-3) + + def test_tensor_parallel_batch_inputs(self): + self.test_tensor_parallel_inference(batch_size=2) + + @is_context_parallel @require_torch_multi_accelerator class ContextParallelTesterMixin: diff --git a/tests/models/transformers/_neuron_tp_worker.py b/tests/models/transformers/_neuron_tp_worker.py new file mode 100644 index 000000000000..681f5686bee9 --- /dev/null +++ b/tests/models/transformers/_neuron_tp_worker.py @@ -0,0 +1,111 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Generic torchrun worker: assert a model's Neuron tensor-parallel output matches its single-device reference. + +Model-agnostic. The model under test is supplied as a `module:function` spec reference on the command line; the +referenced factory returns `(model_class, init_dict, inputs)` with CPU tensors, so all model-specific test data lives +with the launching test rather than here. + +Launched as a subprocess by a `@require_torch_neuron` test (and runnable directly for debugging):: + + torchrun --nproc_per_node=2 _neuron_tp_worker.py \\ + tests.models.transformers.test_models_transformer_flux2:make_neuron_tp_spec + +Each rank builds an identical (seeded) model on CPU, computes a single-device reference, then shards it with +`enable_parallelism(TensorParallelConfig(mesh=neuron_mesh))` — which auto-selects the Neuron pre-shard backend — runs +a forward pass on the Neuron device, and asserts the gathered output matches the reference. Exit code 0 means the TP +path is numerically equivalent to the unsharded model; non-zero means failure. +""" + +import argparse +import importlib +import os +import sys +import traceback + + +# Make the in-repo `diffusers` and `tests` packages importable when run via torchrun from an arbitrary CWD. +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "..", "src")) +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "..")) + +import torch +import torch.distributed as dist +import torch_neuronx # noqa: F401 — registers torch.neuron +from torch.distributed.device_mesh import DeviceMesh + +from diffusers import TensorParallelConfig + + +def main(): + parser = argparse.ArgumentParser(description="Neuron tensor-parallel correctness worker.") + parser.add_argument( + "spec", + help="`module:function` reference returning (model_class, init_dict, cpu_inputs) for the model under test.", + ) + args = parser.parse_args() + module_name, _, fn_name = args.spec.partition(":") + model_class, init_dict, inputs = getattr(importlib.import_module(module_name), fn_name)() + + dist.init_process_group(backend="neuron") + rank = dist.get_rank() + tp_size = dist.get_world_size() + device = torch.neuron.current_device() + tp_mesh = DeviceMesh("neuron", list(range(tp_size))) + + # Identical weights on every rank (same seed), kept on CPU as the Neuron pre-shard backend requires. + torch.manual_seed(0) + model = model_class(**init_dict).eval() + + # Single-device (unsharded) reference on CPU, computed before TP mutates the weights in place. + with torch.no_grad(): + ref_output = model(**inputs, return_dict=False)[0].float().cpu() + + # Shard across all ranks; the Neuron backend is auto-selected from the mesh device type. + model.enable_parallelism(config=TensorParallelConfig(mesh=tp_mesh)) + model = model.to(device) + torch.neuron.synchronize() + + inputs_on_device = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()} + with torch.no_grad(): + tp_output = model(**inputs_on_device, return_dict=False)[0] + torch.neuron.synchronize() + tp_output = tp_output.float().cpu() + + if rank == 0: + assert tp_output.shape == ref_output.shape, f"shape mismatch: {tp_output.shape} vs {ref_output.shape}" + assert torch.isfinite(tp_output).all(), "TP output contains non-finite values" + max_abs = (tp_output - ref_output).abs().max().item() + denom = ref_output.abs().max().item() + 1e-6 + print( + f"[rank0] tp_size={tp_size} output_shape={tuple(tp_output.shape)} " + f"max_abs_diff={max_abs:.4e} max_rel_diff={max_abs / denom:.4e}" + ) + # Neuron runs matmuls in bf16 internally, so compare with a bf16-level tolerance. A wrong shard + # plan produces grossly different output and is caught comfortably within this bound. + torch.testing.assert_close(tp_output, ref_output, atol=2e-2, rtol=2e-2) + print("[rank0] PASS: Neuron tensor-parallel output matches single-device reference.") + + dist.barrier() + dist.destroy_process_group() + + +if __name__ == "__main__": + try: + main() + except Exception: + traceback.print_exc() + # Ensure a non-zero exit so the launching pytest sees the failure. + os._exit(1) diff --git a/tests/models/transformers/test_models_transformer_flux.py b/tests/models/transformers/test_models_transformer_flux.py index b6a3dfae9e3d..a32fde3a8ffc 100644 --- a/tests/models/transformers/test_models_transformer_flux.py +++ b/tests/models/transformers/test_models_transformer_flux.py @@ -13,6 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os +import subprocess +import sys import tempfile from typing import Any @@ -24,7 +27,7 @@ from diffusers.models.transformers.transformer_flux import FluxIPAdapterAttnProcessor from diffusers.utils.torch_utils import randn_tensor -from ...testing_utils import enable_full_determinism, torch_device +from ...testing_utils import enable_full_determinism, is_tensor_parallel, require_torch_neuron, torch_device from ..testing_utils import ( AttentionBackendTesterMixin, AttentionTesterMixin, @@ -51,6 +54,7 @@ SDNQTesterMixin, SingleFileTesterMixin, TaylorSeerCacheTesterMixin, + TensorParallelTesterMixin, TorchAoCompileTesterMixin, TorchAoTesterMixin, TorchCompileTesterMixin, @@ -157,7 +161,7 @@ def get_init_dict(self) -> dict[str, int | list[int]]: "axes_dims_rope": [4, 4, 8], } - def get_dummy_inputs(self, batch_size: int = 1) -> dict[str, torch.Tensor]: + def get_dummy_inputs(self, batch_size: int = 1, device=torch_device) -> dict[str, torch.Tensor]: height = width = 4 num_latent_channels = 4 num_image_channels = 3 @@ -168,34 +172,34 @@ def get_dummy_inputs(self, batch_size: int = 1) -> dict[str, torch.Tensor]: "hidden_states": randn_tensor( (batch_size, height * width, num_latent_channels), generator=self.generator, - device=torch_device, + device=device, dtype=self.torch_dtype, ), "encoder_hidden_states": randn_tensor( (batch_size, sequence_length, embedding_dim), generator=self.generator, - device=torch_device, + device=device, dtype=self.torch_dtype, ), "pooled_projections": randn_tensor( (batch_size, embedding_dim), generator=self.generator, - device=torch_device, + device=device, dtype=self.torch_dtype, ), "img_ids": randn_tensor( (height * width, num_image_channels), generator=self.generator, - device=torch_device, + device=device, dtype=self.torch_dtype, ), "txt_ids": randn_tensor( (sequence_length, num_image_channels), generator=self.generator, - device=torch_device, + device=device, dtype=self.torch_dtype, ), - "timestep": torch.tensor([1.0]).to(torch_device, self.torch_dtype).expand(batch_size), + "timestep": torch.tensor([1.0]).to(device, self.torch_dtype).expand(batch_size), } @@ -262,6 +266,43 @@ class TestFluxTransformerContextParallelAttnBackends( """Context Parallel inference x attention backends tests for Flux Transformer""" +class TestFluxTransformerTensorParallel(FluxTransformerTesterConfig, TensorParallelTesterMixin): + """Tensor Parallel inference tests for Flux Transformer (CUDA/XPU multi-accelerator).""" + + +def make_neuron_tp_spec(): + """Model spec consumed by the generic Neuron TP worker (`_neuron_tp_worker.py`). + + Returns `(model_class, init_dict, cpu_inputs)`. Defined here so all Flux-specific test data lives in this file + while the worker stays model-agnostic. Reuses the shared tester config so the spec never drifts from the rest of + the Flux tests. + """ + config = FluxTransformerTesterConfig() + return FluxTransformer2DModel, config.get_init_dict(), config.get_dummy_inputs(device="cpu") + + +@is_tensor_parallel +@require_torch_neuron +class TestFluxTransformerTensorParallelNeuron: + """Tensor Parallel inference test for Flux Transformer on AWS Neuron. + + Neuron TP runs through `torchrun` with the `"neuron"` distributed backend, so it cannot use the + `torch.multiprocessing`/NCCL spawn path of `TensorParallelTesterMixin`. This launches the generic worker + with the Flux model spec (`make_neuron_tp_spec`); the worker asserts the sharded output matches a single-device + reference, and the test checks its exit code. + """ + + def test_tensor_parallel_neuron_inference(self): + worker = os.path.join(os.path.dirname(__file__), "_neuron_tp_worker.py") + spec = "tests.models.transformers.test_models_transformer_flux:make_neuron_tp_spec" + cmd = [sys.executable, "-m", "torch.distributed.run", "--nproc_per_node=2", worker, spec] + result = subprocess.run(cmd, capture_output=True, text=True) + assert result.returncode == 0, ( + f"Neuron tensor-parallel worker failed (exit {result.returncode}).\n" + f"--- stdout ---\n{result.stdout}\n--- stderr ---\n{result.stderr}" + ) + + class TestFluxTransformerIPAdapter(FluxTransformerTesterConfig, IPAdapterTesterMixin): """IP Adapter tests for Flux Transformer.""" diff --git a/tests/models/transformers/test_models_transformer_flux2.py b/tests/models/transformers/test_models_transformer_flux2.py index 9546fdb5d969..9e541302b05f 100644 --- a/tests/models/transformers/test_models_transformer_flux2.py +++ b/tests/models/transformers/test_models_transformer_flux2.py @@ -13,6 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os +import subprocess +import sys + import torch from diffusers import Flux2Transformer2DModel @@ -24,7 +28,7 @@ ) from diffusers.utils.torch_utils import randn_tensor -from ...testing_utils import enable_full_determinism, torch_device +from ...testing_utils import enable_full_determinism, is_tensor_parallel, require_torch_neuron, torch_device from ..testing_utils import ( AttentionTesterMixin, BaseModelTesterConfig, @@ -36,6 +40,7 @@ LoraTesterMixin, MemoryTesterMixin, ModelTesterMixin, + TensorParallelTesterMixin, TorchAoCompileTesterMixin, TorchAoTesterMixin, TorchCompileTesterMixin, @@ -90,16 +95,18 @@ def get_init_dict(self) -> dict[str, int | list[int]]: "axes_dims_rope": [4, 4, 4, 4], } - def get_dummy_inputs(self, height: int = 4, width: int = 4, batch_size: int = 1) -> dict[str, torch.Tensor]: + def get_dummy_inputs( + self, height: int = 4, width: int = 4, batch_size: int = 1, device: str = torch_device + ) -> dict[str, torch.Tensor]: num_latent_channels = 4 sequence_length = 48 embedding_dim = 32 hidden_states = randn_tensor( - (batch_size, height * width, num_latent_channels), generator=self.generator, device=torch_device + (batch_size, height * width, num_latent_channels), generator=self.generator, device=device ) encoder_hidden_states = randn_tensor( - (batch_size, sequence_length, embedding_dim), generator=self.generator, device=torch_device + (batch_size, sequence_length, embedding_dim), generator=self.generator, device=device ) t_coords = torch.arange(1) @@ -107,17 +114,17 @@ def get_dummy_inputs(self, height: int = 4, width: int = 4, batch_size: int = 1) w_coords = torch.arange(width) l_coords = torch.arange(1) image_ids = torch.cartesian_prod(t_coords, h_coords, w_coords, l_coords) # [height * width, 4] - image_ids = image_ids.unsqueeze(0).expand(batch_size, -1, -1).to(torch_device) + image_ids = image_ids.unsqueeze(0).expand(batch_size, -1, -1).to(device) text_t_coords = torch.arange(1) text_h_coords = torch.arange(1) text_w_coords = torch.arange(1) text_l_coords = torch.arange(sequence_length) text_ids = torch.cartesian_prod(text_t_coords, text_h_coords, text_w_coords, text_l_coords) - text_ids = text_ids.unsqueeze(0).expand(batch_size, -1, -1).to(torch_device) + text_ids = text_ids.unsqueeze(0).expand(batch_size, -1, -1).to(device) - timestep = torch.tensor([1.0]).to(torch_device).expand(batch_size) - guidance = torch.tensor([1.0]).to(torch_device).expand(batch_size) + timestep = torch.tensor([1.0]).to(device).expand(batch_size) + guidance = torch.tensor([1.0]).to(device).expand(batch_size) return { "hidden_states": hidden_states, @@ -153,6 +160,43 @@ class TestFlux2TransformerContextParallel(Flux2TransformerTesterConfig, ContextP """Context Parallel inference tests for Flux2 Transformer.""" +class TestFlux2TransformerTensorParallel(Flux2TransformerTesterConfig, TensorParallelTesterMixin): + """Tensor Parallel inference tests for Flux2 Transformer (CUDA/XPU multi-accelerator).""" + + +def make_neuron_tp_spec(): + """Model spec consumed by the generic Neuron TP worker (`_neuron_tp_worker.py`). + + Returns `(model_class, init_dict, cpu_inputs)`. Defined here so all Flux2-specific test data lives in this file + while the worker stays model-agnostic. Reuses the shared tester config so the spec never drifts from the rest of + the Flux2 tests. + """ + config = Flux2TransformerTesterConfig() + return Flux2Transformer2DModel, config.get_init_dict(), config.get_dummy_inputs(device="cpu") + + +@is_tensor_parallel +@require_torch_neuron +class TestFlux2TransformerTensorParallelNeuron: + """Tensor Parallel inference test for Flux2 Transformer on AWS Neuron. + + Neuron TP runs through `torchrun` with the `"neuron"` distributed backend, so it cannot use the + `torch.multiprocessing`/NCCL spawn path of `TensorParallelTesterMixin`. This launches the generic worker + with the Flux2 model spec (`make_neuron_tp_spec`); the worker asserts the sharded output matches a single-device + reference, and the test checks its exit code. + """ + + def test_tensor_parallel_neuron_inference(self): + worker = os.path.join(os.path.dirname(__file__), "_neuron_tp_worker.py") + spec = "tests.models.transformers.test_models_transformer_flux2:make_neuron_tp_spec" + cmd = [sys.executable, "-m", "torch.distributed.run", "--nproc_per_node=2", worker, spec] + result = subprocess.run(cmd, capture_output=True, text=True) + assert result.returncode == 0, ( + f"Neuron tensor-parallel worker failed (exit {result.returncode}).\n" + f"--- stdout ---\n{result.stdout}\n--- stderr ---\n{result.stderr}" + ) + + class TestFlux2TransformerLoRA(Flux2TransformerTesterConfig, LoraTesterMixin): """LoRA adapter tests for Flux2 Transformer.""" diff --git a/tests/models/transformers/test_models_transformer_qwenimage.py b/tests/models/transformers/test_models_transformer_qwenimage.py index bed2bda4064f..053fed294fac 100644 --- a/tests/models/transformers/test_models_transformer_qwenimage.py +++ b/tests/models/transformers/test_models_transformer_qwenimage.py @@ -13,6 +13,10 @@ # limitations under the License. +import os +import subprocess +import sys + import pytest import torch @@ -20,7 +24,7 @@ from diffusers.models.transformers.transformer_qwenimage import compute_text_seq_len_from_mask from diffusers.utils.torch_utils import randn_tensor -from ...testing_utils import enable_full_determinism, torch_device +from ...testing_utils import enable_full_determinism, is_tensor_parallel, require_torch_neuron, torch_device from ..testing_utils import ( AttentionBackendTesterMixin, AttentionTesterMixin, @@ -32,6 +36,7 @@ LoraTesterMixin, MemoryTesterMixin, ModelTesterMixin, + TensorParallelTesterMixin, TorchAoTesterMixin, TorchCompileTesterMixin, TrainingTesterMixin, @@ -79,20 +84,20 @@ def get_init_dict(self) -> dict[str, int | list[int]]: "axes_dims_rope": (8, 4, 4), } - def get_dummy_inputs(self, batch_size: int = 1) -> dict[str, torch.Tensor]: + def get_dummy_inputs(self, batch_size: int = 1, device=torch_device) -> dict[str, torch.Tensor]: num_latent_channels = embedding_dim = 16 height = width = 4 sequence_length = 8 vae_scale_factor = 4 hidden_states = randn_tensor( - (batch_size, height * width, num_latent_channels), generator=self.generator, device=torch_device + (batch_size, height * width, num_latent_channels), generator=self.generator, device=device ) encoder_hidden_states = randn_tensor( - (batch_size, sequence_length, embedding_dim), generator=self.generator, device=torch_device + (batch_size, sequence_length, embedding_dim), generator=self.generator, device=device ) - encoder_hidden_states_mask = torch.ones((batch_size, sequence_length)).to(torch_device, torch.long) - timestep = torch.tensor([1.0]).to(torch_device).expand(batch_size) + encoder_hidden_states_mask = torch.ones((batch_size, sequence_length)).to(device, torch.long) + timestep = torch.tensor([1.0]).to(device).expand(batch_size) orig_height = height * 2 * vae_scale_factor orig_width = width * 2 * vae_scale_factor img_shapes = [(1, orig_height // vae_scale_factor // 2, orig_width // vae_scale_factor // 2)] * batch_size @@ -298,6 +303,43 @@ def get_dummy_inputs(self, batch_size: int = 1) -> dict[str, torch.Tensor]: return inputs +class TestQwenImageTransformerTensorParallel(QwenImageTransformerTesterConfig, TensorParallelTesterMixin): + """Tensor Parallel inference tests for QwenImage Transformer (CUDA/XPU multi-accelerator).""" + + +def make_neuron_tp_spec(): + """Model spec consumed by the generic Neuron TP worker (``_neuron_tp_worker.py``). + + Returns ``(model_class, init_dict, cpu_inputs)``. Defined here so all QwenImage-specific test data lives in this + file while the worker stays model-agnostic. Reuses the shared tester config so the spec never drifts from the rest + of the QwenImage tests. + """ + config = QwenImageTransformerTesterConfig() + return QwenImageTransformer2DModel, config.get_init_dict(), config.get_dummy_inputs(device="cpu") + + +@is_tensor_parallel +@require_torch_neuron +class TestQwenImageTransformerTensorParallelNeuron: + """Tensor Parallel inference test for QwenImage Transformer on AWS Neuron. + + Neuron TP runs through ``torchrun`` with the ``"neuron"`` distributed backend, so it cannot use the + ``torch.multiprocessing``/NCCL spawn path of ``TensorParallelTesterMixin``. This launches the generic worker + with the QwenImage model spec (``make_neuron_tp_spec``); the worker asserts the sharded output matches a + single-device reference, and the test checks its exit code. + """ + + def test_tensor_parallel_neuron_inference(self): + worker = os.path.join(os.path.dirname(__file__), "_neuron_tp_worker.py") + spec = "tests.models.transformers.test_models_transformer_qwenimage:make_neuron_tp_spec" + cmd = [sys.executable, "-m", "torch.distributed.run", "--nproc_per_node=2", worker, spec] + result = subprocess.run(cmd, capture_output=True, text=True) + assert result.returncode == 0, ( + f"Neuron tensor-parallel worker failed (exit {result.returncode}).\n" + f"--- stdout ---\n{result.stdout}\n--- stderr ---\n{result.stderr}" + ) + + class TestQwenImageTransformerLoRA(QwenImageTransformerTesterConfig, LoraTesterMixin): """LoRA adapter tests for QwenImage Transformer.""" diff --git a/tests/testing_utils.py b/tests/testing_utils.py index 4e9c553365d9..e88cc829b763 100644 --- a/tests/testing_utils.py +++ b/tests/testing_utils.py @@ -494,6 +494,14 @@ def is_context_parallel(test_case): return pytest.mark.context_parallel(test_case) +def is_tensor_parallel(test_case): + """ + Decorator marking a test as a tensor parallel inference test. These tests can be filtered using: + pytest -m "not tensor_parallel" to skip pytest -m tensor_parallel to run only these tests + """ + return pytest.mark.tensor_parallel(test_case) + + def is_cache(test_case): """ Decorator marking a test as a cache test. These tests can be filtered using: