From 764e7c8a85e1a066443b849b54cc7f73e77975cd Mon Sep 17 00:00:00 2001 From: JingyaHuang Date: Fri, 21 Aug 2026 11:08:58 +0000 Subject: [PATCH 01/19] Add tensor-parallel support for MiniMax-H3 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Shards `MiniMaxH3Transformer3DModel` across devices, following the plan already established for Flux1/Flux2/Qwen-Image. Validated on Trainium at TP=2 and TP=8. - `_tp_plan` with twelve entries: the same six shapes for the 50 denoiser blocks and for the two token-refiner blocks, which are the same attention + SwiGLU FFN minus AdaLN and rotary. Q/K/V and the attention output are unfused, so they are plain colwise/rowwise; the SwiGLU input `ff.net.0.proj` is one Linear producing `[value; gate]` in equal halves and takes PackedColwiseParallel([1, 1]). - The attention processor reshaped by the config head count, `unflatten(-1, (attn.heads, -1))`, which mis-splits under sharding: each rank holds `inner_dim / tp_degree` columns, so this yields `head_dim / tp_degree` per head instead of `heads / tp_degree` heads. Reshape by the fixed `attn.head_dim` instead and let `-1` absorb the head count, as Flux does. Numerically identical unsharded, since `inner_dim == heads * head_dim`. - Norms, QK-norms (head_dim-shaped, applied after the head split), AdaLN modulation and the patch/text embedders and output heads stay replicated. `attn.to_qkv` is deliberately not in the plan: it exists only after `fuse_projections()`, and the plan is resolved by attribute lookup. No RoPE change was needed — unlike Qwen-Image, H3's rotary is already real sin/cos and already broadcasts over the head axis. Tests mirror the Flux2/Qwen-Image layout: the CUDA/XPU `TensorParallelTesterMixin` class, a `make_neuron_tp_spec()` factory, and a Neuron launcher that shells out to the model-agnostic `_neuron_tp_worker.py`. `get_dummy_inputs` and `get_packed_layout` take an optional `device` so the Neuron spec can ask for CPU tensors, since its worker shards on CPU and moves to device after. --- .../transformers/transformer_minimax_h3.py | 37 ++++++++- .../test_models_transformer_minimax_h3.py | 83 +++++++++++++++---- 2 files changed, 100 insertions(+), 20 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_minimax_h3.py b/src/diffusers/models/transformers/transformer_minimax_h3.py index 95ee1f2ac511..9d816aed2a15 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 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,34 @@ 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. + # + # 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; the AdaLN modulation (`adaln_proj.linear`, `norm_out.linear`), which indexes the full hidden + # dim; 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", + # 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/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}" + ) From 58d9b51309cb70c6541f07c3bba30c1e103cf366 Mon Sep 17 00:00:00 2001 From: JingyaHuang Date: Fri, 21 Aug 2026 13:45:24 +0000 Subject: [PATCH 02/19] Shard MiniMax-H3's adaln_proj to fit tensor parallelism on one device MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `transformer_blocks.*.adaln_proj.linear` was left replicated on every rank, and at `[96768, 2688]` bf16 per block it is 24.23 GiB of the denoiser's 61.73 GiB — about 40%. That made the per-rank floor 24.40 GiB of weights (plus 5.13 GiB for the two VAEs) regardless of TP degree, so MiniMax-H3 could not fit a 24 GiB NeuronCore at *any* valid TP: 34.20 GiB/rank at TP=8, and still 30.20 GiB/rank at TP=56. Raising TP only divided the 60% that already sharded. (TP=16 is not an option either — 56 attention heads.) Shard it rowwise, over the `time_embed_dim` input, rather than colwise: the six modulation parameters scale and shift the *full* hidden dim of a sequence that is already all-reduced by the time they are applied, so a colwise split would need an all-gather to rebuild that width. Rowwise keeps the output full-width, leaving the module's `view`/`chunk` untouched, and all-reduces a few hundred KB per block per step. Plain `"rowwise"` could not be reused. It is normally the second half of a colwise/rowwise pair, so it defaults to `input_layouts=Shard(-1)` and would read the full-width `temb` as if it were one rank's shard. Hence `ReplicatedInputRowwiseParallel`: input narrowed locally on the way in (no collective), partial output all-reduced on the way out, bias replicated and added after the reduce. It is wired into `_styles`, `_hooks_only_styles` — the path the Neuron backend takes, since `_apply_tp_neuron` pre-shards on CPU and then registers hooks only — and `resolve_tp_shard_specs`. Replicated weights drop from 24.40 GiB to 0.15 GiB, putting TP=8 at 7.84 GiB of transformer plus 5.13 GiB of VAEs, i.e. 12.97 GiB/rank against a 24 GiB budget. Verified on CPU/gloo that both the generic and the pre-sharded hooks-only path shard the weight on its input dim and match a replicated reference to 3.6e-7. Co-Authored-By: Claude Opus 5 (1M context) --- src/diffusers/hooks/tensor_parallel.py | 40 +++++++++++++++---- .../transformers/transformer_minimax_h3.py | 17 ++++++-- 2 files changed, 47 insertions(+), 10 deletions(-) diff --git a/src/diffusers/hooks/tensor_parallel.py b/src/diffusers/hooks/tensor_parallel.py index ce0c76f97081..41d67e098681 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,10 +347,13 @@ 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 diff --git a/src/diffusers/models/transformers/transformer_minimax_h3.py b/src/diffusers/models/transformers/transformer_minimax_h3.py index 9d816aed2a15..0f9e38bdb965 100644 --- a/src/diffusers/models/transformers/transformer_minimax_h3.py +++ b/src/diffusers/models/transformers/transformer_minimax_h3.py @@ -19,7 +19,7 @@ import torch.nn as nn from ...configuration_utils import ConfigMixin, register_to_config -from ...hooks.tensor_parallel import PackedColwiseParallel +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 @@ -459,10 +459,20 @@ class MiniMaxH3Transformer3DModel( # 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; the AdaLN modulation (`adaln_proj.linear`, `norm_out.linear`), which indexes the full hidden - # dim; and the patch/text embedders and the two output heads. + # 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. @@ -474,6 +484,7 @@ class MiniMaxH3Transformer3DModel( "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", From c766b888ae7d85d09ab366b24d00aa070bc20061 Mon Sep 17 00:00:00 2001 From: JingyaHuang Date: Fri, 21 Aug 2026 14:32:06 +0000 Subject: [PATCH 03/19] Build MiniMax-H3's row timestep plan on CPU, as its caller expects MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `build_row_timesteps` allocates `row_timesteps` with `torch.full` and no device, then scatters into it at `video_indices` / `audio_indices`. The layout step hands those index tensors over already on the execution device, and indexing a CPU tensor with an accelerator one is an error — on Neuron it surfaces as "Non-scalar tensor arg0 is on cpu device, expected neuron", and on CUDA it would raise "indices should be either on cpu or on the same device". CPU is the right place for this to run, not the accelerator: `torch.unique` has a data-dependent output shape, which is precisely what a tracing backend cannot handle, and the caller already moves the finished `(timestep, timestep_indices)` pair to the device itself. So bring the two index tensors back to CPU for the scatter rather than allocating `row_timesteps` on their device. Only reachable once the denoiser is actually on an accelerator while the pipeline's execution device resolves there too, which is why it went unnoticed. Co-Authored-By: Claude Opus 5 (1M context) --- .../modular_pipelines/minimax_h3/before_denoise.py | 7 +++++++ 1 file changed, 7 insertions(+) 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 From 4162784267a815757d982295da86c89e2b903f66 Mon Sep 17 00:00:00 2001 From: JingyaHuang Date: Thu, 27 Aug 2026 13:20:19 +0000 Subject: [PATCH 04/19] feat:combine tp+cp (cherry picked from commit a489f03e4e9888c0abb996f45d38885f03772ae4) --- src/diffusers/models/_modeling_parallel.py | 41 ++++- src/diffusers/models/modeling_utils.py | 49 +++++- tests/models/test_parallelism_combined.py | 164 ++++++++++++++++++ tests/models/testing_utils/__init__.py | 2 + tests/models/testing_utils/parallelism.py | 146 +++++++++++++++- .../test_models_transformer_qwenimage.py | 20 +++ 6 files changed, 409 insertions(+), 13 deletions(-) create mode 100644 tests/models/test_parallelism_combined.py diff --git a/src/diffusers/models/_modeling_parallel.py b/src/diffusers/models/_modeling_parallel.py index 58cbc822f8bf..1b61c9186d6e 100644 --- a/src/diffusers/models/_modeling_parallel.py +++ b/src/diffusers/models/_modeling_parallel.py @@ -200,6 +200,11 @@ 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 — `("ring", "ulysses", + "tp")`, 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. @@ -215,12 +220,10 @@ class ParallelConfig: _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`." - ) + @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, @@ -234,10 +237,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..268226d918af 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -1720,15 +1720,30 @@ 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 last, which makes TP ranks adjacent: its all-reduce fires twice per block, more + # often than the CP collectives, so it is the one that wants the closest devices. The CP dimensions are + # then strided, which is fine for them — on Neuron only `all_to_all` (Ulysses) constrains its replica + # groups to be block-aligned; `all_gather` (ring) and `all_reduce` (TP) accept any grouping. + mesh = ( + cp_config.mesh + or tp_config.mesh + or torch.distributed.device_mesh.init_device_mesh( + device_type=device_type, + mesh_shape=(cp_config.ring_degree, cp_config.ulysses_degree, tp_config.tp_degree), + mesh_dim_names=("ring", "ulysses", "tp"), + ) + ) + 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 +1752,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 diff --git a/tests/models/test_parallelism_combined.py b/tests/models/test_parallelism_combined.py new file mode 100644 index 000000000000..76b814c73cca --- /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=(1, ulysses_degree, tp_degree), mesh_dim_names=("ring", "ulysses", "tp") + ) + 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 `mesh_shape=(ring=1, ulysses=2, tp=2)` over four ranks, "tp" is the fastest-varying dimension, so ranks + {0,1} and {2,3} are the TP groups and {0,2} / {1,3} are the CP 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, 1], [0, 1], [2, 3], [2, 3]] + assert [results[r]["cp_ranks"] for r in range(4)] == [[0, 2], [1, 3], [0, 2], [1, 3]] + assert [results[r]["ulysses_local_rank"] for r in range(4)] == [0, 0, 1, 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, 1, 0, 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..ae6cce8964ca 100644 --- a/tests/models/testing_utils/__init__.py +++ b/tests/models/testing_utils/__init__.py @@ -21,6 +21,7 @@ from .lora import LoraHotSwappingForModelTesterMixin, LoraTesterMixin from .memory import CPUOffloadTesterMixin, GroupOffloadTesterMixin, LayerwiseCastingTesterMixin, MemoryTesterMixin from .parallelism import ( + ContextAndTensorParallelTesterMixin, ContextParallelAttentionBackendsTesterMixin, ContextParallelTesterMixin, TensorParallelTesterMixin, @@ -66,6 +67,7 @@ "BitsAndBytesConfigMixin", "BitsAndBytesTesterMixin", "CacheTesterMixin", + "ContextAndTensorParallelTesterMixin", "ContextParallelTesterMixin", "ContextParallelAttentionBackendsTesterMixin", "TensorParallelTesterMixin", diff --git a/tests/models/testing_utils/parallelism.py b/tests/models/testing_utils/parallelism.py index eca6bff5d271..152ba205542a 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: 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``). From 7f649f94e699e3ccb2f59328d23b85ffa420ca2d Mon Sep 17 00:00:00 2001 From: whn09 Date: Mon, 7 Sep 2026 03:08:20 +0000 Subject: [PATCH 05/19] Allow tensor parallelism and context parallelism in one ParallelConfig `ParallelConfig.__post_init__` rejects a config carrying both a `context_parallel_config` and a `tensor_parallel_config`, so a model can only use one form of parallelism at a time. That caps some models well below the hardware they run on: MiniMax-H3 has 56 attention heads, so `tp_degree` cannot exceed 8, and on a 64-accelerator host tensor parallelism alone leaves 56 of them idle even though the model already ships a `_cp_plan`. Both configs already accept a caller-supplied `mesh`, documented as being for "combining CP with other parallelism strategies that share the same mesh", so the intent was there; what was missing was building that shared mesh and splitting it between the two. * `__post_init__` now only rejects a config that carries neither. * `ParallelConfig.setup` hands each config its own dimensions: CP keeps the whole mesh, since its own `setup` selects "ring" and "ulysses" out of it, while TP gets the "tp" dimension alone because it reads its degree from `mesh.size()`. * `enable_parallelism` builds one ("tp", "ring", "ulysses") mesh when both are requested, with the context-parallel dimensions varying fastest. That order is a default rather than a universal optimum, and the comment says so: pass `mesh=` on either config to choose the layout yourself. * The Neuron tensor-parallel backend used the global rank as its shard index. That equals the rank's coordinate on the TP mesh only while that mesh spans the whole world, so sharing a mesh made it index past the end of the weight; the generic backend beside it already reads the coordinate. Measured on a trn2.48xlarge (MiniMax-H3, 1344x768x124f, 30 steps): TP=8 alone is 9.285 s/step on 8 cores, TP=4 x ulysses=4 is 5.66 s/step on 16, and TP=8 x ulysses=4 is 3.941 s/step on 32 -- 2.36x faster than the widest configuration reachable before this change, with output that is visually equivalent to the tensor-parallel-only run. Tests: a `HybridParallelTesterMixin` next to the existing TP and CP mixins, wired into the Flux transformer tests, plus a Neuron `torchrun` worker following the `_neuron_tp_worker.py` convention already in the tree. The Neuron test runs at tp_degree=2 x ulysses_degree=4 and passes on a trn2.48xlarge (max_abs_diff 4.2e-05 against the single-device reference). Co-Authored-By: Claude Opus 5 --- src/diffusers/models/_modeling_parallel.py | 7 + src/diffusers/models/modeling_utils.py | 1 - tests/models/testing_utils/__init__.py | 2 + tests/models/testing_utils/parallelism.py | 129 ++++++++++++++++++ .../transformers/_neuron_hybrid_worker.py | 125 +++++++++++++++++ .../test_models_transformer_flux.py | 48 ++++++- 6 files changed, 310 insertions(+), 2 deletions(-) create mode 100644 tests/models/transformers/_neuron_hybrid_worker.py diff --git a/src/diffusers/models/_modeling_parallel.py b/src/diffusers/models/_modeling_parallel.py index 1b61c9186d6e..0f1da9f3b353 100644 --- a/src/diffusers/models/_modeling_parallel.py +++ b/src/diffusers/models/_modeling_parallel.py @@ -220,6 +220,13 @@ class ParallelConfig: _device: torch.device = None _mesh: torch.distributed.device_mesh.DeviceMesh = None + def __post_init__(self): + if self.context_parallel_config is None and self.tensor_parallel_config is None: + raise ValueError( + "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.""" diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index 268226d918af..a1477589944e 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -1990,7 +1990,6 @@ def _load_pretrained_model( low_cpu_mem_usage=low_cpu_mem_usage, disable_mmap=disable_mmap, ) - if is_parallel_loading_enabled: offload_index, state_dict_index, _mismatched_keys, _error_msgs = load_fn(resolved_model_file) error_msgs += _error_msgs diff --git a/tests/models/testing_utils/__init__.py b/tests/models/testing_utils/__init__.py index ae6cce8964ca..6db810342821 100644 --- a/tests/models/testing_utils/__init__.py +++ b/tests/models/testing_utils/__init__.py @@ -24,6 +24,7 @@ ContextAndTensorParallelTesterMixin, ContextParallelAttentionBackendsTesterMixin, ContextParallelTesterMixin, + HybridParallelTesterMixin, TensorParallelTesterMixin, TensorParallelTPUTesterMixin, ) @@ -70,6 +71,7 @@ "ContextAndTensorParallelTesterMixin", "ContextParallelTesterMixin", "ContextParallelAttentionBackendsTesterMixin", + "HybridParallelTesterMixin", "TensorParallelTesterMixin", "TensorParallelTPUTesterMixin", "CPUOffloadTesterMixin", diff --git a/tests/models/testing_utils/parallelism.py b/tests/models/testing_utils/parallelism.py index 152ba205542a..ef8927d328f5 100644 --- a/tests/models/testing_utils/parallelism.py +++ b/tests/models/testing_utils/parallelism.py @@ -967,3 +967,132 @@ 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 +): + """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. + """ + 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()} + + model.enable_parallelism( + config=ParallelConfig( + tensor_parallel_config=TensorParallelConfig(tp_degree=tp_degree), + context_parallel_config=ContextParallelConfig(**cp_dict), + ) + ) + + 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. + """ + + def test_hybrid_parallel_inference(self, 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()} + + 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, + ), + 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..1c054af44104 --- /dev/null +++ b/tests/models/transformers/_neuron_hybrid_worker.py @@ -0,0 +1,125 @@ +# 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 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() + + model.enable_parallelism( + config=ParallelConfig( + tensor_parallel_config=TensorParallelConfig(tp_degree=tp_degree), + context_parallel_config=ContextParallelConfig(ulysses_degree=ulysses_degree), + ) + ) + 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..8b36e579d833 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,28 @@ 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. + """ + + def test_hybrid_parallel_neuron_inference(self): + 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) + 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.""" From 673a3f10d8c08dde50bc19f61d67643f5677e745 Mon Sep 17 00:00:00 2001 From: Jingya HUANG Date: Wed, 30 Sep 2026 14:01:57 +0000 Subject: [PATCH 06/19] feat: validate the support on TPU --- src/diffusers/hooks/tensor_parallel.py | 49 ++++++++++++++++++- src/diffusers/hooks/tensor_parallel_neuron.py | 31 +----------- src/diffusers/models/modeling_utils.py | 13 +++++ 3 files changed, 63 insertions(+), 30 deletions(-) diff --git a/src/diffusers/hooks/tensor_parallel.py b/src/diffusers/hooks/tensor_parallel.py index 41d67e098681..ce8abd844ab5 100644 --- a/src/diffusers/hooks/tensor_parallel.py +++ b/src/diffusers/hooks/tensor_parallel.py @@ -358,6 +358,47 @@ def _partition_linear_fn(self, name, module, device_mesh): 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 @@ -524,7 +565,7 @@ 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" + backend = tp_mesh.device_type if tp_mesh.device_type in ("neuron", "tpu") 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}.") @@ -544,5 +585,11 @@ def apply_tensor_parallel( _apply_tp_neuron(model, tp_mesh, groups, specs) return + if backend == "tpu": + # `parallelize_module` would materialize every full weight on each chip before scattering it, which runs out + # of HBM on models larger than one chip. Pre-sharding on CPU keeps only this rank's slice on the device. + _pre_shard_and_parallelize(model, tp_mesh, groups, specs, config._device) + return + 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 index ffcba0973d81..a87f28b1edb1 100644 --- a/src/diffusers/hooks/tensor_parallel_neuron.py +++ b/src/diffusers/hooks/tensor_parallel_neuron.py @@ -28,7 +28,7 @@ import torch import torch.nn as nn -from .tensor_parallel import TPShardSpec, _hooks_only_styles, _local_shard +from .tensor_parallel import TPShardSpec, _pre_shard_and_parallelize def _apply_tp_neuron( @@ -45,31 +45,4 @@ def _apply_tp_neuron( 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)) + _pre_shard_and_parallelize(model, tp_mesh, groups, specs, torch.neuron.current_device()) diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index a1477589944e..262e5a8ae699 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -1544,8 +1544,21 @@ 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) + apply_tensor_parallel(model, tp_config, cls._tp_plan, weights_already_sharded=True) else: model.enable_parallelism(config=parallel_config) From 3db653d80fe37330319e4210786e5c438b31d7b9 Mon Sep 17 00:00:00 2001 From: JingyaHuang Date: Thu, 8 Oct 2026 14:12:07 +0000 Subject: [PATCH 07/19] fix: put tp first in the combined TP+CP mesh --- src/diffusers/models/_modeling_parallel.py | 9 +++++---- src/diffusers/models/modeling_utils.py | 10 ++++------ tests/models/test_parallelism_combined.py | 18 +++++++++--------- 3 files changed, 18 insertions(+), 19 deletions(-) diff --git a/src/diffusers/models/_modeling_parallel.py b/src/diffusers/models/_modeling_parallel.py index 0f1da9f3b353..cf576ab6bb77 100644 --- a/src/diffusers/models/_modeling_parallel.py +++ b/src/diffusers/models/_modeling_parallel.py @@ -200,10 +200,11 @@ 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 — `("ring", "ulysses", - "tp")`, 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. + 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*): diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index 262e5a8ae699..332ac65864ac 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -1737,17 +1737,15 @@ def _resolve_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 last, which makes TP ranks adjacent: its all-reduce fires twice per block, more - # often than the CP collectives, so it is the one that wants the closest devices. The CP dimensions are - # then strided, which is fine for them — on Neuron only `all_to_all` (Ulysses) constrains its replica - # groups to be block-aligned; `all_gather` (ring) and `all_reduce` (TP) accept any grouping. + # 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=(cp_config.ring_degree, cp_config.ulysses_degree, tp_config.tp_degree), - mesh_dim_names=("ring", "ulysses", "tp"), + 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: diff --git a/tests/models/test_parallelism_combined.py b/tests/models/test_parallelism_combined.py index 76b814c73cca..4a404a06125e 100644 --- a/tests/models/test_parallelism_combined.py +++ b/tests/models/test_parallelism_combined.py @@ -50,7 +50,7 @@ def _mesh_factorization_worker(rank, world_size, master_port, ulysses_degree, tp dist.init_process_group(backend="gloo", rank=rank, world_size=world_size) mesh = torch.distributed.device_mesh.init_device_mesh( - "cpu", mesh_shape=(1, ulysses_degree, tp_degree), mesh_dim_names=("ring", "ulysses", "tp") + "cpu", mesh_shape=(tp_degree, 1, ulysses_degree), mesh_dim_names=("tp", "ring", "ulysses") ) config = ParallelConfig( context_parallel_config=ContextParallelConfig(ulysses_degree=ulysses_degree), @@ -134,10 +134,10 @@ def test_single_parallelism_is_not_combined(self): 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 `mesh_shape=(ring=1, ulysses=2, tp=2)` over four ranks, "tp" is the fastest-varying dimension, so ranks - {0,1} and {2,3} are the TP groups and {0,2} / {1,3} are the CP 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. + 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) @@ -145,12 +145,12 @@ def test_mesh_is_factored_between_the_two(self): 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, 1], [0, 1], [2, 3], [2, 3]] - assert [results[r]["cp_ranks"] for r in range(4)] == [[0, 2], [1, 3], [0, 2], [1, 3]] - assert [results[r]["ulysses_local_rank"] for r in range(4)] == [0, 0, 1, 1] + 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, 1, 0, 1] + 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): From 9c0f7aaf39fcf7e47e0c66f32fd1c1a8032fd656 Mon Sep 17 00:00:00 2001 From: JingyaHuang Date: Thu, 8 Oct 2026 14:23:18 +0000 Subject: [PATCH 08/19] fix: style --- src/diffusers/hooks/tensor_parallel.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/diffusers/hooks/tensor_parallel.py b/src/diffusers/hooks/tensor_parallel.py index ce8abd844ab5..ccfd609caee8 100644 --- a/src/diffusers/hooks/tensor_parallel.py +++ b/src/diffusers/hooks/tensor_parallel.py @@ -53,14 +53,14 @@ def __init__(self, blocks: "list[int] | None" = None): 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, + 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. + 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. """ From bbc55a76ae338ce235209d716f0dba4335b012de Mon Sep 17 00:00:00 2001 From: JingyaHuang Date: Fri, 9 Oct 2026 13:14:09 +0000 Subject: [PATCH 09/19] fix: apply context parallelism on the sharded-load path --- src/diffusers/hooks/context_parallel.py | 11 ++ src/diffusers/hooks/tensor_parallel.py | 49 +------- src/diffusers/hooks/tensor_parallel_neuron.py | 31 ++++- src/diffusers/models/attention_dispatch.py | 36 ++++++ src/diffusers/models/modeling_utils.py | 107 ++++++++++-------- tests/models/testing_utils/parallelism.py | 55 ++++++--- .../transformers/_neuron_hybrid_worker.py | 25 +++- .../test_models_transformer_flux.py | 5 +- 8 files changed, 202 insertions(+), 117 deletions(-) diff --git a/src/diffusers/hooks/context_parallel.py b/src/diffusers/hooks/context_parallel.py index 1310b20c5c11..ad2becf5cc5c 100644 --- a/src/diffusers/hooks/context_parallel.py +++ b/src/diffusers/hooks/context_parallel.py @@ -345,6 +345,17 @@ def _all_gather_anything(tensor: torch.Tensor, dim: int, group: dist.device_mesh rank_dim = shape[dim] gather_dims = gather_size_by_comm(rank_dim, group) + max_dim = max(gather_dims) + if any(d != max_dim for d in gather_dims) and tensor.device.type == "neuron": + # Neuron's all-gather requires every rank to send the same size, so pad each chunk to the longest one, + # gather, and trim each chunk back to its real length. + pad_shape = list(shape) + pad_shape[dim] = max_dim - rank_dim + padded = torch.cat([tensor, tensor.new_zeros(pad_shape)], dim=dim) + gathered_tensors = [torch.empty_like(padded) for _ in range(world_size)] + dist.all_gather(gathered_tensors, padded, group=group) + return torch.cat([t.narrow(dim, 0, d) for t, d in zip(gathered_tensors, gather_dims)], dim=dim) + gather_shapes = _fill_gather_shapes(tuple(shape), tuple(gather_dims), dim, world_size) gathered_tensors = [torch.empty(shape, device=tensor.device, dtype=tensor.dtype) for shape in gather_shapes] diff --git a/src/diffusers/hooks/tensor_parallel.py b/src/diffusers/hooks/tensor_parallel.py index ccfd609caee8..b83d87d9f0b6 100644 --- a/src/diffusers/hooks/tensor_parallel.py +++ b/src/diffusers/hooks/tensor_parallel.py @@ -358,47 +358,6 @@ def _partition_linear_fn(self, name, module, device_mesh): 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 @@ -565,7 +524,7 @@ def apply_tensor_parallel( f"or from the active accelerator when the mesh is built from `tp_degree`." ) - backend = tp_mesh.device_type if tp_mesh.device_type in ("neuron", "tpu") else "default" + 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}.") @@ -585,11 +544,5 @@ def apply_tensor_parallel( _apply_tp_neuron(model, tp_mesh, groups, specs) return - if backend == "tpu": - # `parallelize_module` would materialize every full weight on each chip before scattering it, which runs out - # of HBM on models larger than one chip. Pre-sharding on CPU keeps only this rank's slice on the device. - _pre_shard_and_parallelize(model, tp_mesh, groups, specs, config._device) - return - 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 index a87f28b1edb1..ffcba0973d81 100644 --- a/src/diffusers/hooks/tensor_parallel_neuron.py +++ b/src/diffusers/hooks/tensor_parallel_neuron.py @@ -28,7 +28,7 @@ import torch import torch.nn as nn -from .tensor_parallel import TPShardSpec, _pre_shard_and_parallelize +from .tensor_parallel import TPShardSpec, _hooks_only_styles, _local_shard def _apply_tp_neuron( @@ -45,4 +45,31 @@ def _apply_tp_neuron( Model weights must be on CPU when this is called. """ - _pre_shard_and_parallelize(model, tp_mesh, groups, specs, torch.neuron.current_device()) + 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/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py index f45c9e0d1107..7a9c6b7a4d9c 100644 --- a/src/diffusers/models/attention_dispatch.py +++ b/src/diffusers/models/attention_dispatch.py @@ -2337,6 +2337,22 @@ def all_to_all_single_any_qkv_async( # since we don't know the actual shape before this timing, thus, # we have to use all gather to collect the S_LOCAL first. output_split_sizes = gather_size_by_comm(S_LOCAL, group) + S_MAX = max(output_split_sizes) + if x.device.type == "neuron" and S_MAX != min(output_split_sizes): + # Neuron's all-to-all only handles equal splits (uneven ones silently return wrong data), so pad every + # rank's chunk to the longest, exchange equal splits, and drop the padding on arrival. + x = torch.cat([x, x.new_zeros(world_size, S_MAX - S_LOCAL, B, H_LOCAL, D)], dim=1).flatten(0, 1) + x = funcol.all_to_all_single(x, None, None, group) + + def wait_padded() -> torch.Tensor: + nonlocal x, H_PAD + x = _wait_tensor(x).unflatten(0, (world_size, S_MAX)) # (world_size, S_MAX, B, H_LOCAL, D) + x = torch.cat([x[i, :size] for i, size in enumerate(output_split_sizes)]) # (S_GLOBAL, B, H_LOCAL, D) + x = x.permute(1, 0, 2, 3).contiguous() + return _maybe_unpad_qkv_head(x, H_PAD, group) + + return wait_padded + x = x.flatten(0, 1) # (world_size * S_LOCAL, B, H_LOCAL, D) x = funcol.all_to_all_single(x, output_split_sizes, input_split_sizes, group) @@ -2377,6 +2393,26 @@ def all_to_all_single_any_o_async(x: torch.Tensor, group: dist.ProcessGroup, **k S_LOCAL = kwargs.get("Q_S_LOCAL") input_split_sizes = gather_size_by_comm(S_LOCAL, group) x = x.permute(1, 0, 2, 3).contiguous() # (S_GLOBAL, B, H_LOCAL, D) + S_MAX = max(input_split_sizes) + if x.device.type == "neuron" and S_MAX != min(input_split_sizes): + # Same equal-split constraint as in `all_to_all_single_any_qkv_async`: pad each outgoing chunk to the longest. + x = torch.cat( + [ + torch.cat([chunk, chunk.new_zeros(S_MAX - chunk.shape[0], B, H_LOCAL, D)]) + for chunk in x.split(input_split_sizes) + ] + ) + x = funcol.all_to_all_single(x, None, None, group) + + def wait_padded() -> torch.Tensor: + nonlocal x, H_PAD + x = _wait_tensor(x).reshape(world_size, S_MAX, B, H_LOCAL, D)[:, :S_LOCAL] + x = x.permute(2, 1, 0, 3, 4).contiguous() + x = x.reshape(B, S_LOCAL, world_size * H_LOCAL, D) + return _maybe_unpad_o_head(x, H_PAD, group) + + return wait_padded + output_split_sizes = [S_LOCAL] * world_size x = funcol.all_to_all_single(x, output_split_sizes, input_split_sizes, group) diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index 332ac65864ac..5051853be46a 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -1559,6 +1559,9 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None 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) @@ -1802,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 " @@ -1833,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( @@ -2001,6 +2013,7 @@ def _load_pretrained_model( low_cpu_mem_usage=low_cpu_mem_usage, disable_mmap=disable_mmap, ) + if is_parallel_loading_enabled: offload_index, state_dict_index, _mismatched_keys, _error_msgs = load_fn(resolved_model_file) error_msgs += _error_msgs diff --git a/tests/models/testing_utils/parallelism.py b/tests/models/testing_utils/parallelism.py index ef8927d328f5..d6c57b84c587 100644 --- a/tests/models/testing_utils/parallelism.py +++ b/tests/models/testing_utils/parallelism.py @@ -970,7 +970,17 @@ def test_context_parallel_attn_backend_inference(self, cp_type, attention_backen def _hybrid_parallel_worker( - rank, world_size, master_port, model_class, init_dict, cp_dict, tp_degree, inputs_dict, return_dict, state_dict + 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. @@ -978,6 +988,9 @@ def _hybrid_parallel_worker( 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" @@ -994,20 +1007,29 @@ def _hybrid_parallel_worker( 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() + 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()} - model.enable_parallelism( - config=ParallelConfig( - tensor_parallel_config=TensorParallelConfig(tp_degree=tp_degree), - context_parallel_config=ContextParallelConfig(**cp_dict), - ) - ) - with torch.no_grad(): output = model(**inputs_on_device, return_dict=False)[0] @@ -1035,7 +1057,8 @@ class HybridParallelTesterMixin: Needs `tp_degree * ulysses_degree` accelerators (4 at the degrees used here), so it skips on a 2-device runner. """ - def test_hybrid_parallel_inference(self, batch_size: int = 1): + @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.") @@ -1068,6 +1091,11 @@ def test_hybrid_parallel_inference(self, batch_size: int = 1): 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() @@ -1084,6 +1112,7 @@ def test_hybrid_parallel_inference(self, batch_size: int = 1): inputs_dict, return_dict, state_dict, + checkpoint_dir, ), nprocs=world_size, join=True, diff --git a/tests/models/transformers/_neuron_hybrid_worker.py b/tests/models/transformers/_neuron_hybrid_worker.py index 1c054af44104..0c0e44896beb 100644 --- a/tests/models/transformers/_neuron_hybrid_worker.py +++ b/tests/models/transformers/_neuron_hybrid_worker.py @@ -37,6 +37,7 @@ import importlib import os import sys +import tempfile import traceback @@ -82,12 +83,26 @@ def main(): with torch.no_grad(): ref_output = model(**inputs, return_dict=False)[0].float().cpu() - model.enable_parallelism( - config=ParallelConfig( - tensor_parallel_config=TensorParallelConfig(tp_degree=tp_degree), - context_parallel_config=ContextParallelConfig(ulysses_degree=ulysses_degree), - ) + 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() diff --git a/tests/models/transformers/test_models_transformer_flux.py b/tests/models/transformers/test_models_transformer_flux.py index 8b36e579d833..b8ad9cdc6a83 100644 --- a/tests/models/transformers/test_models_transformer_flux.py +++ b/tests/models/transformers/test_models_transformer_flux.py @@ -353,11 +353,12 @@ class TestFluxTransformerHybridParallelNeuron: `HybridParallelTesterMixin`. Runs at `tp_degree=2 x ulysses_degree=4`, i.e. 8 ranks. """ - def test_hybrid_parallel_neuron_inference(self): + @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) + 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}" From 6ce34c772be4d22db85c70db8f5651acbca5dad5 Mon Sep 17 00:00:00 2001 From: JingyaHuang Date: Fri, 9 Oct 2026 13:24:15 +0000 Subject: [PATCH 10/19] revert: attention head padding --- src/diffusers/models/attention_dispatch.py | 36 ---------------------- 1 file changed, 36 deletions(-) diff --git a/src/diffusers/models/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py index 7a9c6b7a4d9c..f45c9e0d1107 100644 --- a/src/diffusers/models/attention_dispatch.py +++ b/src/diffusers/models/attention_dispatch.py @@ -2337,22 +2337,6 @@ def all_to_all_single_any_qkv_async( # since we don't know the actual shape before this timing, thus, # we have to use all gather to collect the S_LOCAL first. output_split_sizes = gather_size_by_comm(S_LOCAL, group) - S_MAX = max(output_split_sizes) - if x.device.type == "neuron" and S_MAX != min(output_split_sizes): - # Neuron's all-to-all only handles equal splits (uneven ones silently return wrong data), so pad every - # rank's chunk to the longest, exchange equal splits, and drop the padding on arrival. - x = torch.cat([x, x.new_zeros(world_size, S_MAX - S_LOCAL, B, H_LOCAL, D)], dim=1).flatten(0, 1) - x = funcol.all_to_all_single(x, None, None, group) - - def wait_padded() -> torch.Tensor: - nonlocal x, H_PAD - x = _wait_tensor(x).unflatten(0, (world_size, S_MAX)) # (world_size, S_MAX, B, H_LOCAL, D) - x = torch.cat([x[i, :size] for i, size in enumerate(output_split_sizes)]) # (S_GLOBAL, B, H_LOCAL, D) - x = x.permute(1, 0, 2, 3).contiguous() - return _maybe_unpad_qkv_head(x, H_PAD, group) - - return wait_padded - x = x.flatten(0, 1) # (world_size * S_LOCAL, B, H_LOCAL, D) x = funcol.all_to_all_single(x, output_split_sizes, input_split_sizes, group) @@ -2393,26 +2377,6 @@ def all_to_all_single_any_o_async(x: torch.Tensor, group: dist.ProcessGroup, **k S_LOCAL = kwargs.get("Q_S_LOCAL") input_split_sizes = gather_size_by_comm(S_LOCAL, group) x = x.permute(1, 0, 2, 3).contiguous() # (S_GLOBAL, B, H_LOCAL, D) - S_MAX = max(input_split_sizes) - if x.device.type == "neuron" and S_MAX != min(input_split_sizes): - # Same equal-split constraint as in `all_to_all_single_any_qkv_async`: pad each outgoing chunk to the longest. - x = torch.cat( - [ - torch.cat([chunk, chunk.new_zeros(S_MAX - chunk.shape[0], B, H_LOCAL, D)]) - for chunk in x.split(input_split_sizes) - ] - ) - x = funcol.all_to_all_single(x, None, None, group) - - def wait_padded() -> torch.Tensor: - nonlocal x, H_PAD - x = _wait_tensor(x).reshape(world_size, S_MAX, B, H_LOCAL, D)[:, :S_LOCAL] - x = x.permute(2, 1, 0, 3, 4).contiguous() - x = x.reshape(B, S_LOCAL, world_size * H_LOCAL, D) - return _maybe_unpad_o_head(x, H_PAD, group) - - return wait_padded - output_split_sizes = [S_LOCAL] * world_size x = funcol.all_to_all_single(x, output_split_sizes, input_split_sizes, group) From 986ea9a404b2fd6491377cb88e33bd5ae89f6e2b Mon Sep 17 00:00:00 2001 From: JingyaHuang Date: Sat, 10 Oct 2026 10:33:07 +0000 Subject: [PATCH 11/19] revert: drop the Neuron all-gather padding and restore the TPU pre-shard path --- src/diffusers/hooks/context_parallel.py | 11 ----- src/diffusers/hooks/tensor_parallel.py | 49 ++++++++++++++++++- src/diffusers/hooks/tensor_parallel_neuron.py | 31 +----------- 3 files changed, 50 insertions(+), 41 deletions(-) diff --git a/src/diffusers/hooks/context_parallel.py b/src/diffusers/hooks/context_parallel.py index ad2becf5cc5c..1310b20c5c11 100644 --- a/src/diffusers/hooks/context_parallel.py +++ b/src/diffusers/hooks/context_parallel.py @@ -345,17 +345,6 @@ def _all_gather_anything(tensor: torch.Tensor, dim: int, group: dist.device_mesh rank_dim = shape[dim] gather_dims = gather_size_by_comm(rank_dim, group) - max_dim = max(gather_dims) - if any(d != max_dim for d in gather_dims) and tensor.device.type == "neuron": - # Neuron's all-gather requires every rank to send the same size, so pad each chunk to the longest one, - # gather, and trim each chunk back to its real length. - pad_shape = list(shape) - pad_shape[dim] = max_dim - rank_dim - padded = torch.cat([tensor, tensor.new_zeros(pad_shape)], dim=dim) - gathered_tensors = [torch.empty_like(padded) for _ in range(world_size)] - dist.all_gather(gathered_tensors, padded, group=group) - return torch.cat([t.narrow(dim, 0, d) for t, d in zip(gathered_tensors, gather_dims)], dim=dim) - gather_shapes = _fill_gather_shapes(tuple(shape), tuple(gather_dims), dim, world_size) gathered_tensors = [torch.empty(shape, device=tensor.device, dtype=tensor.dtype) for shape in gather_shapes] diff --git a/src/diffusers/hooks/tensor_parallel.py b/src/diffusers/hooks/tensor_parallel.py index b83d87d9f0b6..ccfd609caee8 100644 --- a/src/diffusers/hooks/tensor_parallel.py +++ b/src/diffusers/hooks/tensor_parallel.py @@ -358,6 +358,47 @@ def _partition_linear_fn(self, name, module, device_mesh): 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 @@ -524,7 +565,7 @@ 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" + backend = tp_mesh.device_type if tp_mesh.device_type in ("neuron", "tpu") 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}.") @@ -544,5 +585,11 @@ def apply_tensor_parallel( _apply_tp_neuron(model, tp_mesh, groups, specs) return + if backend == "tpu": + # `parallelize_module` would materialize every full weight on each chip before scattering it, which runs out + # of HBM on models larger than one chip. Pre-sharding on CPU keeps only this rank's slice on the device. + _pre_shard_and_parallelize(model, tp_mesh, groups, specs, config._device) + return + 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 index ffcba0973d81..a87f28b1edb1 100644 --- a/src/diffusers/hooks/tensor_parallel_neuron.py +++ b/src/diffusers/hooks/tensor_parallel_neuron.py @@ -28,7 +28,7 @@ import torch import torch.nn as nn -from .tensor_parallel import TPShardSpec, _hooks_only_styles, _local_shard +from .tensor_parallel import TPShardSpec, _pre_shard_and_parallelize def _apply_tp_neuron( @@ -45,31 +45,4 @@ def _apply_tp_neuron( 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)) + _pre_shard_and_parallelize(model, tp_mesh, groups, specs, torch.neuron.current_device()) From 6845f33362dae794ec2d1e1e3479cae7413f1a67 Mon Sep 17 00:00:00 2001 From: JingyaHuang Date: Sat, 10 Oct 2026 13:09:51 +0000 Subject: [PATCH 12/19] refactor: share the CPU pre-shard path between Neuron and TPU tensor parallelism --- src/diffusers/hooks/tensor_parallel.py | 19 +++----- src/diffusers/hooks/tensor_parallel_neuron.py | 48 ------------------- 2 files changed, 7 insertions(+), 60 deletions(-) delete mode 100644 src/diffusers/hooks/tensor_parallel_neuron.py diff --git a/src/diffusers/hooks/tensor_parallel.py b/src/diffusers/hooks/tensor_parallel.py index ccfd609caee8..e3f4e05dc7bf 100644 --- a/src/diffusers/hooks/tensor_parallel.py +++ b/src/diffusers/hooks/tensor_parallel.py @@ -565,9 +565,8 @@ def apply_tensor_parallel( f"or from the active accelerator when the mesh is built from `tp_degree`." ) - backend = tp_mesh.device_type if tp_mesh.device_type in ("neuron", "tpu") 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 @@ -579,16 +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) - return - - if backend == "tpu": - # `parallelize_module` would materialize every full weight on each chip before scattering it, which runs out - # of HBM on models larger than one chip. Pre-sharding on CPU keeps only this rank's slice on the device. - _pre_shard_and_parallelize(model, tp_mesh, groups, specs, config._device) + 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 a87f28b1edb1..000000000000 --- a/src/diffusers/hooks/tensor_parallel_neuron.py +++ /dev/null @@ -1,48 +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, _pre_shard_and_parallelize - - -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. - """ - _pre_shard_and_parallelize(model, tp_mesh, groups, specs, torch.neuron.current_device()) From 2590829f5eddfa44fcb89874239016b89e65c2e6 Mon Sep 17 00:00:00 2001 From: JingyaHuang Date: Tue, 13 Oct 2026 13:31:36 +0000 Subject: [PATCH 13/19] fix: new device count API for TPU since deprecation --- src/diffusers/models/modeling_utils.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index 5051853be46a..5cb4ee6ca903 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -1732,8 +1732,7 @@ def _resolve_parallel_config( rank = torch.distributed.get_rank() world_size = torch.distributed.get_world_size() device_type = torch._C._get_accelerator().type - device_module = torch.get_device_module(device_type) - device = torch.device(device_type, rank % device_module.device_count()) + device = torch.device(device_type, rank % torch.accelerator.device_count()) mesh = None cp_config = config.context_parallel_config From ce4198acba8011ed8e9f0e191e87979aa5e14ecc Mon Sep 17 00:00:00 2001 From: JingyaHuang Date: Tue, 13 Oct 2026 14:47:33 +0000 Subject: [PATCH 14/19] fix: pad uneven all-gathers on Neuron and TPU in _all_gather_anything --- src/diffusers/hooks/context_parallel.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/src/diffusers/hooks/context_parallel.py b/src/diffusers/hooks/context_parallel.py index 1310b20c5c11..4b0a90dac4c2 100644 --- a/src/diffusers/hooks/context_parallel.py +++ b/src/diffusers/hooks/context_parallel.py @@ -345,6 +345,17 @@ def _all_gather_anything(tensor: torch.Tensor, dim: int, group: dist.device_mesh rank_dim = shape[dim] gather_dims = gather_size_by_comm(rank_dim, group) + max_dim = max(gather_dims) + if any(d != max_dim for d in gather_dims) and tensor.device.type in ("neuron", "tpu"): + # Neuron and TPU all-gathers require every rank to send the same size, so pad each chunk to the longest one, + # gather, and trim each chunk back to its real length. + pad_shape = list(shape) + pad_shape[dim] = max_dim - rank_dim + padded = torch.cat([tensor, tensor.new_zeros(pad_shape)], dim=dim) + gathered_tensors = [torch.empty_like(padded) for _ in range(world_size)] + dist.all_gather(gathered_tensors, padded, group=group) + return torch.cat([t.narrow(dim, 0, d) for t, d in zip(gathered_tensors, gather_dims)], dim=dim) + gather_shapes = _fill_gather_shapes(tuple(shape), tuple(gather_dims), dim, world_size) gathered_tensors = [torch.empty(shape, device=tensor.device, dtype=tensor.dtype) for shape in gather_shapes] From cd611662e4148ab485c37166d80bac33c5816c23 Mon Sep 17 00:00:00 2001 From: JingyaHuang Date: Wed, 14 Oct 2026 15:37:36 +0000 Subject: [PATCH 15/19] fix: fall back to one device when no accelerator is present Co-Authored-By: Claude Opus 5.5 (1M context) --- src/diffusers/models/modeling_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index 5cb4ee6ca903..e70eb541eddc 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -1732,7 +1732,7 @@ def _resolve_parallel_config( rank = torch.distributed.get_rank() world_size = torch.distributed.get_world_size() device_type = torch._C._get_accelerator().type - device = torch.device(device_type, rank % torch.accelerator.device_count()) + device = torch.device(device_type, rank % max(torch.accelerator.device_count(), 1)) mesh = None cp_config = config.context_parallel_config From cc851100d65c3a0f2ad8734d6794e3ecb55a0b73 Mon Sep 17 00:00:00 2001 From: JingyaHuang Date: Wed, 14 Oct 2026 15:37:54 +0000 Subject: [PATCH 16/19] Move TP+CP composition to #14725 Keep this PR to tensor parallelism for MiniMax-H3. Combining TP and CP in one ParallelConfig, with its tests, now lives in #14725. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/diffusers/hooks/context_parallel.py | 11 - src/diffusers/models/_modeling_parallel.py | 43 +-- src/diffusers/models/modeling_utils.py | 153 +++------ tests/models/test_parallelism_combined.py | 164 ---------- tests/models/testing_utils/__init__.py | 4 - tests/models/testing_utils/parallelism.py | 304 +----------------- .../transformers/_neuron_hybrid_worker.py | 140 -------- .../test_models_transformer_flux.py | 49 +-- .../test_models_transformer_qwenimage.py | 20 -- 9 files changed, 58 insertions(+), 830 deletions(-) delete mode 100644 tests/models/test_parallelism_combined.py delete mode 100644 tests/models/transformers/_neuron_hybrid_worker.py diff --git a/src/diffusers/hooks/context_parallel.py b/src/diffusers/hooks/context_parallel.py index 4b0a90dac4c2..1310b20c5c11 100644 --- a/src/diffusers/hooks/context_parallel.py +++ b/src/diffusers/hooks/context_parallel.py @@ -345,17 +345,6 @@ def _all_gather_anything(tensor: torch.Tensor, dim: int, group: dist.device_mesh rank_dim = shape[dim] gather_dims = gather_size_by_comm(rank_dim, group) - max_dim = max(gather_dims) - if any(d != max_dim for d in gather_dims) and tensor.device.type in ("neuron", "tpu"): - # Neuron and TPU all-gathers require every rank to send the same size, so pad each chunk to the longest one, - # gather, and trim each chunk back to its real length. - pad_shape = list(shape) - pad_shape[dim] = max_dim - rank_dim - padded = torch.cat([tensor, tensor.new_zeros(pad_shape)], dim=dim) - gathered_tensors = [torch.empty_like(padded) for _ in range(world_size)] - dist.all_gather(gathered_tensors, padded, group=group) - return torch.cat([t.narrow(dim, 0, d) for t, d in zip(gathered_tensors, gather_dims)], dim=dim) - gather_shapes = _fill_gather_shapes(tuple(shape), tuple(gather_dims), dim, world_size) gathered_tensors = [torch.empty(shape, device=tensor.device, dtype=tensor.dtype) for shape in gather_shapes] diff --git a/src/diffusers/models/_modeling_parallel.py b/src/diffusers/models/_modeling_parallel.py index cf576ab6bb77..58cbc822f8bf 100644 --- a/src/diffusers/models/_modeling_parallel.py +++ b/src/diffusers/models/_modeling_parallel.py @@ -200,12 +200,6 @@ 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. @@ -222,17 +216,12 @@ class ParallelConfig: _mesh: torch.distributed.device_mesh.DeviceMesh = None def __post_init__(self): - if self.context_parallel_config is None and self.tensor_parallel_config is None: + if self.context_parallel_config is not None and self.tensor_parallel_config is not None: raise ValueError( - "A `ParallelConfig` must specify at least one of `context_parallel_config` or " - "`tensor_parallel_config`." + "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`." ) - @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, @@ -245,32 +234,10 @@ 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, cp_mesh) + 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, tp_mesh) + 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 e70eb541eddc..731fb402a32c 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -1559,9 +1559,6 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None 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) @@ -1735,28 +1732,15 @@ def _resolve_parallel_config( device = torch.device(device_type, rank % max(torch.accelerator.device_count(), 1)) mesh = None - 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: + if config.context_parallel_config is not None: + cp_config = config.context_parallel_config 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 tp_config is not None: + 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,), @@ -1765,32 +1749,6 @@ 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 @@ -1804,6 +1762,11 @@ 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 " @@ -1830,73 +1793,59 @@ 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) - for module in self.modules(): - if not isinstance(module, attention_classes): - continue + 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, "_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. - 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 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 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 + + apply_tensor_parallel(self, config.tensor_parallel_config, self._tp_plan) @classmethod def _load_pretrained_model( diff --git a/tests/models/test_parallelism_combined.py b/tests/models/test_parallelism_combined.py deleted file mode 100644 index 4a404a06125e..000000000000 --- a/tests/models/test_parallelism_combined.py +++ /dev/null @@ -1,164 +0,0 @@ -# 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 6db810342821..ee0ab21e8719 100644 --- a/tests/models/testing_utils/__init__.py +++ b/tests/models/testing_utils/__init__.py @@ -21,10 +21,8 @@ from .lora import LoraHotSwappingForModelTesterMixin, LoraTesterMixin from .memory import CPUOffloadTesterMixin, GroupOffloadTesterMixin, LayerwiseCastingTesterMixin, MemoryTesterMixin from .parallelism import ( - ContextAndTensorParallelTesterMixin, ContextParallelAttentionBackendsTesterMixin, ContextParallelTesterMixin, - HybridParallelTesterMixin, TensorParallelTesterMixin, TensorParallelTPUTesterMixin, ) @@ -68,10 +66,8 @@ "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 d6c57b84c587..eca6bff5d271 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, ParallelConfig, TensorParallelConfig +from diffusers.models._modeling_parallel import ContextParallelConfig, TensorParallelConfig from diffusers.models.attention_dispatch import AttentionBackendName, _AttentionBackendRegistry from diffusers.utils import SAFE_WEIGHTS_INDEX_NAME @@ -558,150 +558,6 @@ 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: @@ -967,161 +823,3 @@ 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 deleted file mode 100644 index 0c0e44896beb..000000000000 --- a/tests/models/transformers/_neuron_hybrid_worker.py +++ /dev/null @@ -1,140 +0,0 @@ -# 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 b8ad9cdc6a83..04542253c7f6 100644 --- a/tests/models/transformers/test_models_transformer_flux.py +++ b/tests/models/transformers/test_models_transformer_flux.py @@ -27,13 +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, - is_context_parallel, - is_tensor_parallel, - require_torch_neuron, - torch_device, -) +from ...testing_utils import enable_full_determinism, is_tensor_parallel, require_torch_neuron, torch_device from ..testing_utils import ( AttentionBackendTesterMixin, AttentionTesterMixin, @@ -46,7 +40,6 @@ FirstBlockCacheTesterMixin, GGUFCompileTesterMixin, GGUFTesterMixin, - HybridParallelTesterMixin, IPAdapterTesterMixin, LoraHotSwappingForModelTesterMixin, LoraTesterMixin, @@ -292,23 +285,6 @@ 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`). @@ -342,29 +318,6 @@ 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_qwenimage.py b/tests/models/transformers/test_models_transformer_qwenimage.py index 10d36c8155ba..6210aecff87f 100644 --- a/tests/models/transformers/test_models_transformer_qwenimage.py +++ b/tests/models/transformers/test_models_transformer_qwenimage.py @@ -30,7 +30,6 @@ AttentionTesterMixin, BaseModelTesterConfig, BitsAndBytesTesterMixin, - ContextAndTensorParallelTesterMixin, ContextParallelAttentionBackendsTesterMixin, ContextParallelTesterMixin, LoraHotSwappingForModelTesterMixin, @@ -321,25 +320,6 @@ 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``). From 8526d9427f9e93d101488866494d29e4936f6c02 Mon Sep 17 00:00:00 2001 From: JingyaHuang Date: Thu, 15 Oct 2026 13:52:21 +0000 Subject: [PATCH 17/19] review: use `config._device` for the Neuron TP device `config._device` now resolves to the local Neuron core on every rank, so the `torch.neuron.current_device()` special case is no longer needed when pre-sharding, loading sharded checkpoints and moving buffers. --- src/diffusers/hooks/tensor_parallel.py | 3 +-- src/diffusers/models/model_loading_utils.py | 7 +------ src/diffusers/models/modeling_utils.py | 4 +--- 3 files changed, 3 insertions(+), 11 deletions(-) diff --git a/src/diffusers/hooks/tensor_parallel.py b/src/diffusers/hooks/tensor_parallel.py index 5ff6ac9333ce..1663528b783c 100644 --- a/src/diffusers/hooks/tensor_parallel.py +++ b/src/diffusers/hooks/tensor_parallel.py @@ -582,8 +582,7 @@ def apply_tensor_parallel( # 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) + _pre_shard_and_parallelize(model, tp_mesh, groups, specs, config._device) return for submodule, relative_plan in groups: diff --git a/src/diffusers/models/model_loading_utils.py b/src/diffusers/models/model_loading_utils.py index b6d9d7997d81..655f50b5378a 100644 --- a/src/diffusers/models/model_loading_utils.py +++ b/src/diffusers/models/model_loading_utils.py @@ -407,12 +407,7 @@ def _load_shard_file_tp( from ..hooks.tensor_parallel import _local_shard tp_mesh = tp_config._mesh - # `TensorParallelConfig._device` is derived from the default accelerator, which is not meaningful on - # Neuron; resolve it the way the Neuron pre-shard backend does. - if tp_mesh.device_type == "neuron": - device = torch.neuron.current_device() - else: - device = tp_config._device + device = tp_config._device mismatched_keys = [] diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index 731fb402a32c..35774a129b8d 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -1550,9 +1550,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None # 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 - ) + tp_device = 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(".") From ed452ef87d1f69cd1349f078748c90957fe218c6 Mon Sep 17 00:00:00 2001 From: JingyaHuang Date: Thu, 15 Oct 2026 12:29:06 +0000 Subject: [PATCH 18/19] review: replace ReplicatedInputRowwiseParallel with a "rowwise_split_input" style The marker class carried no state, so make it a plain string style next to "colwise" and "rowwise". The name matches the transformers style of the same name, which is the same RowwiseParallel(input_layouts=Replicate()). Add it to the TP styles table in the distributed inference docs, with the AdaLN modulation projection as the example. Co-Authored-By: Claude Opus 5.5 --- .../en/training/distributed_inference.md | 1 + src/diffusers/hooks/tensor_parallel.py | 38 ++++++------------- .../transformers/transformer_minimax_h3.py | 6 +-- 3 files changed, 16 insertions(+), 29 deletions(-) diff --git a/docs/source/en/training/distributed_inference.md b/docs/source/en/training/distributed_inference.md index 818488b4cfce..25de837a8c25 100644 --- a/docs/source/en/training/distributed_inference.md +++ b/docs/source/en/training/distributed_inference.md @@ -511,6 +511,7 @@ Each key may contain **at most one `*`**, and the prefix before it must resolve |---|---|---|---| | `"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 | +| `"rowwise_split_input"` | input features (`weight` dim 1) | slices its replicated input, then computes a partial sum | a standalone projection that reads a replicated activation and must output the full width, e.g. an AdaLN modulation projection off the timestep embedding | 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. diff --git a/src/diffusers/hooks/tensor_parallel.py b/src/diffusers/hooks/tensor_parallel.py index 1663528b783c..abe22eed7980 100644 --- a/src/diffusers/hooks/tensor_parallel.py +++ b/src/diffusers/hooks/tensor_parallel.py @@ -50,20 +50,6 @@ 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" @@ -163,7 +149,7 @@ 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" or isinstance(style, ReplicatedInputRowwiseParallel): + elif style in ("rowwise", "rowwise_split_input"): # 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) @@ -178,8 +164,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, PackedRowwiseParallel, or " - f"ReplicatedInputRowwiseParallel." + f"Expected 'colwise', 'rowwise', 'rowwise_split_input', PackedColwiseParallel, or " + f"PackedRowwiseParallel." ) specs[f"{path}.weight"] = weight_spec @@ -258,9 +244,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` / - `ReplicatedInputRowwiseParallel` 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"` / `"rowwise_split_input"`) or `PackedColwiseParallel` / + `PackedRowwiseParallel` 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 @@ -301,7 +287,7 @@ def _partition_linear_fn(self, name, module, device_mesh): resolved[path] = ColwiseParallel() elif style == "rowwise": resolved[path] = RowwiseParallel() - elif isinstance(style, ReplicatedInputRowwiseParallel): + elif style == "rowwise_split_input": # `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. @@ -313,8 +299,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, PackedRowwiseParallel, or " - f"ReplicatedInputRowwiseParallel." + f"Expected 'colwise', 'rowwise', 'rowwise_split_input', PackedColwiseParallel, or " + f"PackedRowwiseParallel." ) return resolved @@ -347,13 +333,13 @@ 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): + elif style == "rowwise_split_input": resolved[path] = _NoPartitionRowwise(input_layouts=Replicate()) else: raise ValueError( f"Unsupported tensor-parallel style '{style}' for '{path}'. " - f"Expected 'colwise', 'rowwise', PackedColwiseParallel, PackedRowwiseParallel, or " - f"ReplicatedInputRowwiseParallel." + f"Expected 'colwise', 'rowwise', 'rowwise_split_input', PackedColwiseParallel, or " + f"PackedRowwiseParallel." ) return resolved diff --git a/src/diffusers/models/transformers/transformer_minimax_h3.py b/src/diffusers/models/transformers/transformer_minimax_h3.py index 0f9e38bdb965..5e167ccfc36e 100644 --- a/src/diffusers/models/transformers/transformer_minimax_h3.py +++ b/src/diffusers/models/transformers/transformer_minimax_h3.py @@ -19,7 +19,7 @@ import torch.nn as nn from ...configuration_utils import ConfigMixin, register_to_config -from ...hooks.tensor_parallel import PackedColwiseParallel, ReplicatedInputRowwiseParallel +from ...hooks.tensor_parallel import PackedColwiseParallel from ...loaders import FromOriginalModelMixin, PeftAdapterMixin from ...utils import BaseOutput, apply_lora_scale, logging from .._modeling_parallel import ContextParallelInput, ContextParallelOutput @@ -462,7 +462,7 @@ class MiniMaxH3Transformer3DModel( # `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 + # sharded `"rowwise_split_input"` 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 * @@ -484,7 +484,7 @@ class MiniMaxH3Transformer3DModel( "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(), + "transformer_blocks.*.adaln_proj.linear": "rowwise_split_input", # 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", From bc35ec2d513653a38bf9a6bf00f25eb212cd3cf3 Mon Sep 17 00:00:00 2001 From: JingyaHuang Date: Thu, 15 Oct 2026 14:24:16 +0000 Subject: [PATCH 19/19] review: explain when to use "rowwise_split_input" Replace the comment on how torch narrows the input with when the style is needed, plus two examples: H3's AdaLN modulation projection and an output projection after a gathered QK-norm. Co-Authored-By: Claude Opus 5.5 --- src/diffusers/hooks/tensor_parallel.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/diffusers/hooks/tensor_parallel.py b/src/diffusers/hooks/tensor_parallel.py index abe22eed7980..951450a8dffb 100644 --- a/src/diffusers/hooks/tensor_parallel.py +++ b/src/diffusers/hooks/tensor_parallel.py @@ -288,9 +288,10 @@ def _partition_linear_fn(self, name, module, device_mesh): elif style == "rowwise": resolved[path] = RowwiseParallel() elif style == "rowwise_split_input": - # `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. + # Same weight split as "rowwise", for a Linear whose input arrives whole instead of split by a "colwise" + # layer before it; each rank slices its own columns locally. E.g. an AdaLN modulation projection reading + # the timestep embedding (MiniMax-H3's `adaln_proj.linear`), or an output projection after a gathered + # QK-norm. Same as the transformers style of the same name. resolved[path] = RowwiseParallel(input_layouts=Replicate()) elif isinstance(style, PackedColwiseParallel): resolved[path] = _make_packed_col(style)