diff --git a/src/diffusers/hooks/tensor_parallel.py b/src/diffusers/hooks/tensor_parallel.py index ce0c76f97081..e3f4e05dc7bf 100644 --- a/src/diffusers/hooks/tensor_parallel.py +++ b/src/diffusers/hooks/tensor_parallel.py @@ -50,6 +50,20 @@ def __init__(self, blocks: "list[int] | None" = None): self.blocks = blocks +class ReplicatedInputRowwiseParallel: + """Row-wise sharding for a Linear whose input arrives replicated instead of column-sharded. + + Plain `"rowwise"` is the second half of a colwise/rowwise pair, so it expects its input to already be `Shard(-1)` — + which it is when the preceding Linear was colwise-sharded. A Linear that instead reads a replicated activation, + such as a modulation projection off the shared timestep embedding, needs its input sharded on the way in (a local + narrow, no collective) and its partial output all-reduced on the way out. + + Weight and bias shard exactly as for plain `"rowwise"`: the weight over its input columns, the bias replicated and + added after the all-reduce. Use this to shard a large standalone projection whose output must keep the full feature + dimension, where colwise sharding would need an extra all-gather to rebuild it. + """ + + def _packed_blocks(style: "PackedColwiseParallel | PackedRowwiseParallel", module, path: str) -> "list[int]": """The blocks of a packed style: its own `blocks`, else the `_tp_packed_*_blocks` attribute on the Linear.""" attr = "_tp_packed_col_blocks" if isinstance(style, PackedColwiseParallel) else "_tp_packed_row_blocks" @@ -149,7 +163,8 @@ def resolve_tp_shard_specs(model: torch.nn.Module, tp_plan: dict, tp_degree: int if style == "colwise": weight_spec = TPShardSpec(0, [submodule.weight.shape[0]]) bias_spec = weight_spec - elif style == "rowwise": + elif style == "rowwise" or isinstance(style, ReplicatedInputRowwiseParallel): + # Both place the weight the same way; they differ only in the forward input/output hooks. weight_spec = TPShardSpec(1, [submodule.weight.shape[1]]) bias_spec = TPShardSpec(None, None) elif isinstance(style, PackedColwiseParallel): @@ -163,7 +178,8 @@ def resolve_tp_shard_specs(model: torch.nn.Module, tp_plan: dict, tp_degree: int else: raise ValueError( f"Unsupported tensor-parallel style '{style}' for '{path}'. " - f"Expected 'colwise', 'rowwise', PackedColwiseParallel, or PackedRowwiseParallel." + f"Expected 'colwise', 'rowwise', PackedColwiseParallel, PackedRowwiseParallel, or " + f"ReplicatedInputRowwiseParallel." ) specs[f"{path}.weight"] = weight_spec @@ -242,9 +258,9 @@ def _shard_packed_param(param, dim: int, blocks: "list[int]", device_mesh, src_d 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() | }`. Divisibility by the TP - degree is checked earlier, by `resolve_tp_shard_specs`. + Values may be plain strings (`"colwise"` / `"rowwise"`) or `PackedColwiseParallel` / `PackedRowwiseParallel` / + `ReplicatedInputRowwiseParallel` marker instances. Returns `{relative_path: ColwiseParallel() | RowwiseParallel() | + }`. Divisibility by the TP degree is checked earlier, by `resolve_tp_shard_specs`. """ from torch.distributed.tensor import Replicate, distribute_tensor from torch.distributed.tensor.parallel import ColwiseParallel, RowwiseParallel @@ -285,6 +301,11 @@ def _partition_linear_fn(self, name, module, device_mesh): resolved[path] = ColwiseParallel() elif style == "rowwise": resolved[path] = RowwiseParallel() + elif isinstance(style, ReplicatedInputRowwiseParallel): + # `input_layouts=Replicate()` makes `prepare_input` narrow the replicated activation down to this rank's + # columns rather than trusting it to already be `Shard(-1)`; the default would read a full-width tensor + # as if it were one rank's shard. + resolved[path] = RowwiseParallel(input_layouts=Replicate()) elif isinstance(style, PackedColwiseParallel): resolved[path] = _make_packed_col(style) elif isinstance(style, PackedRowwiseParallel): @@ -292,7 +313,8 @@ def _partition_linear_fn(self, name, module, device_mesh): else: raise ValueError( f"Unsupported tensor-parallel style '{style}' for '{path}'. " - f"Expected 'colwise', 'rowwise', PackedColwiseParallel, or PackedRowwiseParallel." + f"Expected 'colwise', 'rowwise', PackedColwiseParallel, PackedRowwiseParallel, or " + f"ReplicatedInputRowwiseParallel." ) return resolved @@ -308,6 +330,7 @@ def _hooks_only_styles(relative_plan: dict) -> dict: targeted module into a `Replicate()` DTensor via a broadcast. Callers should therefore place every planned parameter themselves, and must ensure none is left on `meta` — the broadcast would be issued on a meta tensor. """ + from torch.distributed.tensor import Replicate from torch.distributed.tensor.parallel import ColwiseParallel, RowwiseParallel class _NoPartitionColwise(ColwiseParallel): @@ -324,14 +347,58 @@ def _partition_linear_fn(self, name, module, device_mesh): resolved[path] = _NoPartitionColwise() elif style == "rowwise" or isinstance(style, PackedRowwiseParallel): resolved[path] = _NoPartitionRowwise() + elif isinstance(style, ReplicatedInputRowwiseParallel): + resolved[path] = _NoPartitionRowwise(input_layouts=Replicate()) else: raise ValueError( f"Unsupported tensor-parallel style '{style}' for '{path}'. " - f"Expected 'colwise', 'rowwise', PackedColwiseParallel, or PackedRowwiseParallel." + f"Expected 'colwise', 'rowwise', PackedColwiseParallel, PackedRowwiseParallel, or " + f"ReplicatedInputRowwiseParallel." ) return resolved +def _pre_shard_and_parallelize( + model: torch.nn.Module, + tp_mesh: "torch.distributed.device_mesh.DeviceMesh", + groups: list, + specs: "dict[str, TPShardSpec]", + device: torch.device, +) -> None: + """Slice every planned parameter on CPU, place only this rank's shard on `device`, then register the hooks. + + Unlike the default path this does not broadcast from a single rank, so every rank must already hold the same + weights. Model weights must be on CPU when this is called. + """ + import torch.nn as nn + from torch.distributed.tensor import DTensor, Replicate, Shard + from torch.distributed.tensor.parallel import parallelize_module + + for name, spec in specs.items(): + path, _, param_name = name.rpartition(".") + module = model.get_submodule(path) + param = getattr(module, param_name) + + if spec.dim is None: + # A rowwise bias is added after the all-reduce, so every rank needs the whole vector. + local, placement = param.data, Replicate() + else: + local, placement = _local_shard(param.data, spec.dim, spec.block_sizes, tp_mesh), Shard(spec.dim) + + module.register_parameter( + param_name, + nn.Parameter( + DTensor.from_local(local.to(device), tp_mesh, [placement]), + requires_grad=param.requires_grad, + ), + ) + + # `parallelize_module` is now a no-op for weight distribution (they are already DTensors) but still registers the + # input/output hooks required for the forward pass. + for block, relative_plan in groups: + parallelize_module(block, tp_mesh, _hooks_only_styles(relative_plan)) + + def _tp_degree(tp_config) -> int: """The TP degree, read the same way `_resolve_parallel_config` will, for checks run before the mesh is built.""" return tp_config.mesh.size() if tp_config.mesh is not None else tp_config.tp_degree @@ -498,9 +565,8 @@ def apply_tensor_parallel( 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}.") + logger.debug(f"Applying tensor parallel over {len(groups)} module group(s) on mesh {tp_mesh}.") from torch.distributed.tensor.parallel import parallelize_module @@ -512,10 +578,12 @@ def apply_tensor_parallel( # Also validates every planned shard, so an uneven split raises before any module is sharded. specs = resolve_tp_shard_specs(model, tp_plan, tp_mesh.size()) - if backend == "neuron": - from .tensor_parallel_neuron import _apply_tp_neuron - - _apply_tp_neuron(model, tp_mesh, groups, specs) + if tp_mesh.device_type in ("neuron", "tpu"): + # Pre-shard on CPU instead of letting `parallelize_module` distribute the weights. On Neuron, consecutive + # `reduce_scatter`s of large weights (>= 5120x5120) in one `parallelize_module` call can fail in the runtime. On + # TPU, it would materialize every full weight on each chip first, which runs out of HBM on large models. + device = torch.neuron.current_device() if tp_mesh.device_type == "neuron" else config._device + _pre_shard_and_parallelize(model, tp_mesh, groups, specs, device) return for submodule, relative_plan in groups: diff --git a/src/diffusers/hooks/tensor_parallel_neuron.py b/src/diffusers/hooks/tensor_parallel_neuron.py deleted file mode 100644 index ffcba0973d81..000000000000 --- a/src/diffusers/hooks/tensor_parallel_neuron.py +++ /dev/null @@ -1,75 +0,0 @@ -# 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 and skips the collective for weights, while still registering the required -input/output hooks for the forward pass. - -Only needed for a model that is already in memory. `from_pretrained` with a tensor-parallel `parallel_config` streams -each rank's slice straight off disk into its DTensor, which issues no weight collectives at all and so cannot hit the -bug in the first place. -""" - -import torch -import torch.nn as nn - -from .tensor_parallel import TPShardSpec, _hooks_only_styles, _local_shard - - -def _apply_tp_neuron( - model: nn.Module, - tp_mesh: "torch.distributed.device_mesh.DeviceMesh", - groups: list, - specs: "dict[str, TPShardSpec]", -) -> None: - """Pre-shard the planned parameters via `DTensor.from_local`, then register the forward hooks. - - `groups` and `specs` both come from the model's `_tp_plan` via `diffusers.hooks.tensor_parallel._resolve_tp_plan` / - `resolve_tp_shard_specs`, the same source of truth the generic path uses, so the two backends shard identical - layers. - - Model weights must be on CPU when this is called. - """ - from torch.distributed.tensor import DTensor, Replicate, Shard - from torch.distributed.tensor.parallel import parallelize_module - - device = torch.neuron.current_device() - - for name, spec in specs.items(): - path, _, param_name = name.rpartition(".") - module = model.get_submodule(path) - param = getattr(module, param_name) - - if spec.dim is None: - # A rowwise bias is added after the all-reduce, so every rank needs the whole vector. - local, placement = param.data, Replicate() - else: - local, placement = _local_shard(param.data, spec.dim, spec.block_sizes, tp_mesh), Shard(spec.dim) - - module.register_parameter( - param_name, - nn.Parameter( - DTensor.from_local(local.to(device), tp_mesh, [placement]), - requires_grad=param.requires_grad, - ), - ) - - # `parallelize_module` is now a no-op for weight distribution (they are already DTensors) but still registers the - # input/output hooks required for the forward pass. - for block, relative_plan in groups: - parallelize_module(block, tp_mesh, _hooks_only_styles(relative_plan)) diff --git a/src/diffusers/models/_modeling_parallel.py b/src/diffusers/models/_modeling_parallel.py index 58cbc822f8bf..cf576ab6bb77 100644 --- a/src/diffusers/models/_modeling_parallel.py +++ b/src/diffusers/models/_modeling_parallel.py @@ -200,6 +200,12 @@ class ParallelConfig: """ Configuration for applying different parallelisms. + Both may be set at once. The two are then applied over one device mesh with a dimension each — `("tp", "ring", + "ulysses")`, built by `enable_parallelism` — so their collectives stay in separate process groups. This is what + lets a model too large for one device (TP shards the weights) also run a sequence too long for one device's + attention scratchpad (CP splits the sequence): TP alone cannot divide the sequence, and CP alone cannot divide the + weights. + Args: context_parallel_config (`ContextParallelConfig`, *optional*): Configuration for context parallelism. @@ -216,12 +222,17 @@ class ParallelConfig: _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: + if self.context_parallel_config is None and self.tensor_parallel_config is 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`." + "A `ParallelConfig` must specify at least one of `context_parallel_config` or " + "`tensor_parallel_config`." ) + @property + def _is_combined(self) -> bool: + """Whether both context and tensor parallelism are requested, i.e. they must share one mesh.""" + return self.context_parallel_config is not None and self.tensor_parallel_config is not None + def setup( self, rank: int, @@ -234,10 +245,32 @@ def setup( self._world_size = world_size self._device = device self._mesh = mesh + + # Context and tensor parallelism compose because they cut the model along different axes: CP splits the + # sequence (and, under Ulysses, trades sequence for heads inside attention) while TP shards the Linear + # weights. Composing them means giving each its own mesh *dimension*, so that every collective one issues + # stays inside its own process group: TP's all-reduce must not reach a rank holding a different sequence + # chunk, and CP's sequence all-gather must not reach a rank holding a different weight shard. + cp_mesh = tp_mesh = mesh + if self._is_combined: + dim_names = mesh.mesh_dim_names if mesh is not None else None + missing = [d for d in ("ring", "ulysses", "tp") if dim_names is None or d not in dim_names] + if missing: + raise ValueError( + f"Combining context and tensor parallelism requires a device mesh with 'ring', 'ulysses' and " + f"'tp' dimensions, but {missing} {'is' if len(missing) == 1 else 'are'} missing (got " + f"{dim_names}). `enable_parallelism` builds this mesh from `ring_degree`/`ulysses_degree`/" + f"`tp_degree`; if you build it yourself, set `mesh=` on one of the two configs." + ) + # `ContextParallelConfig.setup` slices ("ring", "ulysses") off whatever mesh it is handed, so it takes + # the full mesh. TP instead needs its own 1-D submesh: `parallelize_module` shards a weight over every + # rank of the mesh it is given, so handing it the 3-D mesh would shard over the CP ranks as well. + tp_mesh = mesh["tp"] + if self.context_parallel_config is not None: - self.context_parallel_config.setup(rank, world_size, device, mesh) + self.context_parallel_config.setup(rank, world_size, device, cp_mesh) if self.tensor_parallel_config is not None: - self.tensor_parallel_config.setup(rank, world_size, device, mesh) + self.tensor_parallel_config.setup(rank, world_size, device, tp_mesh) @dataclass(frozen=True) diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index a6881e971edf..5051853be46a 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -1544,8 +1544,24 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None if tp_shard_specs is not None: # The weights are already sharded, so this only registers the forward hooks. `_parallel_config` # was recorded by `_resolve_parallel_config` before loading. + from torch.distributed.tensor import DTensor + from ..hooks.tensor_parallel import apply_tensor_parallel + # Non-persistent buffers are absent from both the state dict and the checkpoint, and + # `init_empty_weights` leaves them as real CPU tensors, so move them across explicitly. + tp_device = ( + torch.neuron.current_device() if tp_config._mesh.device_type == "neuron" else tp_config._device + ) + for name, buffer in model.named_buffers(): + if buffer.device != tp_device and not isinstance(buffer, DTensor): + module_path, _, buffer_name = name.rpartition(".") + module = model.get_submodule(module_path) if module_path else model + module._buffers[buffer_name] = buffer.to(tp_device) + + # Same order as `enable_parallelism`: context parallel hooks first, then tensor parallel ones. + if parallel_config.context_parallel_config is not None: + model._apply_context_parallel(parallel_config) apply_tensor_parallel(model, tp_config, cls._tp_plan, weights_already_sharded=True) else: model.enable_parallelism(config=parallel_config) @@ -1720,15 +1736,28 @@ def _resolve_parallel_config( device = torch.device(device_type, rank % device_module.device_count()) mesh = None - if config.context_parallel_config is not None: - cp_config = config.context_parallel_config + cp_config = config.context_parallel_config + tp_config = config.tensor_parallel_config + if cp_config is not None and tp_config is not None: + # One mesh, one dimension per parallelism, so `ParallelConfig.setup` can hand each config its own + # submesh. "tp" goes first so the CP dimensions vary fastest: on Neuron, `all_to_all` (Ulysses) rejects + # strided replica groups. TP's all-reduce accepts them, so TP is the axis that gives up contiguity. + mesh = ( + cp_config.mesh + or tp_config.mesh + or torch.distributed.device_mesh.init_device_mesh( + device_type=device_type, + mesh_shape=(tp_config.tp_degree, cp_config.ring_degree, cp_config.ulysses_degree), + mesh_dim_names=("tp", "ring", "ulysses"), + ) + ) + elif cp_config is not None: mesh = cp_config.mesh or torch.distributed.device_mesh.init_device_mesh( device_type=device_type, 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 + elif tp_config is not None: mesh = tp_config.mesh or torch.distributed.device_mesh.init_device_mesh( device_type=device_type, mesh_shape=(tp_config.tp_degree,), @@ -1737,6 +1766,32 @@ def _resolve_parallel_config( # `config.setup()` records the mesh resolved above onto the config; see `ParallelConfig.setup`. config.setup(rank, world_size, device, mesh=mesh) + + # Validate the combination up front — after `setup`, which resolves `_tp_degree` from the mesh, but before + # anything is recorded on the model or applied to it. CP hooks are applied before TP, so a check left to the + # TP branch would raise on a model already carrying CP hooks: half-parallelised, and not usable. + if cp_config is not None and tp_config is not None: + tp_degree = tp_config._tp_degree + requested = cp_config.ring_degree * cp_config.ulysses_degree * tp_degree + if requested > world_size: + raise ValueError( + f"Combining context and tensor parallelism needs `ring_degree` ({cp_config.ring_degree}) * " + f"`ulysses_degree` ({cp_config.ulysses_degree}) * `tp_degree` ({tp_degree}) = {requested} " + f"devices, which exceeds the world size ({world_size})." + ) + num_heads = getattr(self.config, "num_attention_heads", None) + divisor = tp_degree * cp_config.ulysses_degree + if num_heads is not None and not cp_config.ulysses_anything and num_heads % divisor != 0: + # Ulysses trades sequence for heads inside attention (`SeqAllToAllDim` scatters over dim 2), and it + # only ever sees the heads TP left on this rank, so the count has to survive both splits. + raise ValueError( + f"Combining tensor parallelism (`tp_degree`={tp_degree}) with Ulysses context parallelism " + f"(`ulysses_degree`={cp_config.ulysses_degree}) requires the number of attention heads " + f"({num_heads}) to be divisible by their product ({divisor}): TP shards the heads first, and " + f"Ulysses splits what is left on each rank. Pass `ulysses_anything=True` to pad the head " + f"dimension instead, or pick degrees whose product divides {num_heads}." + ) + self._parallel_config = config return config @@ -1750,11 +1805,6 @@ def enable_parallelism( "`enable_parallelism` is an experimental feature. The API may change in the future and breaking changes may be introduced at any time without warning." ) - from ..hooks.context_parallel import apply_context_parallel - from .attention import AttentionModuleMixin - from .attention_dispatch import AttentionBackendName, _AttentionBackendRegistry - from .attention_processor import Attention, MochiAttention - if self._parallel_config is not None: raise RuntimeError( f"Parallelism is already applied to this {self.__class__.__name__}. `enable_parallelism` cannot be " @@ -1781,59 +1831,73 @@ def enable_parallelism( config = self._resolve_parallel_config(config) + if config.context_parallel_config is not None: + self._apply_context_parallel(config, cp_plan) + + if config.tensor_parallel_config is not None: + from ..hooks.tensor_parallel import apply_tensor_parallel + + apply_tensor_parallel(self, config.tensor_parallel_config, self._tp_plan) + + def _apply_context_parallel( + self, config: ParallelConfig, cp_plan: dict[str, ContextParallelModelPlan] | None = None + ) -> None: + """Apply `config.context_parallel_config` to a model whose `config` was already resolved. + + Shared by `enable_parallelism` and the sharded `from_pretrained` path, which resolves the config before + loading. + """ + from ..hooks.context_parallel import apply_context_parallel + from .attention import AttentionModuleMixin + from .attention_dispatch import AttentionBackendName, _AttentionBackendRegistry + from .attention_processor import Attention, MochiAttention + attention_classes = (Attention, MochiAttention, AttentionModuleMixin) - if config.context_parallel_config is not None: - for module in self.modules(): - if not isinstance(module, attention_classes): - continue + for module in self.modules(): + if not isinstance(module, attention_classes): + continue - processor = module.processor - if processor is None or not hasattr(processor, "_attention_backend"): - continue + processor = module.processor + if processor is None or not hasattr(processor, "_attention_backend"): + continue - attention_backend = processor._attention_backend - if attention_backend is None: - attention_backend, _ = _AttentionBackendRegistry.get_active_backend() - else: - attention_backend = AttentionBackendName(attention_backend) + attention_backend = processor._attention_backend + if attention_backend is None: + attention_backend, _ = _AttentionBackendRegistry.get_active_backend() + else: + attention_backend = AttentionBackendName(attention_backend) - if not _AttentionBackendRegistry._is_context_parallel_available(attention_backend): - compatible_backends = sorted(_AttentionBackendRegistry._supports_context_parallel) - raise ValueError( - f"Context parallelism is enabled but the attention processor '{processor.__class__.__name__}' " - f"is using backend '{attention_backend.value}' which does not support context parallelism. " - f"Please set a compatible attention backend: {compatible_backends} using `model.set_attention_backend()` before " - f"calling `model.enable_parallelism()`." - ) + if not _AttentionBackendRegistry._is_context_parallel_available(attention_backend): + compatible_backends = sorted(_AttentionBackendRegistry._supports_context_parallel) + raise ValueError( + f"Context parallelism is enabled but the attention processor '{processor.__class__.__name__}' " + f"is using backend '{attention_backend.value}' which does not support context parallelism. " + f"Please set a compatible attention backend: {compatible_backends} using `model.set_attention_backend()` before " + f"calling `model.enable_parallelism()`." + ) - # All modules use the same attention processor and backend. We don't need to - # iterate over all modules after checking the first processor - break + # All modules use the same attention processor and backend. We don't need to + # iterate over all modules after checking the first processor + break # 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." - ) - 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: - from ..hooks.tensor_parallel import apply_tensor_parallel + 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 - apply_tensor_parallel(self, config.tensor_parallel_config, self._tp_plan) + 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." + ) + cp_plan = cp_plan if cp_plan is not None else self._cp_plan + apply_context_parallel(self, config.context_parallel_config, cp_plan) @classmethod def _load_pretrained_model( diff --git a/src/diffusers/models/transformers/transformer_minimax_h3.py b/src/diffusers/models/transformers/transformer_minimax_h3.py index 95ee1f2ac511..0f9e38bdb965 100644 --- a/src/diffusers/models/transformers/transformer_minimax_h3.py +++ b/src/diffusers/models/transformers/transformer_minimax_h3.py @@ -19,6 +19,7 @@ import torch.nn as nn from ...configuration_utils import ConfigMixin, register_to_config +from ...hooks.tensor_parallel import PackedColwiseParallel, ReplicatedInputRowwiseParallel from ...loaders import FromOriginalModelMixin, PeftAdapterMixin from ...utils import BaseOutput, apply_lora_scale, logging from .._modeling_parallel import ContextParallelInput, ContextParallelOutput @@ -179,9 +180,11 @@ def __call__( key = attn.to_k(hidden_states) value = attn.to_v(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) @@ -451,6 +454,45 @@ class MiniMaxH3Transformer3DModel( "audio_proj_out", "rope", ] + # Tensor-parallel plan: how each block's Linears shard across the TP mesh. Q/K/V and the attention output + # are unfused, so they are plain "colwise"/"rowwise" (torch's ColwiseParallel / RowwiseParallel). The one + # packed projection is the SwiGLU input `ff.net.0.proj`, a single Linear producing `[value; gate]` in equal + # halves, hence PackedColwiseParallel([1, 1]) so each half is sharded independently. + # + # `adaln_proj.linear` is the one projection that has to keep its full output width: the six modulation + # parameters it produces scale and shift the full hidden dim of the packed sequence, which is already all-reduced + # by the time they are applied. Sharding it colwise would need an all-gather to rebuild that width, so it is + # sharded `ReplicatedInputRowwiseParallel` instead — over its `time_embed_dim` input, reading the replicated + # `temb` and all-reducing the result. It is worth the collective: at `6 * hidden_size * MINIMAX_H3_MODALITY_NUM` + # outputs per block it is ~40% of the denoiser's weights, and leaving it replicated puts the per-rank floor above + # a single device's memory at any TP degree. The all-reduce itself is over `num_timesteps * 6 * hidden_size * + # MINIMAX_H3_MODALITY_NUM` elements, i.e. hundreds of KB, once per block per step. + # + # Intentionally absent, i.e. replicated on every rank: the RMSNorms (`norm1`, `norm2`, + # `token_refiner.final_norm`, `norm_out.norm`); the QK-norms, which apply over `head_dim` after the heads are + # already split; `norm_out.linear`, which indexes the full hidden dim as `adaln_proj` does but is a single + # `2 * hidden_size` projection rather than one per block, so sharding it would buy nothing; and the patch/text + # embedders and the two output heads. + # + # `attn.to_qkv` is deliberately not listed: it only exists after `fuse_projections()`, and the plan is + # resolved by attribute lookup, so an unconditional entry would break the ordinary unfused model. + _tp_plan = { + # denoiser block stack + "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.*.ff.net.0.proj": PackedColwiseParallel([1, 1]), + "transformer_blocks.*.ff.net.2": "rowwise", + "transformer_blocks.*.adaln_proj.linear": ReplicatedInputRowwiseParallel(), + # the token-refiner blocks are the same attention + SwiGLU FFN, minus AdaLN and rotary + "token_refiner.refiner_blocks.*.attn.to_q": "colwise", + "token_refiner.refiner_blocks.*.attn.to_k": "colwise", + "token_refiner.refiner_blocks.*.attn.to_v": "colwise", + "token_refiner.refiner_blocks.*.attn.to_out.0": "rowwise", + "token_refiner.refiner_blocks.*.ff.net.0.proj": PackedColwiseParallel([1, 1]), + "token_refiner.refiner_blocks.*.ff.net.2": "rowwise", + } # Context parallelism shards the packed sequence, so the split cannot happen on the inputs of `forward`: the rows # of the three modalities are scattered into the packed buffer with sequence-wide indices, which only address the # full sequence. The split therefore happens once the buffer is built, at the first block, and everything that is diff --git a/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py b/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py index 94395a7c3ebe..bb1c1cd54b4f 100644 --- a/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py +++ b/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py @@ -1228,6 +1228,13 @@ def build_row_timesteps( Returns: `tuple[torch.Tensor, torch.Tensor]`: the distinct timesteps, sorted, and the index of every row into them. """ + # The row plan is built on CPU and the caller moves the finished pair to the denoiser's device: `unique` + # has a data-dependent output shape, which an accelerator would rather not trace. The layout step hands the + # index tensors over already on that device, so bring them back for the scatter below — indexing a CPU + # tensor with an accelerator one does not work. + video_indices = video_indices.cpu() + audio_indices = audio_indices.cpu() + sequence_length = int(video_indices.numel() + audio_indices.numel() + num_text_tokens) row_timesteps = torch.full((sequence_length,), video_timestep, dtype=torch.float32) row_timesteps[video_indices[:num_condition_video_rows]] = condition_video_timestep diff --git a/tests/models/test_parallelism_combined.py b/tests/models/test_parallelism_combined.py new file mode 100644 index 000000000000..4a404a06125e --- /dev/null +++ b/tests/models/test_parallelism_combined.py @@ -0,0 +1,164 @@ +# 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. + +"""Tests for combining context parallelism with tensor parallelism on one device mesh. + +These cover the wiring rather than the numerics: that `ParallelConfig` hands each parallelism its own submesh, and +that the resulting process groups are the intended factorisation of the world. They run on CPU over `gloo`, so they +need no accelerator — `gloo` has no `all_to_all`, so an end-to-end Ulysses forward pass cannot run here; that is +covered by `ContextAndTensorParallelTesterMixin` in `testing_utils/parallelism.py`, which needs four accelerators. +""" + +import os +import socket + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from diffusers.models._modeling_parallel import ContextParallelConfig, ParallelConfig, TensorParallelConfig + + +def _find_free_port(): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("", 0)) + s.listen(1) + return s.getsockname()[1] + + +def _mesh_factorization_worker(rank, world_size, master_port, ulysses_degree, tp_degree, return_dict): + """Build a combined mesh, run `ParallelConfig.setup`, and report the process groups each parallelism landed on.""" + 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) + + dist.init_process_group(backend="gloo", rank=rank, world_size=world_size) + + mesh = torch.distributed.device_mesh.init_device_mesh( + "cpu", mesh_shape=(tp_degree, 1, ulysses_degree), mesh_dim_names=("tp", "ring", "ulysses") + ) + config = ParallelConfig( + context_parallel_config=ContextParallelConfig(ulysses_degree=ulysses_degree), + tensor_parallel_config=TensorParallelConfig(tp_degree=tp_degree), + ) + config.setup(rank, world_size, torch.device("cpu"), mesh=mesh) + + cp_config = config.context_parallel_config + tp_config = config.tensor_parallel_config + return_dict[rank] = { + "status": "success", + "cp_ranks": dist.get_process_group_ranks(cp_config._flattened_mesh.get_group()), + "ulysses_local_rank": cp_config._ulysses_local_rank, + "tp_ranks": dist.get_process_group_ranks(tp_config._mesh.get_group()), + "tp_local_rank": tp_config._mesh.get_local_rank(), + "tp_degree": tp_config._tp_degree, + } + except Exception as e: # noqa: BLE001 — surfaced as a test failure by the caller + return_dict[rank] = {"status": "error", "error": f"{type(e).__name__}: {e}"} + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +def _missing_tp_dim_worker(rank, world_size, master_port, return_dict): + """A combined config handed a CP-only mesh must be rejected, naming the missing 'tp' dimension.""" + 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) + + dist.init_process_group(backend="gloo", rank=rank, world_size=world_size) + + mesh = torch.distributed.device_mesh.init_device_mesh( + "cpu", mesh_shape=(1, world_size), mesh_dim_names=("ring", "ulysses") + ) + config = ParallelConfig( + context_parallel_config=ContextParallelConfig(ulysses_degree=world_size), + tensor_parallel_config=TensorParallelConfig(tp_degree=1), + ) + try: + config.setup(rank, world_size, torch.device("cpu"), mesh=mesh) + except ValueError as e: + return_dict[rank] = {"status": "raised", "message": str(e)} + else: + return_dict[rank] = {"status": "no_raise"} + except Exception as e: # noqa: BLE001 + return_dict[rank] = {"status": "error", "error": f"{type(e).__name__}: {e}"} + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +def _spawn(worker, world_size, *args): + if not dist.is_available(): + pytest.skip("torch.distributed is not available.") + manager = mp.Manager() + return_dict = manager.dict() + mp.spawn(worker, args=(world_size, _find_free_port(), *args, return_dict), nprocs=world_size, join=True) + return return_dict + + +class TestCombinedParallelConfig: + """`ParallelConfig` accepting both parallelisms at once, and splitting the mesh between them.""" + + def test_both_configs_are_accepted(self): + # Combining the two used to raise outright; the mesh dimension per parallelism is what makes it work. + config = ParallelConfig( + context_parallel_config=ContextParallelConfig(ulysses_degree=2), + tensor_parallel_config=TensorParallelConfig(tp_degree=2), + ) + assert config._is_combined + assert config.context_parallel_config.ulysses_degree == 2 + assert config.tensor_parallel_config.tp_degree == 2 + + def test_single_parallelism_is_not_combined(self): + assert not ParallelConfig(context_parallel_config=ContextParallelConfig(ulysses_degree=2))._is_combined + assert not ParallelConfig(tensor_parallel_config=TensorParallelConfig(tp_degree=2))._is_combined + + def test_mesh_is_factored_between_the_two(self): + """Each parallelism must end up on its own process group: TP within a CP chunk, CP across the TP groups. + + With the default `mesh_shape=(tp=2, ring=1, ulysses=2)` over four ranks, "ulysses" is the fastest-varying + dimension, so ranks {0,1} and {2,3} are the CP groups and {0,2} / {1,3} are the TP groups. A TP all-reduce + must not reach a rank holding a different sequence chunk, and a CP collective must not reach a rank holding a + different weight shard — this asserts exactly that split. + """ + world_size = 4 + results = _spawn(_mesh_factorization_worker, world_size, 2, 2) + + for rank in range(world_size): + assert results[rank]["status"] == "success", results[rank].get("error") + + assert [results[r]["tp_ranks"] for r in range(4)] == [[0, 2], [1, 3], [0, 2], [1, 3]] + assert [results[r]["cp_ranks"] for r in range(4)] == [[0, 1], [0, 1], [2, 3], [2, 3]] + assert [results[r]["ulysses_local_rank"] for r in range(4)] == [0, 1, 0, 1] + # The TP shard index is the coordinate inside the TP group, which stops tracking the global rank as soon as + # the mesh has more than one dimension. Weight sharding keys off this, so it is the value that matters. + assert [results[r]["tp_local_rank"] for r in range(4)] == [0, 0, 1, 1] + assert all(results[r]["tp_degree"] == 2 for r in range(4)) + + def test_mesh_without_tp_dimension_is_rejected(self): + world_size = 2 + results = _spawn(_missing_tp_dim_worker, world_size) + + for rank in range(world_size): + assert results[rank]["status"] == "raised", ( + f"rank {rank}: expected a ValueError for a mesh with no 'tp' dimension, got {results[rank]}" + ) + assert "tp" in results[rank]["message"] diff --git a/tests/models/testing_utils/__init__.py b/tests/models/testing_utils/__init__.py index ee0ab21e8719..6db810342821 100644 --- a/tests/models/testing_utils/__init__.py +++ b/tests/models/testing_utils/__init__.py @@ -21,8 +21,10 @@ from .lora import LoraHotSwappingForModelTesterMixin, LoraTesterMixin from .memory import CPUOffloadTesterMixin, GroupOffloadTesterMixin, LayerwiseCastingTesterMixin, MemoryTesterMixin from .parallelism import ( + ContextAndTensorParallelTesterMixin, ContextParallelAttentionBackendsTesterMixin, ContextParallelTesterMixin, + HybridParallelTesterMixin, TensorParallelTesterMixin, TensorParallelTPUTesterMixin, ) @@ -66,8 +68,10 @@ "BitsAndBytesConfigMixin", "BitsAndBytesTesterMixin", "CacheTesterMixin", + "ContextAndTensorParallelTesterMixin", "ContextParallelTesterMixin", "ContextParallelAttentionBackendsTesterMixin", + "HybridParallelTesterMixin", "TensorParallelTesterMixin", "TensorParallelTPUTesterMixin", "CPUOffloadTesterMixin", diff --git a/tests/models/testing_utils/parallelism.py b/tests/models/testing_utils/parallelism.py index eca6bff5d271..d6c57b84c587 100644 --- a/tests/models/testing_utils/parallelism.py +++ b/tests/models/testing_utils/parallelism.py @@ -21,7 +21,7 @@ import torch.distributed as dist import torch.multiprocessing as mp -from diffusers.models._modeling_parallel import ContextParallelConfig, TensorParallelConfig +from diffusers.models._modeling_parallel import ContextParallelConfig, ParallelConfig, TensorParallelConfig from diffusers.models.attention_dispatch import AttentionBackendName, _AttentionBackendRegistry from diffusers.utils import SAFE_WEIGHTS_INDEX_NAME @@ -558,6 +558,150 @@ def test_tensor_parallel_tpu_inference(self, atol=1e-3, rtol=1e-3): torch.testing.assert_close(ref_output, tp_output, atol=atol, rtol=rtol) +def _context_and_tensor_parallel_worker( + rank, world_size, master_port, model_class, init_dict, cp_dict, tp_degree, inputs_dict, return_dict, state_dict +): + """Worker for combined context + tensor parallel inference. + + Both configs go into one `ParallelConfig`, which puts them on one mesh with a dimension each. The result should + still match the single-device reference: TP is mathematically equivalent to the unsharded model, and CP splits the + sequence and gathers it back, so neither changes the function being computed. + """ + 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"]) + dist.init_process_group(backend=device_config["backend"], rank=rank, world_size=world_size) + + device_config["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()} + + model.enable_parallelism( + config=ParallelConfig( + context_parallel_config=ContextParallelConfig(**cp_dict), + tensor_parallel_config=TensorParallelConfig(tp_degree=tp_degree), + ) + ) + + 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) + return_dict["output"] = output.float().cpu().tolist() + + except Exception as e: + if rank == 0: + return_dict["status"] = "error" + return_dict["error"] = f"{type(e).__name__}: {e}" + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +@is_context_parallel +@is_tensor_parallel +@require_torch_multi_accelerator +class ContextAndTensorParallelTesterMixin: + """Context and tensor parallelism enabled together, over one device mesh. + + The two are orthogonal — CP cuts the sequence, TP cuts the weights — so composing them should leave the computed + function unchanged. Both axes are covered: Ulysses, which trades sequence for heads inside attention, and ring, + which does not touch the head dimension at all. + + Needs `cp_degree` x `tp_degree` accelerators (four by default). Head-count requirements differ per CP type, so a + model whose dummy config has too few heads skips rather than failing. + """ + + cp_degree = 2 + tp_degree = 2 + + @pytest.mark.parametrize("cp_type", ["ulysses_degree", "ring_degree"], ids=["ulysses", "ring"]) + def test_context_and_tensor_parallel_inference(self, cp_type, 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.") + if getattr(self.model_class, "_cp_plan", None) is None: + pytest.skip("Model does not define a `_cp_plan` for context parallel inference.") + + if cp_type == "ring_degree": + active_backend, _ = _AttentionBackendRegistry.get_active_backend() + if active_backend == AttentionBackendName.NATIVE: + pytest.skip("Ring attention is not supported with the native attention backend.") + + world_size = self.cp_degree * self.tp_degree + device_module = DEVICE_CONFIG.get(torch_device, DEVICE_CONFIG["cuda"])["module"] + if device_module.device_count() < world_size: + pytest.skip(f"Combined CP x TP needs {world_size} accelerators, found {device_module.device_count()}.") + + init_dict = self.get_init_dict() + num_heads = init_dict.get("num_attention_heads") + # TP shards the heads; Ulysses then splits what TP left on each rank, so the two multiply. Ring leaves the + # head dimension alone, so only the TP degree has to divide the head count. + required_head_multiple = self.tp_degree * (self.cp_degree if cp_type == "ulysses_degree" else 1) + if num_heads is not None and num_heads % required_head_multiple != 0: + pytest.skip( + f"`num_attention_heads` ({num_heads}) is not divisible by {required_head_multiple}, required for " + f"{cp_type.removesuffix('_degree')}={self.cp_degree} x tp_degree={self.tp_degree}." + ) + + inputs_dict = self.get_dummy_inputs(batch_size=batch_size) + + # Single-device reference, captured before anything is sharded. + 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() + + inputs_dict = {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in inputs_dict.items()} + + manager = mp.Manager() + return_dict = manager.dict() + mp.spawn( + _context_and_tensor_parallel_worker, + args=( + world_size, + _find_free_port(), + self.model_class, + init_dict, + {cp_type: self.cp_degree}, + self.tp_degree, + inputs_dict, + return_dict, + state_dict, + ), + nprocs=world_size, + join=True, + ) + + assert return_dict.get("status") == "success", ( + f"Combined context + tensor parallel inference failed: {return_dict.get('error', 'Unknown error')}" + ) + + combined_output = torch.tensor(return_dict["output"]) + assert list(ref_output.shape) == return_dict["output_shape"] + # Two sets of collectives reorder the summation on top of the sharded matmuls, so the tolerance matches the + # TP-only test rather than the tighter CP-only one. + torch.testing.assert_close(ref_output, combined_output, atol=1e-3, rtol=1e-3) + + @pytest.mark.parametrize("cp_type", ["ulysses_degree", "ring_degree"], ids=["ulysses", "ring"]) + def test_context_and_tensor_parallel_batch_inputs(self, cp_type): + self.test_context_and_tensor_parallel_inference(cp_type, batch_size=2) + + @is_context_parallel @require_torch_multi_accelerator class ContextParallelTesterMixin: @@ -823,3 +967,161 @@ def test_context_parallel_attn_backend_inference(self, cp_type, attention_backen cp_output = torch.tensor(return_dict["output"], dtype=ref_output.dtype) torch.testing.assert_close(ref_output, cp_output, atol=1e-2, rtol=1e-2) + + +def _hybrid_parallel_worker( + rank, + world_size, + master_port, + model_class, + init_dict, + cp_dict, + tp_degree, + inputs_dict, + return_dict, + state_dict, + checkpoint_dir=None, +): + """Worker function for combined tensor + context parallel inference testing. + + Both parallelisms are requested through a single `ParallelConfig`, which shares one device mesh between them, so + each rank holds `1 / tp_degree` of every sharded weight *and* `1 / (ring_degree * ulysses_degree)` of the + sequence. Rank 0 reports its output so the caller can compare it against a single-device reference: the + composition is mathematically equivalent to the unsharded model up to floating-point reduction order. + + With `checkpoint_dir`, the model is loaded with `from_pretrained(..., parallel_config=...)` instead, which shards + while reading the checkpoint and must still apply the context parallel hooks. + """ + 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}") + + parallel_config = ParallelConfig( + tensor_parallel_config=TensorParallelConfig(tp_degree=tp_degree), + context_parallel_config=ContextParallelConfig(**cp_dict), + ) + if checkpoint_dir is not None: + model = model_class.from_pretrained(checkpoint_dir, parallel_config=parallel_config).eval() + cp_hooks = [ + name + for module in model.modules() + if getattr(module, "_diffusers_hook", None) is not None + for name in module._diffusers_hook.hooks + if name.startswith("cp_") + ] + assert cp_hooks, "`from_pretrained(parallel_config=...)` applied no context parallel hooks." + else: + model = model_class(**init_dict) + model.load_state_dict(state_dict) + model.to(device) + model.eval() + model.enable_parallelism(config=parallel_config) + + inputs_on_device = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in inputs_dict.items()} + + 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_context_parallel +@is_tensor_parallel +@require_torch_multi_accelerator +class HybridParallelTesterMixin: + """Tensor parallelism and context parallelism together, from one `ParallelConfig`. + + Needs `tp_degree * ulysses_degree` accelerators (4 at the degrees used here), so it skips on a 2-device runner. + """ + + @pytest.mark.parametrize("load", ["enable_parallelism", "from_pretrained"]) + def test_hybrid_parallel_inference(self, tmp_path, load, batch_size: int = 1): + if not torch.distributed.is_available(): + pytest.skip("torch.distributed is not available.") + + for plan in ("_tp_plan", "_cp_plan"): + if getattr(self.model_class, plan, None) is None: + pytest.skip(f"Model does not define a `{plan}`, which hybrid parallelism requires.") + + tp_degree, ulysses_degree = 2, 2 + world_size = tp_degree * ulysses_degree + device_count = DEVICE_CONFIG.get(torch_device, DEVICE_CONFIG["cuda"])["module"].device_count() + if device_count < world_size: + pytest.skip( + f"tp_degree={tp_degree} x ulysses_degree={ulysses_degree} needs {world_size} accelerators, " + f"found {device_count}." + ) + + init_dict = self.get_init_dict() + num_heads = init_dict.get("num_attention_heads") + # Each rank keeps `num_heads // tp_degree` heads, which Ulysses splits again. + if num_heads is not None and num_heads % world_size != 0: + pytest.skip(f"`num_attention_heads` ({num_heads}) is not divisible by {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() + + inputs_dict = {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in inputs_dict.items()} + + checkpoint_dir = None + if load == "from_pretrained": + checkpoint_dir = str(tmp_path / "checkpoint") + model.save_pretrained(checkpoint_dir) + + master_port = _find_free_port() + manager = mp.Manager() + return_dict = manager.dict() + + mp.spawn( + _hybrid_parallel_worker, + args=( + world_size, + master_port, + self.model_class, + init_dict, + {"ulysses_degree": ulysses_degree}, + tp_degree, + inputs_dict, + return_dict, + state_dict, + checkpoint_dir, + ), + nprocs=world_size, + join=True, + ) + + assert return_dict.get("status") == "success", ( + f"Hybrid parallel inference failed: {return_dict.get('error', 'Unknown error')}" + ) + + output = torch.tensor(return_dict["output"]) + # Sharded matmuls plus the Ulysses all-to-all reorder the summation, hence the tolerance. + torch.testing.assert_close(ref_output, output, atol=1e-3, rtol=1e-3) diff --git a/tests/models/transformers/_neuron_hybrid_worker.py b/tests/models/transformers/_neuron_hybrid_worker.py new file mode 100644 index 000000000000..0c0e44896beb --- /dev/null +++ b/tests/models/transformers/_neuron_hybrid_worker.py @@ -0,0 +1,140 @@ +# 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 x context-parallel output matches its reference. + +The counterpart of `_neuron_tp_worker.py` for the two parallelisms composed in a single `ParallelConfig`. Same +contract: the model under test is supplied as a `module:function` spec reference on the command line, and the +referenced factory returns `(model_class, init_dict, inputs)` with CPU tensors. + + torchrun --nproc_per_node=8 _neuron_hybrid_worker.py \\ + tests.models.transformers.test_models_transformer_flux:make_neuron_hybrid_spec + +`tp_degree` and `ulysses_degree` are read from `TP_DEGREE` / `ULYSSES_DEGREE` (defaults 2 and 4, whose product is +the launched world size). `ulysses_degree` cannot be 2 on Neuron: its all-to-all only accepts group sizes of 4, 8, +16 or multiples of 32. + +No mesh is passed, so this also exercises the default mesh layout that `enable_parallelism` builds for the combined +case -- which matters on Neuron, where the all-to-all Ulysses depends on rejects strided replica groups and so the +context-parallel dimensions have to vary fastest. + +Exit code 0 means the composed path is numerically equivalent to the unsharded model; non-zero means failure. +""" + +import argparse +import importlib +import os +import sys +import tempfile +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 diffusers import ContextParallelConfig, ParallelConfig, TensorParallelConfig + + +def main(): + parser = argparse.ArgumentParser(description="Neuron tensor-parallel x context-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)() + + tp_degree = int(os.environ.get("TP_DEGREE", "2")) + ulysses_degree = int(os.environ.get("ULYSSES_DEGREE", "4")) + + dist.init_process_group(backend="neuron") + rank = dist.get_rank() + world_size = dist.get_world_size() + device = torch.neuron.current_device() + + if tp_degree * ulysses_degree != world_size: + raise ValueError( + f"tp_degree ({tp_degree}) x ulysses_degree ({ulysses_degree}) must equal the world size ({world_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 the shard plan mutates the weights in place. + with torch.no_grad(): + ref_output = model(**inputs, return_dict=False)[0].float().cpu() + + parallel_config = ParallelConfig( + tensor_parallel_config=TensorParallelConfig(tp_degree=tp_degree), + context_parallel_config=ContextParallelConfig(ulysses_degree=ulysses_degree), + ) + if os.environ.get("LOAD_MODE", "enable_parallelism") == "from_pretrained": + # The sharded-load path: every rank reads only its own slice of the checkpoint. + # Every rank holds the same weights, so each saves and reloads its own copy. + with tempfile.TemporaryDirectory() as tmpdir: + model.save_pretrained(tmpdir) + model = model_class.from_pretrained(tmpdir, parallel_config=parallel_config).eval() + cp_hooks = [ + name + for module in model.modules() + if getattr(module, "_diffusers_hook", None) is not None + for name in module._diffusers_hook.hooks + if name.startswith("cp_") + ] + assert cp_hooks, "from_pretrained(parallel_config=TP+CP) applied no context parallel hooks" + else: + model.enable_parallelism(config=parallel_config) + 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(): + output = model(**inputs_on_device, return_dict=False)[0] + torch.neuron.synchronize() + output = output.float().cpu() + + if rank == 0: + assert output.shape == ref_output.shape, f"shape mismatch: {output.shape} vs {ref_output.shape}" + assert torch.isfinite(output).all(), "output contains non-finite values" + max_abs = (output - ref_output).abs().max().item() + denom = ref_output.abs().max().item() + 1e-6 + print( + f"[rank0] tp_degree={tp_degree} ulysses_degree={ulysses_degree} " + f"output_shape={tuple(output.shape)} 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, as `_neuron_tp_worker` + # does. A wrong shard plan or a mis-ordered mesh produces grossly different output and is caught well + # inside this bound. + torch.testing.assert_close(output, ref_output, atol=2e-2, rtol=2e-2) + print("[rank0] PASS: Neuron hybrid-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 04542253c7f6..b8ad9cdc6a83 100644 --- a/tests/models/transformers/test_models_transformer_flux.py +++ b/tests/models/transformers/test_models_transformer_flux.py @@ -27,7 +27,13 @@ from diffusers.models.transformers.transformer_flux import FluxIPAdapterAttnProcessor from diffusers.utils.torch_utils import randn_tensor -from ...testing_utils import enable_full_determinism, is_tensor_parallel, require_torch_neuron, torch_device +from ...testing_utils import ( + enable_full_determinism, + is_context_parallel, + is_tensor_parallel, + require_torch_neuron, + torch_device, +) from ..testing_utils import ( AttentionBackendTesterMixin, AttentionTesterMixin, @@ -40,6 +46,7 @@ FirstBlockCacheTesterMixin, GGUFCompileTesterMixin, GGUFTesterMixin, + HybridParallelTesterMixin, IPAdapterTesterMixin, LoraHotSwappingForModelTesterMixin, LoraTesterMixin, @@ -285,6 +292,23 @@ def get_init_dict(self): return {**super().get_init_dict(), "num_attention_heads": self.tp_world_size} +class TestFluxTransformerHybridParallel(FluxTransformerTesterConfig, HybridParallelTesterMixin): + """Tensor Parallel x Context Parallel inference tests for Flux Transformer (needs 4 accelerators).""" + + +def make_neuron_hybrid_spec(): + """Model spec consumed by the generic Neuron hybrid worker (`_neuron_hybrid_worker.py`). + + Same contract as `make_neuron_tp_spec`, but `num_attention_heads` is raised to 8 so the head count survives + being divided twice: `tp_degree=2` leaves 4 heads per rank and `ulysses_degree=4` splits those into 1 each. + (`ulysses_degree` cannot be 2 on Neuron, whose all-to-all only accepts group sizes of 4, 8, 16 or multiples + of 32.) + """ + config = FluxTransformerTesterConfig() + init_dict = config.get_init_dict() | {"num_attention_heads": 8} + return FluxTransformer2DModel, init_dict, config.get_dummy_inputs(device="cpu") + + def make_neuron_tp_spec(): """Model spec consumed by the generic Neuron TP worker (`_neuron_tp_worker.py`). @@ -318,6 +342,29 @@ def test_tensor_parallel_neuron_inference(self): ) +@is_context_parallel +@is_tensor_parallel +@require_torch_neuron +class TestFluxTransformerHybridParallelNeuron: + """Tensor Parallel x Context Parallel inference test for Flux Transformer on AWS Neuron. + + Same launching pattern as `TestFluxTransformerTensorParallelNeuron`: Neuron needs `torchrun` with the + `"neuron"` distributed backend, so it cannot use the `torch.multiprocessing`/NCCL spawn path of + `HybridParallelTesterMixin`. Runs at `tp_degree=2 x ulysses_degree=4`, i.e. 8 ranks. + """ + + @pytest.mark.parametrize("load", ["enable_parallelism", "from_pretrained"]) + def test_hybrid_parallel_neuron_inference(self, load): + worker = os.path.join(os.path.dirname(__file__), "_neuron_hybrid_worker.py") + spec = "tests.models.transformers.test_models_transformer_flux:make_neuron_hybrid_spec" + cmd = [sys.executable, "-m", "torch.distributed.run", "--nproc_per_node=8", worker, spec] + result = subprocess.run(cmd, capture_output=True, text=True, env={**os.environ, "LOAD_MODE": load}) + assert result.returncode == 0, ( + f"Neuron hybrid-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_minimax_h3.py b/tests/models/transformers/test_models_transformer_minimax_h3.py index 4165f9b5392d..7c2eaad617f5 100644 --- a/tests/models/transformers/test_models_transformer_minimax_h3.py +++ b/tests/models/transformers/test_models_transformer_minimax_h3.py @@ -13,13 +13,17 @@ # 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 MiniMaxH3Transformer3DModel from diffusers.models.transformers.transformer_minimax_h3 import MiniMaxH3TransformerOutput 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, @@ -28,6 +32,7 @@ MemoryTesterMixin, ModelTesterMixin, SingleFileTesterMixin, + TensorParallelTesterMixin, TorchCompileTesterMixin, TrainingTesterMixin, ) @@ -85,7 +90,7 @@ def get_init_dict(self) -> dict: "rope_freq_dim": 2, } - def get_packed_layout(self, num_video_tokens: int = NUM_VIDEO_TOKENS) -> dict: + def get_packed_layout(self, num_video_tokens: int = NUM_VIDEO_TOKENS, device: str | torch.device = None) -> dict: r""" Build the structural arguments of one packed sequence. @@ -93,29 +98,33 @@ def get_packed_layout(self, num_video_tokens: int = NUM_VIDEO_TOKENS) -> dict: modality and its noise level, and hands over the `(t, h, w)` grid plus the three index tensors. The layout here mirrors what the pipelines pack, with two distinct timesteps so the `(timestep, modality)` AdaLN table is addressed on more than one row. + + `device` defaults to the test device; the Neuron TP spec asks for CPU because its worker builds the model on + CPU and moves it only after sharding. """ + device = torch_device if device is None else device sequence_length = NUM_TEXT_TOKENS + NUM_AUDIO_TOKENS + num_video_tokens - text_indices = torch.arange(NUM_TEXT_TOKENS, device=torch_device) - audio_indices = torch.arange(NUM_TEXT_TOKENS, NUM_TEXT_TOKENS + NUM_AUDIO_TOKENS, device=torch_device) - video_indices = torch.arange(NUM_TEXT_TOKENS + NUM_AUDIO_TOKENS, sequence_length, device=torch_device) + text_indices = torch.arange(NUM_TEXT_TOKENS, device=device) + audio_indices = torch.arange(NUM_TEXT_TOKENS, NUM_TEXT_TOKENS + NUM_AUDIO_TOKENS, device=device) + video_indices = torch.arange(NUM_TEXT_TOKENS + NUM_AUDIO_TOKENS, sequence_length, device=device) # 0 = video, 1 = text, 2 = audio. - token_tags = torch.empty(sequence_length, dtype=torch.long, device=torch_device) + token_tags = torch.empty(sequence_length, dtype=torch.long, device=device) token_tags[text_indices] = 1 token_tags[audio_indices] = 2 token_tags[video_indices] = 0 # The conditioning-free rows share the video timestep; the audio rows step down their own schedule. - timestep_indices = torch.zeros(sequence_length, dtype=torch.long, device=torch_device) + timestep_indices = torch.zeros(sequence_length, dtype=torch.long, device=device) timestep_indices[audio_indices] = 1 - position_ids = torch.zeros(sequence_length, 3, dtype=torch.float32, device=torch_device) - position_ids[:, 0] = torch.arange(sequence_length, dtype=torch.float32, device=torch_device) - position_ids[video_indices, 1] = torch.arange(num_video_tokens, dtype=torch.float32, device=torch_device) % 4 - position_ids[video_indices, 2] = torch.arange(num_video_tokens, dtype=torch.float32, device=torch_device) % 2 + position_ids = torch.zeros(sequence_length, 3, dtype=torch.float32, device=device) + position_ids[:, 0] = torch.arange(sequence_length, dtype=torch.float32, device=device) + position_ids[video_indices, 1] = torch.arange(num_video_tokens, dtype=torch.float32, device=device) % 4 + position_ids[video_indices, 2] = torch.arange(num_video_tokens, dtype=torch.float32, device=device) % 2 return { - "timestep": torch.tensor([0.7, 0.3], device=torch_device), + "timestep": torch.tensor([0.7, 0.3], device=device), "timestep_indices": timestep_indices, "token_tags": token_tags, "position_ids": position_ids, @@ -124,7 +133,10 @@ def get_packed_layout(self, num_video_tokens: int = NUM_VIDEO_TOKENS) -> dict: "text_indices": text_indices, } - def get_dummy_inputs(self, num_video_tokens: int = NUM_VIDEO_TOKENS, batch_size: int = 2) -> dict: + def get_dummy_inputs( + self, num_video_tokens: int = NUM_VIDEO_TOKENS, batch_size: int = 2, device: str | torch.device = None + ) -> dict: + device = torch_device if device is None else device generator = self.generator init_dict = self.get_init_dict() patch_size = init_dict["patch_size"] @@ -132,17 +144,17 @@ def get_dummy_inputs(self, num_video_tokens: int = NUM_VIDEO_TOKENS, batch_size: return { "hidden_states": randn_tensor( - (batch_size, num_video_tokens, video_patch_dim), generator=generator, device=torch_device + (batch_size, num_video_tokens, video_patch_dim), generator=generator, device=device ), "audio_hidden_states": randn_tensor( (batch_size, NUM_AUDIO_TOKENS, init_dict["audio_in_channels"]), generator=generator, - device=torch_device, + device=device, ), "encoder_hidden_states": randn_tensor( - (batch_size, NUM_TEXT_TOKENS, init_dict["text_dim"]), generator=generator, device=torch_device + (batch_size, NUM_TEXT_TOKENS, init_dict["text_dim"]), generator=generator, device=device ), - **self.get_packed_layout(num_video_tokens), + **self.get_packed_layout(num_video_tokens, device=device), } @@ -206,3 +218,40 @@ def pretrained_model_name_or_path(self): @property def pretrained_model_kwargs(self): return {"subfolder": "transformer"} + + +class TestMiniMaxH3TransformerTensorParallel(MiniMaxH3TransformerTesterConfig, TensorParallelTesterMixin): + """Tensor Parallel inference tests for the MiniMax-H3 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 MiniMax-H3-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 MiniMax-H3 tests. + """ + config = MiniMaxH3TransformerTesterConfig() + return MiniMaxH3Transformer3DModel, config.get_init_dict(), config.get_dummy_inputs(device="cpu") + + +@is_tensor_parallel +@require_torch_neuron +class TestMiniMaxH3TransformerTensorParallelNeuron: + """Tensor Parallel inference test for the MiniMax-H3 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 MiniMax-H3 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_minimax_h3: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}" + ) diff --git a/tests/models/transformers/test_models_transformer_qwenimage.py b/tests/models/transformers/test_models_transformer_qwenimage.py index 6210aecff87f..10d36c8155ba 100644 --- a/tests/models/transformers/test_models_transformer_qwenimage.py +++ b/tests/models/transformers/test_models_transformer_qwenimage.py @@ -30,6 +30,7 @@ AttentionTesterMixin, BaseModelTesterConfig, BitsAndBytesTesterMixin, + ContextAndTensorParallelTesterMixin, ContextParallelAttentionBackendsTesterMixin, ContextParallelTesterMixin, LoraHotSwappingForModelTesterMixin, @@ -320,6 +321,25 @@ def test_tensor_parallel_tpu_inference(self): super().test_tensor_parallel_tpu_inference(atol=2e-2, rtol=2e-2) +class TestQwenImageTransformerContextAndTensorParallel( + QwenImageTransformerTesterConfig, ContextAndTensorParallelTesterMixin +): + """Context Parallel x Tensor Parallel inference tests for QwenImage Transformer (4 accelerators). + + QwenImage is the model this runs on because its dummy config has four attention heads, and Ulysses splits the + heads TP already sharded — so `ulysses_degree` x `tp_degree` has to divide the head count. + """ + + def get_dummy_inputs(self, batch_size: int = 1) -> dict[str, torch.Tensor]: + inputs = super().get_dummy_inputs(batch_size=batch_size) + encoder_hidden_states_mask = inputs["encoder_hidden_states_mask"] + encoder_hidden_states_mask[:, 1] = 0 + encoder_hidden_states_mask[:, 3] = 0 + encoder_hidden_states_mask[:, 5:] = 0 + inputs["encoder_hidden_states_mask"] = encoder_hidden_states_mask + return inputs + + def make_neuron_tp_spec(): """Model spec consumed by the generic Neuron TP worker (``_neuron_tp_worker.py``).