diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index f05667986f11..cb0dde8c9790 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -199,6 +199,8 @@ title: Intel Gaudi - local: optimization/neuron title: AWS Neuron + - local: optimization/tpu + title: TPU title: Model accelerators and hardware - isExpanded: false sections: diff --git a/docs/source/en/api/parallel.md b/docs/source/en/api/parallel.md index 5f300d5dd566..49c1ca3e4fb9 100644 --- a/docs/source/en/api/parallel.md +++ b/docs/source/en/api/parallel.md @@ -28,3 +28,13 @@ Parallelism strategies help speed up diffusion transformers by distributing comp [[autodoc]] TensorParallelConfig [[autodoc]] hooks.apply_tensor_parallel + +## TPU + +### `enable_tpu_compile` + +[[autodoc]] diffusers.DiffusionPipeline.enable_tpu_compile + +### `tpu_warmup` + +[[autodoc]] diffusers.DiffusionPipeline.tpu_warmup \ No newline at end of file diff --git a/docs/source/en/optimization/tpu.md b/docs/source/en/optimization/tpu.md new file mode 100644 index 000000000000..57090cc00eda --- /dev/null +++ b/docs/source/en/optimization/tpu.md @@ -0,0 +1,122 @@ + + +# TorchTPU + +[TorchTPU](https://github.com/google-pytorch/torch_tpu/) is a PyTorch backend for Google's Tensor Processing Units (TPUs), which lets you run Diffusers pipelines on Cloud TPUs (v6e, v5p, etc.) with minimal code changes. + +Two execution modes are available: + +| Mode | Constant | How to activate | Notes | +|---|---|---|---| +| Strict eager (default) | `EagerMode.DEFER_NEVER` | `import torch_tpu` | Operations dispatched one at a time, asynchronous | +| Compile | — | `pipe.enable_tpu_compile()` | AOT compilation with `TpuBackend` | + +Follow the [TorchTPU installation guide](https://github.com/google-pytorch/torch_tpu/). After installation, +`import torch_tpu` registers the `"tpu"` device automatically. + +## Eager mode + +```python +import gc +import torch +import torch_tpu # noqa: F401 + +from diffusers import FluxPipeline + +pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16) + +# 1. Encode on TPU. +pipe.text_encoder.to("tpu") +pipe.text_encoder_2.to("tpu") +with torch.no_grad(): + prompt_embeds, pooled_prompt_embeds, _ = pipe.encode_prompt( + prompt="a golden retriever surfing a wave, photorealistic", + prompt_2="a golden retriever surfing a wave, photorealistic", + device=torch.device("tpu"), + max_sequence_length=512, + ) + +# 2. Free the text encoders — nothing below needs them. +pipe.text_encoder = None +pipe.text_encoder_2 = None +gc.collect() + +# 3. Move the transformer and VAE in, then denoise with the precomputed embeddings. +pipe.transformer.to("tpu") +pipe.vae.to("tpu") +image = pipe( + prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + height=1024, + width=1024, + num_inference_steps=4, + guidance_scale=0.0, +).images[0] + +image.save("output.png") +``` + +If the text encoder alone is too large for a single chip(eg. FLUX.2-dev's Mistral-3-Small is ~45GB), +shard it across multiple chips with [`~diffusers.hooks.tensor_parallel.apply_tensor_parallel`], the +same mechanism [`~ModelMixin.enable_parallelism`] uses for the transformer (see [Tensor +parallelism](../training/distributed_inference#tensor-parallelism)). It only requires `model: +torch.nn.Module`, so it works directly on a `transformers.PreTrainedModel` text encoder too, not +just a diffusers `ModelMixin`. The text encoder doesn't define a `_tp_plan`, so supply one: pair +each attention/MLP projection that expands the hidden dimension (`"colwise"`) with the one that +contracts it back (`"rowwise"`), matching the `transformers` model's actual module names. + +## Compiled mode + +[`enable_tpu_compile`] runs `torch.compile` with `TpuBackend` on each pipeline module that is already on TPU. The first call (warmup) is slow because it compiles. Later calls reuse the compiled graph. Where it's supported, it replaces SDP-based attention with `AttnProcessor` for XLA tracing. + +> [!IMPORTANT] +> TorchTPU requires **static shapes** — `torch.compile` is called with `dynamic=False` +> internally. Every time `height`, `width`, or `num_inference_steps` changes, the graph is +> recompiled from scratch. Keep these values constant across all calls after warmup, or call +> [`tpu_warmup`] again before changing them. + +```python +import torch +import torch_tpu # noqa: F401 + +from diffusers import FluxPipeline + +pipe = FluxPipeline.from_pretrained( + "black-forest-labs/FLUX.1-schnell", + torch_dtype=torch.bfloat16, +) +pipe.transformer.to("tpu") +pipe.vae.to("tpu") + +pipe.enable_tpu_compile() + +# Warmup — triggers static graph compilation. +pipe.tpu_warmup( + prompt="warmup", + height=1024, + width=1024, + num_inference_steps=4, + guidance_scale=0.0, +) + +# Timed inference reuses the compiled graph. +image = pipe( + prompt="a golden retriever surfing a wave, photorealistic", + height=1024, + width=1024, + num_inference_steps=4, + guidance_scale=0.0, +).images[0] + +image.save("output.png") +``` diff --git a/src/diffusers/hooks/tensor_parallel.py b/src/diffusers/hooks/tensor_parallel.py index b90a5761d043..c21a4cc66bea 100644 --- a/src/diffusers/hooks/tensor_parallel.py +++ b/src/diffusers/hooks/tensor_parallel.py @@ -20,7 +20,7 @@ logger = get_logger(__name__) # pylint: disable=invalid-name -_SUPPORTED_TP_DEVICES = ("cuda", "neuron") +_SUPPORTED_TP_DEVICES = ("cuda", "neuron", "tpu") class PackedColwiseParallel: @@ -257,7 +257,12 @@ 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" + if tp_mesh.device_type == "neuron": + backend = "neuron" + elif tp_mesh.device_type == "tpu": + backend = "tpu" + else: + backend = "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}.") @@ -267,6 +272,12 @@ def apply_tensor_parallel( _apply_tp_neuron(model, tp_mesh, groups) return + if backend == "tpu": + from .tensor_parallel_tpu import _apply_tp_tpu + + _apply_tp_tpu(model, tp_mesh, groups) + return + from torch.distributed.tensor.parallel import parallelize_module for submodule, relative_plan in groups: diff --git a/src/diffusers/hooks/tensor_parallel_tpu.py b/src/diffusers/hooks/tensor_parallel_tpu.py new file mode 100644 index 000000000000..6805e5814523 --- /dev/null +++ b/src/diffusers/hooks/tensor_parallel_tpu.py @@ -0,0 +1,116 @@ +# 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. + +"""TPU backend for tensor parallelism, dispatched from ``apply_tensor_parallel(backend="tpu")``. + +The structure mirrors the Neuron backend (``tensor_parallel_neuron.py``). The motivation differs: on Neuron the +pre-shard path works around an NRT consecutive-reduce-scatter bug; here it prevents OOM. Without pre-sharding, +``parallelize_module`` calls ``distribute_tensor`` internally, which loads the full weight matrix on every TPU chip +before scattering. For large diffusion models this exhausts HBM. Pre-sharding each weight on CPU via +``DTensor.from_local`` first means each chip only receives its local shard, then ``parallelize_module`` is called as a +no-op for weights (they are already DTensors) but still registers the input/output hooks for the forward pass. +""" + +import torch +import torch.distributed as dist +import torch.nn as nn + + +def _pre_shard_and_tp( + module: nn.Module, + tp_mesh: "torch.distributed.device_mesh.DeviceMesh", + relative_plan: dict, + rank: int, + tp_size: int, +) -> None: + """Pre-shard plain colwise/rowwise Linear weights via ``DTensor.from_local``, then call ``parallelize_module``. + + Args: + module: The block whose Linear sub-modules are being sharded. + tp_mesh: Device mesh for TP (1-D, size == tp_size). + relative_plan: ``{relative_path: "colwise" | "rowwise" | PackedColwiseParallel | PackedRowwiseParallel}``, + the raw per-block plan as produced by ``_resolve_tp_plan`` (before ``_styles`` resolves it to + `parallelize_module` style instances). + rank: Current rank (``dist.get_rank()``). + tp_size: Total TP degree (``tp_mesh.size()``). + """ + from torch.distributed.tensor import DTensor, Shard + from torch.distributed.tensor.parallel import ColwiseParallel, RowwiseParallel, parallelize_module + + from .tensor_parallel import PackedColwiseParallel, PackedRowwiseParallel, _styles + + # Each torchrun worker owns one TPU chip (its local device). Use "tpu" without an explicit + # index — specifying tpu:rank would address chip `rank` from the current process's view, + # which fails because each worker only has access to its own assigned chip. + device = torch.device("tpu") + styles = _styles(relative_plan) + + for path, raw_style in relative_plan.items(): + # Packed (fused) projections resolve to a `_partition_linear_fn` that first replicates the full weight + # before re-slicing it into blocks (see `_styles` in tensor_parallel.py). Pre-sharding it here would + # hand that function an already-sharded DTensor with a placement it never asked for, which raises + # "Cannot distribute a DTensor with placements (Shard(dim=0),) to a different placements [Replicate()]". + # Leave these to `parallelize_module`, which materializes and shards them directly. + if isinstance(raw_style, (PackedColwiseParallel, PackedRowwiseParallel)): + continue + + submod = module + for part in path.split("."): + submod = getattr(submod, part) + + if not hasattr(submod, "weight"): + continue + + w = submod.weight.data # CPU at this point + style = styles[path] + if isinstance(style, ColwiseParallel): + rows = w.shape[0] // tp_size + shard = w[rank * rows : (rank + 1) * rows, :].contiguous().to(device) + submod.weight = nn.Parameter(DTensor.from_local(shard, tp_mesh, [Shard(0)])) + elif isinstance(style, RowwiseParallel): + cols = w.shape[1] // tp_size + shard = w[:, rank * cols : (rank + 1) * cols].contiguous().to(device) + submod.weight = nn.Parameter(DTensor.from_local(shard, tp_mesh, [Shard(1)])) + + # parallelize_module is now a no-op for weight distribution on the modules pre-sharded above (their + # weights are already DTensors with a matching placement), but still registers the input/output hooks + # required for the forward pass, and fully materializes+shards the packed modules skipped above. + parallelize_module(module, tp_mesh, styles) + + +def _apply_tp_tpu( + model: nn.Module, + tp_mesh: "torch.distributed.device_mesh.DeviceMesh", + groups: list, +) -> None: + """Apply tensor parallelism on TPU from resolved ``_tp_plan`` groups. + + ``groups`` is produced by ``diffusers.hooks.tensor_parallel._resolve_tp_plan`` — the same source of truth used by + the generic path, so the two backends shard identical layers. For each ``(block, relative_plan)`` group this: + 1. permutes the model's fused weights (via ``model._tp_fused_block_permuters``, the same backend-agnostic permuters + the generic path uses) so column/row slicing gives each rank a correct chunk, + 2. pre-shards the weights via ``DTensor.from_local`` (avoids full-weight materialisation on TPU HBM), then calls + ``parallelize_module`` to register the forward hooks. + + Model weights must be on CPU when this is called. + """ + rank = dist.get_rank() + tp_size = tp_mesh.size() + permuters = getattr(model, "_tp_fused_block_permuters", None) or {} + + for block, relative_plan in groups: + permuter = permuters.get(block.__class__.__name__) + if permuter is not None: + permuter(block, tp_size) + _pre_shard_and_tp(block, tp_mesh, relative_plan, rank, tp_size) diff --git a/src/diffusers/modular_pipelines/flux2/encoders.py b/src/diffusers/modular_pipelines/flux2/encoders.py index 09615c4becb6..99d7f92547fa 100644 --- a/src/diffusers/modular_pipelines/flux2/encoders.py +++ b/src/diffusers/modular_pipelines/flux2/encoders.py @@ -330,8 +330,9 @@ def _get_qwen3_prompt_embeds( all_input_ids.append(inputs["input_ids"]) all_attention_masks.append(inputs["attention_mask"]) - input_ids = torch.cat(all_input_ids, dim=0).to(device) - attention_mask = torch.cat(all_attention_masks, dim=0).to(device) + model_device = text_encoder.device + input_ids = torch.cat(all_input_ids, dim=0).to(model_device) + attention_mask = torch.cat(all_attention_masks, dim=0).to(model_device) # Forward pass through the model output = text_encoder( @@ -471,8 +472,9 @@ def _get_qwen3_prompt_embeds( all_input_ids.append(inputs["input_ids"]) all_attention_masks.append(inputs["attention_mask"]) - input_ids = torch.cat(all_input_ids, dim=0).to(device) - attention_mask = torch.cat(all_attention_masks, dim=0).to(device) + model_device = text_encoder.device + input_ids = torch.cat(all_input_ids, dim=0).to(model_device) + attention_mask = torch.cat(all_attention_masks, dim=0).to(model_device) # Forward pass through the model output = text_encoder( diff --git a/src/diffusers/pipelines/anyflow/pipeline_anyflow.py b/src/diffusers/pipelines/anyflow/pipeline_anyflow.py index c3e1dbf3a459..39d47e0e9174 100644 --- a/src/diffusers/pipelines/anyflow/pipeline_anyflow.py +++ b/src/diffusers/pipelines/anyflow/pipeline_anyflow.py @@ -161,7 +161,8 @@ def _get_t5_prompt_embeds( text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask seq_lens = mask.gt(0).sum(dim=1).long() - prompt_embeds = self.text_encoder(text_input_ids.to(device), mask.to(device)).last_hidden_state + model_device = self.text_encoder.device + prompt_embeds = self.text_encoder(text_input_ids.to(model_device), mask.to(model_device)).last_hidden_state prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] prompt_embeds = torch.stack( diff --git a/src/diffusers/pipelines/anyflow/pipeline_anyflow_far.py b/src/diffusers/pipelines/anyflow/pipeline_anyflow_far.py index 96edc07a0043..e6179bf69076 100644 --- a/src/diffusers/pipelines/anyflow/pipeline_anyflow_far.py +++ b/src/diffusers/pipelines/anyflow/pipeline_anyflow_far.py @@ -178,7 +178,8 @@ def _get_t5_prompt_embeds( text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask seq_lens = mask.gt(0).sum(dim=1).long() - prompt_embeds = self.text_encoder(text_input_ids.to(device), mask.to(device)).last_hidden_state + model_device = self.text_encoder.device + prompt_embeds = self.text_encoder(text_input_ids.to(model_device), mask.to(model_device)).last_hidden_state prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] prompt_embeds = torch.stack( diff --git a/src/diffusers/pipelines/chronoedit/pipeline_chronoedit.py b/src/diffusers/pipelines/chronoedit/pipeline_chronoedit.py index 1e0cc0ea5c2a..bb7ea1cfe464 100644 --- a/src/diffusers/pipelines/chronoedit/pipeline_chronoedit.py +++ b/src/diffusers/pipelines/chronoedit/pipeline_chronoedit.py @@ -210,7 +210,8 @@ def _get_t5_prompt_embeds( text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask seq_lens = mask.gt(0).sum(dim=1).long() - prompt_embeds = self.text_encoder(text_input_ids.to(device), mask.to(device)).last_hidden_state + model_device = self.text_encoder.device + prompt_embeds = self.text_encoder(text_input_ids.to(model_device), mask.to(model_device)).last_hidden_state prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] prompt_embeds = torch.stack( diff --git a/src/diffusers/pipelines/ernie_image/pipeline_ernie_image.py b/src/diffusers/pipelines/ernie_image/pipeline_ernie_image.py index 11fce6a204bf..df5b27ace653 100644 --- a/src/diffusers/pipelines/ernie_image/pipeline_ernie_image.py +++ b/src/diffusers/pipelines/ernie_image/pipeline_ernie_image.py @@ -114,7 +114,7 @@ def _enhance_prompt_with_pe( tokenize=False, add_generation_prompt=False, # "Output:" is already in the user block ) - inputs = self.pe_tokenizer(input_text, return_tensors="pt").to(device) + inputs = self.pe_tokenizer(input_text, return_tensors="pt").to(self.pe.device) output_ids = self.pe.generate( **inputs, max_new_tokens=self.pe_tokenizer.model_max_length, @@ -155,7 +155,7 @@ def encode_prompt( else: ids = [0] - input_ids = torch.tensor([ids], device=device) + input_ids = torch.tensor([ids], device=self.text_encoder.device) with torch.no_grad(): outputs = self.text_encoder( input_ids=input_ids, diff --git a/src/diffusers/pipelines/flux/pipeline_flux.py b/src/diffusers/pipelines/flux/pipeline_flux.py index d3e0682c5419..016e557ad19d 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux.py +++ b/src/diffusers/pipelines/flux/pipeline_flux.py @@ -250,7 +250,8 @@ def _get_t5_prompt_embeds( f" {max_sequence_length} tokens: {removed_text}" ) - prompt_embeds = self.text_encoder_2(text_input_ids.to(device), output_hidden_states=False)[0] + model_device = self.text_encoder_2.device + prompt_embeds = self.text_encoder_2(text_input_ids.to(model_device), output_hidden_states=False)[0] dtype = self.text_encoder_2.dtype prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) @@ -295,7 +296,8 @@ def _get_clip_prompt_embeds( "The following part of your input was truncated because CLIP can only handle sequences up to" f" {self.tokenizer_max_length} tokens: {removed_text}" ) - prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False) + model_device = self.text_encoder.device + prompt_embeds = self.text_encoder(text_input_ids.to(model_device), output_hidden_states=False) # Use pooled output of CLIPTextModel prompt_embeds = prompt_embeds.pooler_output diff --git a/src/diffusers/pipelines/flux/pipeline_flux_control.py b/src/diffusers/pipelines/flux/pipeline_flux_control.py index 46671c44cca8..f8e270cba83b 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_control.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_control.py @@ -262,7 +262,8 @@ def _get_t5_prompt_embeds( f" {max_sequence_length} tokens: {removed_text}" ) - prompt_embeds = self.text_encoder_2(text_input_ids.to(device), output_hidden_states=False)[0] + model_device = self.text_encoder_2.device + prompt_embeds = self.text_encoder_2(text_input_ids.to(model_device), output_hidden_states=False)[0] dtype = self.text_encoder_2.dtype prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) @@ -308,7 +309,8 @@ def _get_clip_prompt_embeds( "The following part of your input was truncated because CLIP can only handle sequences up to" f" {self.tokenizer_max_length} tokens: {removed_text}" ) - prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False) + model_device = self.text_encoder.device + prompt_embeds = self.text_encoder(text_input_ids.to(model_device), output_hidden_states=False) # Use pooled output of CLIPTextModel prompt_embeds = prompt_embeds.pooler_output diff --git a/src/diffusers/pipelines/flux/pipeline_flux_control_img2img.py b/src/diffusers/pipelines/flux/pipeline_flux_control_img2img.py index b455c611e0ae..663390114d3e 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_control_img2img.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_control_img2img.py @@ -273,7 +273,8 @@ def _get_t5_prompt_embeds( f" {max_sequence_length} tokens: {removed_text}" ) - prompt_embeds = self.text_encoder_2(text_input_ids.to(device), output_hidden_states=False)[0] + model_device = self.text_encoder_2.device + prompt_embeds = self.text_encoder_2(text_input_ids.to(model_device), output_hidden_states=False)[0] dtype = self.text_encoder_2.dtype prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) @@ -319,7 +320,8 @@ def _get_clip_prompt_embeds( "The following part of your input was truncated because CLIP can only handle sequences up to" f" {self.tokenizer_max_length} tokens: {removed_text}" ) - prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False) + model_device = self.text_encoder.device + prompt_embeds = self.text_encoder(text_input_ids.to(model_device), output_hidden_states=False) # Use pooled output of CLIPTextModel prompt_embeds = prompt_embeds.pooler_output diff --git a/src/diffusers/pipelines/flux/pipeline_flux_control_inpaint.py b/src/diffusers/pipelines/flux/pipeline_flux_control_inpaint.py index 15e27653c3e2..44a0b85336e4 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_control_inpaint.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_control_inpaint.py @@ -312,7 +312,8 @@ def _get_t5_prompt_embeds( f" {max_sequence_length} tokens: {removed_text}" ) - prompt_embeds = self.text_encoder_2(text_input_ids.to(device), output_hidden_states=False)[0] + model_device = self.text_encoder_2.device + prompt_embeds = self.text_encoder_2(text_input_ids.to(model_device), output_hidden_states=False)[0] dtype = self.text_encoder_2.dtype prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) @@ -358,7 +359,8 @@ def _get_clip_prompt_embeds( "The following part of your input was truncated because CLIP can only handle sequences up to" f" {self.tokenizer_max_length} tokens: {removed_text}" ) - prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False) + model_device = self.text_encoder.device + prompt_embeds = self.text_encoder(text_input_ids.to(model_device), output_hidden_states=False) # Use pooled output of CLIPTextModel prompt_embeds = prompt_embeds.pooler_output diff --git a/src/diffusers/pipelines/flux/pipeline_flux_controlnet.py b/src/diffusers/pipelines/flux/pipeline_flux_controlnet.py index e7792d667f16..7bedf12a7c7c 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_controlnet.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_controlnet.py @@ -282,7 +282,8 @@ def _get_t5_prompt_embeds( f" {max_sequence_length} tokens: {removed_text}" ) - prompt_embeds = self.text_encoder_2(text_input_ids.to(device), output_hidden_states=False)[0] + model_device = self.text_encoder_2.device + prompt_embeds = self.text_encoder_2(text_input_ids.to(model_device), output_hidden_states=False)[0] dtype = self.text_encoder_2.dtype prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) @@ -327,7 +328,8 @@ def _get_clip_prompt_embeds( "The following part of your input was truncated because CLIP can only handle sequences up to" f" {self.tokenizer_max_length} tokens: {removed_text}" ) - prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False) + model_device = self.text_encoder.device + prompt_embeds = self.text_encoder(text_input_ids.to(model_device), output_hidden_states=False) # Use pooled output of CLIPTextModel prompt_embeds = prompt_embeds.pooler_output diff --git a/src/diffusers/pipelines/flux/pipeline_flux_controlnet_image_to_image.py b/src/diffusers/pipelines/flux/pipeline_flux_controlnet_image_to_image.py index 61c9da0c9496..be3253c6e634 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_controlnet_image_to_image.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_controlnet_image_to_image.py @@ -274,7 +274,8 @@ def _get_t5_prompt_embeds( f" {max_sequence_length} tokens: {removed_text}" ) - prompt_embeds = self.text_encoder_2(text_input_ids.to(device), output_hidden_states=False)[0] + model_device = self.text_encoder_2.device + prompt_embeds = self.text_encoder_2(text_input_ids.to(model_device), output_hidden_states=False)[0] dtype = self.text_encoder_2.dtype prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) @@ -320,7 +321,8 @@ def _get_clip_prompt_embeds( "The following part of your input was truncated because CLIP can only handle sequences up to" f" {self.tokenizer_max_length} tokens: {removed_text}" ) - prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False) + model_device = self.text_encoder.device + prompt_embeds = self.text_encoder(text_input_ids.to(model_device), output_hidden_states=False) # Use pooled output of CLIPTextModel prompt_embeds = prompt_embeds.pooler_output diff --git a/src/diffusers/pipelines/flux/pipeline_flux_controlnet_inpainting.py b/src/diffusers/pipelines/flux/pipeline_flux_controlnet_inpainting.py index eed671152bc9..b93c3f969f91 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_controlnet_inpainting.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_controlnet_inpainting.py @@ -285,7 +285,8 @@ def _get_t5_prompt_embeds( f" {max_sequence_length} tokens: {removed_text}" ) - prompt_embeds = self.text_encoder_2(text_input_ids.to(device), output_hidden_states=False)[0] + model_device = self.text_encoder_2.device + prompt_embeds = self.text_encoder_2(text_input_ids.to(model_device), output_hidden_states=False)[0] dtype = self.text_encoder_2.dtype prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) @@ -331,7 +332,8 @@ def _get_clip_prompt_embeds( "The following part of your input was truncated because CLIP can only handle sequences up to" f" {self.tokenizer_max_length} tokens: {removed_text}" ) - prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False) + model_device = self.text_encoder.device + prompt_embeds = self.text_encoder(text_input_ids.to(model_device), output_hidden_states=False) # Use pooled output of CLIPTextModel prompt_embeds = prompt_embeds.pooler_output diff --git a/src/diffusers/pipelines/flux/pipeline_flux_fill.py b/src/diffusers/pipelines/flux/pipeline_flux_fill.py index ab4431b5b768..a3f6ffc9349b 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_fill.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_fill.py @@ -277,7 +277,8 @@ def _get_t5_prompt_embeds( f" {max_sequence_length} tokens: {removed_text}" ) - prompt_embeds = self.text_encoder_2(text_input_ids.to(device), output_hidden_states=False)[0] + model_device = self.text_encoder_2.device + prompt_embeds = self.text_encoder_2(text_input_ids.to(model_device), output_hidden_states=False)[0] dtype = self.text_encoder_2.dtype prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) @@ -323,7 +324,8 @@ def _get_clip_prompt_embeds( "The following part of your input was truncated because CLIP can only handle sequences up to" f" {self.tokenizer_max_length} tokens: {removed_text}" ) - prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False) + model_device = self.text_encoder.device + prompt_embeds = self.text_encoder(text_input_ids.to(model_device), output_hidden_states=False) # Use pooled output of CLIPTextModel prompt_embeds = prompt_embeds.pooler_output diff --git a/src/diffusers/pipelines/flux/pipeline_flux_img2img.py b/src/diffusers/pipelines/flux/pipeline_flux_img2img.py index 94582a84cf84..a08355ea426a 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_img2img.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_img2img.py @@ -271,7 +271,8 @@ def _get_t5_prompt_embeds( f" {max_sequence_length} tokens: {removed_text}" ) - prompt_embeds = self.text_encoder_2(text_input_ids.to(device), output_hidden_states=False)[0] + model_device = self.text_encoder_2.device + prompt_embeds = self.text_encoder_2(text_input_ids.to(model_device), output_hidden_states=False)[0] dtype = self.text_encoder_2.dtype prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) @@ -317,7 +318,8 @@ def _get_clip_prompt_embeds( "The following part of your input was truncated because CLIP can only handle sequences up to" f" {self.tokenizer_max_length} tokens: {removed_text}" ) - prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False) + model_device = self.text_encoder.device + prompt_embeds = self.text_encoder(text_input_ids.to(model_device), output_hidden_states=False) # Use pooled output of CLIPTextModel prompt_embeds = prompt_embeds.pooler_output diff --git a/src/diffusers/pipelines/flux/pipeline_flux_inpaint.py b/src/diffusers/pipelines/flux/pipeline_flux_inpaint.py index 4c35ffefe088..8935ae114c78 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_inpaint.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_inpaint.py @@ -275,7 +275,8 @@ def _get_t5_prompt_embeds( f" {max_sequence_length} tokens: {removed_text}" ) - prompt_embeds = self.text_encoder_2(text_input_ids.to(device), output_hidden_states=False)[0] + model_device = self.text_encoder_2.device + prompt_embeds = self.text_encoder_2(text_input_ids.to(model_device), output_hidden_states=False)[0] dtype = self.text_encoder_2.dtype prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) @@ -321,7 +322,8 @@ def _get_clip_prompt_embeds( "The following part of your input was truncated because CLIP can only handle sequences up to" f" {self.tokenizer_max_length} tokens: {removed_text}" ) - prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False) + model_device = self.text_encoder.device + prompt_embeds = self.text_encoder(text_input_ids.to(model_device), output_hidden_states=False) # Use pooled output of CLIPTextModel prompt_embeds = prompt_embeds.pooler_output diff --git a/src/diffusers/pipelines/flux/pipeline_flux_kontext.py b/src/diffusers/pipelines/flux/pipeline_flux_kontext.py index 849d9686de62..7b7d58e8848a 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_kontext.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_kontext.py @@ -296,7 +296,8 @@ def _get_t5_prompt_embeds( f" {max_sequence_length} tokens: {removed_text}" ) - prompt_embeds = self.text_encoder_2(text_input_ids.to(device), output_hidden_states=False)[0] + model_device = self.text_encoder_2.device + prompt_embeds = self.text_encoder_2(text_input_ids.to(model_device), output_hidden_states=False)[0] dtype = self.text_encoder_2.dtype prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) @@ -342,7 +343,8 @@ def _get_clip_prompt_embeds( "The following part of your input was truncated because CLIP can only handle sequences up to" f" {self.tokenizer_max_length} tokens: {removed_text}" ) - prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False) + model_device = self.text_encoder.device + prompt_embeds = self.text_encoder(text_input_ids.to(model_device), output_hidden_states=False) # Use pooled output of CLIPTextModel prompt_embeds = prompt_embeds.pooler_output diff --git a/src/diffusers/pipelines/flux/pipeline_flux_kontext_inpaint.py b/src/diffusers/pipelines/flux/pipeline_flux_kontext_inpaint.py index 982581c01b3c..4a3bd7c2f7c5 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_kontext_inpaint.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_kontext_inpaint.py @@ -329,7 +329,8 @@ def _get_t5_prompt_embeds( f" {max_sequence_length} tokens: {removed_text}" ) - prompt_embeds = self.text_encoder_2(text_input_ids.to(device), output_hidden_states=False)[0] + model_device = self.text_encoder_2.device + prompt_embeds = self.text_encoder_2(text_input_ids.to(model_device), output_hidden_states=False)[0] dtype = self.text_encoder_2.dtype prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) @@ -375,7 +376,8 @@ def _get_clip_prompt_embeds( "The following part of your input was truncated because CLIP can only handle sequences up to" f" {self.tokenizer_max_length} tokens: {removed_text}" ) - prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False) + model_device = self.text_encoder.device + prompt_embeds = self.text_encoder(text_input_ids.to(model_device), output_hidden_states=False) # Use pooled output of CLIPTextModel prompt_embeds = prompt_embeds.pooler_output diff --git a/src/diffusers/pipelines/flux/pipeline_flux_prior_redux.py b/src/diffusers/pipelines/flux/pipeline_flux_prior_redux.py index f173fdef88c6..08861437f99f 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_prior_redux.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_prior_redux.py @@ -234,7 +234,8 @@ def _get_t5_prompt_embeds( f" {max_sequence_length} tokens: {removed_text}" ) - prompt_embeds = self.text_encoder_2(text_input_ids.to(device), output_hidden_states=False)[0] + model_device = self.text_encoder_2.device + prompt_embeds = self.text_encoder_2(text_input_ids.to(model_device), output_hidden_states=False)[0] dtype = self.text_encoder_2.dtype prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) @@ -280,7 +281,8 @@ def _get_clip_prompt_embeds( "The following part of your input was truncated because CLIP can only handle sequences up to" f" {self.tokenizer_max_length} tokens: {removed_text}" ) - prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False) + model_device = self.text_encoder.device + prompt_embeds = self.text_encoder(text_input_ids.to(model_device), output_hidden_states=False) # Use pooled output of CLIPTextModel prompt_embeds = prompt_embeds.pooler_output diff --git a/src/diffusers/pipelines/flux2/pipeline_flux2_klein.py b/src/diffusers/pipelines/flux2/pipeline_flux2_klein.py index d768e6127f26..98a87eacbbeb 100644 --- a/src/diffusers/pipelines/flux2/pipeline_flux2_klein.py +++ b/src/diffusers/pipelines/flux2/pipeline_flux2_klein.py @@ -241,8 +241,9 @@ def _get_qwen3_prompt_embeds( all_input_ids.append(inputs["input_ids"]) all_attention_masks.append(inputs["attention_mask"]) - input_ids = torch.cat(all_input_ids, dim=0).to(device) - attention_mask = torch.cat(all_attention_masks, dim=0).to(device) + model_device = text_encoder.device + input_ids = torch.cat(all_input_ids, dim=0).to(model_device) + attention_mask = torch.cat(all_attention_masks, dim=0).to(model_device) # Forward pass through the model output = text_encoder( diff --git a/src/diffusers/pipelines/flux2/pipeline_flux2_klein_inpaint.py b/src/diffusers/pipelines/flux2/pipeline_flux2_klein_inpaint.py index fd9467003a71..3f0b42d15f90 100644 --- a/src/diffusers/pipelines/flux2/pipeline_flux2_klein_inpaint.py +++ b/src/diffusers/pipelines/flux2/pipeline_flux2_klein_inpaint.py @@ -288,8 +288,9 @@ def _get_qwen3_prompt_embeds( all_input_ids.append(inputs["input_ids"]) all_attention_masks.append(inputs["attention_mask"]) - input_ids = torch.cat(all_input_ids, dim=0).to(device) - attention_mask = torch.cat(all_attention_masks, dim=0).to(device) + model_device = text_encoder.device + input_ids = torch.cat(all_input_ids, dim=0).to(model_device) + attention_mask = torch.cat(all_attention_masks, dim=0).to(model_device) # Forward pass through the model output = text_encoder( diff --git a/src/diffusers/pipelines/lucy/pipeline_lucy_edit.py b/src/diffusers/pipelines/lucy/pipeline_lucy_edit.py index 69eb2a02be5c..b9c9a8fb74f0 100644 --- a/src/diffusers/pipelines/lucy/pipeline_lucy_edit.py +++ b/src/diffusers/pipelines/lucy/pipeline_lucy_edit.py @@ -221,7 +221,8 @@ def _get_t5_prompt_embeds( text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask seq_lens = mask.gt(0).sum(dim=1).long() - prompt_embeds = self.text_encoder(text_input_ids.to(device), mask.to(device)).last_hidden_state + model_device = self.text_encoder.device + prompt_embeds = self.text_encoder(text_input_ids.to(model_device), mask.to(model_device)).last_hidden_state prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] prompt_embeds = torch.stack( diff --git a/src/diffusers/pipelines/pipeline_utils.py b/src/diffusers/pipelines/pipeline_utils.py index 8986553eda3d..eb4ea0b2d563 100644 --- a/src/diffusers/pipelines/pipeline_utils.py +++ b/src/diffusers/pipelines/pipeline_utils.py @@ -22,7 +22,7 @@ import types from dataclasses import dataclass from pathlib import Path -from typing import Any, Callable, Dict, List, Union, get_args, get_origin, get_type_hints +from typing import Any, Callable, Dict, List, Optional, Union, get_args, get_origin, get_type_hints import httpx import numpy as np @@ -73,6 +73,7 @@ is_transformers_version, logging, numpy_to_pil, + requires_backends, ) from ..utils.distributed_utils import is_torch_dist_rank_zero from ..utils.hub_utils import ( @@ -195,6 +196,37 @@ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) +def _supports_generic_attn_processor(module: "torch.nn.Module") -> bool: + """Whether every attention submodule in `module` accepts the generic, non-SDPA `AttnProcessor`. + + Two independent module shapes are incompatible with the generic `AttnProcessor`: + + 1. Newer attention classes (`AttentionModuleMixin` subclasses, e.g. `FluxAttention`, `Flux2Attention`, + `WanAttention`) declare a fixed `_available_processors` list of model-specific processor classes; the + generic `AttnProcessor` isn't among them, because its `__call__` assumes attributes (e.g. `spatial_norm`) + only the legacy `Attention` class defines — forcing it onto one of these raises `AttributeError` on the + very next forward pass. Those models' own default processors (`FluxAttnProcessor`, `Flux2AttnProcessor`, + `WanAttnProcessor`, ...) are SDPA-based too, so they're not a safe substitute for the crash + `enable_tpu_compile` is working around either — the fix is simply to leave this class of module on its own + native processor, which compiles under `TpuBackend` without incident. + 2. Legacy `Attention` modules configured for joint/dual-stream attention (`added_kv_proj_dim` set, e.g. + QwenImage's transformer blocks) use a custom processor (e.g. `QwenDoubleStreamAttnProcessor2_0`) that + returns a separate `(image, text)` output pair. The generic `AttnProcessor` returns a single tensor, so the + caller's unpacking of the paired output raises `ValueError` — this is a functional protocol mismatch, not + merely a numerical one, and no `_available_processors` restriction catches it since these aren't + `AttentionModuleMixin` subclasses. + """ + from ..models.attention_processor import AttnProcessor + + for submodule in module.modules(): + available = getattr(submodule, "_available_processors", None) + if available is not None and AttnProcessor not in available: + return False + if getattr(submodule, "added_kv_proj_dim", None) is not None: + return False + return True + + class DiffusionPipeline(ConfigMixin, PushToHubMixin): r""" Base class for all pipelines. @@ -2257,6 +2289,91 @@ def _is_pipeline_device_mapped(self): return not is_device_type_map and isinstance(device_map, dict) and len(device_map) > 1 + def enable_tpu_compile( + self, + model_names: Optional[List[str]] = None, + **compile_kwargs, + ) -> None: + """Compile pipeline components that are on TPU using ``torch.compile`` with the ``TpuBackend``. + + Before compiling, each component that exposes ``set_attn_processor`` has ``AttnProcessor`` applied. This + replaces ``AttnProcessor2_0`` (SDP-based) which triggers XLA fusion-emitter crashes in eager/lazy mode. + ``TpuBackend`` handles the resulting ``torch.cat`` layout internally during static tracing, so no additional + wrapper is needed at compile time. + + Args: + model_names (`list[str]`, *optional*): + Names of pipeline components to compile. Defaults to all ``torch.nn.Module`` components currently + resident on a TPU device. + **compile_kwargs: + Extra keyword arguments forwarded to ``torch.compile``. ``backend`` defaults to ``TpuBackend()`` and + ``dynamic`` defaults to ``False`` (required for static tracing). + + Example: + ```python + import torch + import torch_tpu # noqa: F401 + + pipe.transformer.to("tpu") + pipe.vae.to("tpu") + pipe.enable_tpu_compile() + ``` + """ + requires_backends(self, "torch_tpu") + from torch_tpu._internal.compile import TpuBackend + + from ..models.attention_processor import AttnProcessor + + if model_names is None: + model_names = [ + name + for name, comp in self.components.items() + if isinstance(comp, torch.nn.Module) and comp.device.type == "tpu" + ] + + for name in model_names: + component = getattr(self, name, None) + if not isinstance(component, torch.nn.Module): + logger.warning(f"`enable_tpu_compile`: component '{name}' is not a nn.Module, skipping.") + continue + if is_compiled_module(component): + logger.warning(f"`enable_tpu_compile`: component '{name}' is already compiled, skipping.") + continue + if hasattr(component, "set_attn_processor") and _supports_generic_attn_processor(component): + component.set_attn_processor(AttnProcessor()) + compile_kwargs.setdefault("backend", TpuBackend()) + compile_kwargs.setdefault("dynamic", False) + logger.info(f"Compiling '{name}' with TpuBackend.") + setattr(self, name, torch.compile(component, **compile_kwargs)) + + def tpu_warmup(self, *args, **kwargs) -> None: + """Run a single forward pass to trigger XLA / ``TpuBackend`` compilation. + + Call this after ``enable_tpu_compile`` and before timed inference. The warmup pass compiles the static + computation graphs; subsequent calls reuse the compiled graphs and run at full speed. + + Args: + *args: Positional arguments forwarded to the pipeline ``__call__``. + **kwargs: Keyword arguments forwarded to the pipeline ``__call__``. + + Example: + ```python + pipe.tpu_warmup( + prompt="warmup", + height=1024, + width=1024, + num_inference_steps=4, + guidance_scale=0.0, + ) + ``` + """ + logger.info("Running TPU warmup pass to trigger XLA compilation...") + with torch.no_grad(): + self(*args, **kwargs) + if hasattr(torch, "tpu") and hasattr(torch.tpu, "synchronize"): + torch.tpu.synchronize() + logger.info("TPU warmup complete.") + class StableDiffusionMixin: r""" diff --git a/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2.py b/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2.py index faad0fb14086..848cbd6d200c 100644 --- a/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2.py +++ b/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2.py @@ -179,7 +179,8 @@ def _get_t5_prompt_embeds( text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask seq_lens = mask.gt(0).sum(dim=1).long() - prompt_embeds = self.text_encoder(text_input_ids.to(device), mask.to(device)).last_hidden_state + model_device = self.text_encoder.device + prompt_embeds = self.text_encoder(text_input_ids.to(model_device), mask.to(model_device)).last_hidden_state prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] prompt_embeds = torch.stack( diff --git a/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing.py b/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing.py index 8751240a1af9..fc4345fab832 100644 --- a/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing.py +++ b/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing.py @@ -200,7 +200,8 @@ def _get_t5_prompt_embeds( text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask seq_lens = mask.gt(0).sum(dim=1).long() - prompt_embeds = self.text_encoder(text_input_ids.to(device), mask.to(device)).last_hidden_state + model_device = self.text_encoder.device + prompt_embeds = self.text_encoder(text_input_ids.to(model_device), mask.to(model_device)).last_hidden_state prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] prompt_embeds = torch.stack( diff --git a/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing_i2v.py b/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing_i2v.py index 335a38e53004..1171b21eb36b 100644 --- a/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing_i2v.py +++ b/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing_i2v.py @@ -205,7 +205,8 @@ def _get_t5_prompt_embeds( text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask seq_lens = mask.gt(0).sum(dim=1).long() - prompt_embeds = self.text_encoder(text_input_ids.to(device), mask.to(device)).last_hidden_state + model_device = self.text_encoder.device + prompt_embeds = self.text_encoder(text_input_ids.to(model_device), mask.to(model_device)).last_hidden_state prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] prompt_embeds = torch.stack( diff --git a/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing_v2v.py b/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing_v2v.py index 810b0d019ddb..8701ab823889 100644 --- a/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing_v2v.py +++ b/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing_v2v.py @@ -261,7 +261,8 @@ def _get_t5_prompt_embeds( text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask seq_lens = mask.gt(0).sum(dim=1).long() - prompt_embeds = self.text_encoder(text_input_ids.to(device), mask.to(device)).last_hidden_state + model_device = self.text_encoder.device + prompt_embeds = self.text_encoder(text_input_ids.to(model_device), mask.to(model_device)).last_hidden_state prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] prompt_embeds = torch.stack( diff --git a/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_i2v.py b/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_i2v.py index 91c09a56fcfb..1aab5864c480 100644 --- a/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_i2v.py +++ b/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_i2v.py @@ -209,7 +209,8 @@ def _get_t5_prompt_embeds( text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask seq_lens = mask.gt(0).sum(dim=1).long() - prompt_embeds = self.text_encoder(text_input_ids.to(device), mask.to(device)).last_hidden_state + model_device = self.text_encoder.device + prompt_embeds = self.text_encoder(text_input_ids.to(model_device), mask.to(model_device)).last_hidden_state prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] prompt_embeds = torch.stack( diff --git a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py index ba94e3051fd3..fd7fcc9c7a15 100644 --- a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py +++ b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py @@ -402,7 +402,7 @@ def encode_prompt( f" {tokenizer.model_max_length} tokens: {removed_text}" ) - prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=True) + prompt_embeds = text_encoder(text_input_ids.to(text_encoder.device), output_hidden_states=True) # We are only ALWAYS interested in the pooled output of the final text encoder if pooled_prompt_embeds is None and prompt_embeds[0].ndim == 2: @@ -463,7 +463,7 @@ def encode_prompt( ) negative_prompt_embeds = text_encoder( - uncond_input.input_ids.to(device), + uncond_input.input_ids.to(text_encoder.device), output_hidden_states=True, ) diff --git a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py index c7a13ca02524..7147b71bab0f 100644 --- a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py +++ b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py @@ -420,7 +420,7 @@ def encode_prompt( f" {tokenizer.model_max_length} tokens: {removed_text}" ) - prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=True) + prompt_embeds = text_encoder(text_input_ids.to(text_encoder.device), output_hidden_states=True) # We are only ALWAYS interested in the pooled output of the final text encoder if pooled_prompt_embeds is None and prompt_embeds[0].ndim == 2: @@ -481,7 +481,7 @@ def encode_prompt( ) negative_prompt_embeds = text_encoder( - uncond_input.input_ids.to(device), + uncond_input.input_ids.to(text_encoder.device), output_hidden_states=True, ) diff --git a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py index 3f18cbe21d0f..4f67f0d81af5 100644 --- a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py +++ b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py @@ -524,7 +524,7 @@ def encode_prompt( f" {tokenizer.model_max_length} tokens: {removed_text}" ) - prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=True) + prompt_embeds = text_encoder(text_input_ids.to(text_encoder.device), output_hidden_states=True) # We are only ALWAYS interested in the pooled output of the final text encoder if pooled_prompt_embeds is None and prompt_embeds[0].ndim == 2: @@ -585,7 +585,7 @@ def encode_prompt( ) negative_prompt_embeds = text_encoder( - uncond_input.input_ids.to(device), + uncond_input.input_ids.to(text_encoder.device), output_hidden_states=True, ) diff --git a/src/diffusers/pipelines/visualcloze/pipeline_visualcloze_generation.py b/src/diffusers/pipelines/visualcloze/pipeline_visualcloze_generation.py index ed2cd519df25..5940d591166d 100644 --- a/src/diffusers/pipelines/visualcloze/pipeline_visualcloze_generation.py +++ b/src/diffusers/pipelines/visualcloze/pipeline_visualcloze_generation.py @@ -228,7 +228,8 @@ def _get_t5_prompt_embeds( f" {max_sequence_length} tokens: {removed_text}" ) - prompt_embeds = self.text_encoder_2(text_input_ids.to(device), output_hidden_states=False)[0] + model_device = self.text_encoder_2.device + prompt_embeds = self.text_encoder_2(text_input_ids.to(model_device), output_hidden_states=False)[0] dtype = self.text_encoder_2.dtype prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) @@ -274,7 +275,8 @@ def _get_clip_prompt_embeds( "The following part of your input was truncated because CLIP can only handle sequences up to" f" {self.tokenizer_max_length} tokens: {removed_text}" ) - prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False) + model_device = self.text_encoder.device + prompt_embeds = self.text_encoder(text_input_ids.to(model_device), output_hidden_states=False) # Use pooled output of CLIPTextModel prompt_embeds = prompt_embeds.pooler_output diff --git a/src/diffusers/pipelines/wan/pipeline_wan.py b/src/diffusers/pipelines/wan/pipeline_wan.py index be2d53f17932..b7e8a71cd50c 100644 --- a/src/diffusers/pipelines/wan/pipeline_wan.py +++ b/src/diffusers/pipelines/wan/pipeline_wan.py @@ -182,7 +182,8 @@ def _get_t5_prompt_embeds( text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask seq_lens = mask.gt(0).sum(dim=1).long() - prompt_embeds = self.text_encoder(text_input_ids.to(device), mask.to(device)).last_hidden_state + model_device = self.text_encoder.device + prompt_embeds = self.text_encoder(text_input_ids.to(model_device), mask.to(model_device)).last_hidden_state prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] prompt_embeds = torch.stack( @@ -654,7 +655,7 @@ def __call__( self._current_timestep = None if not output_type == "latent": - latents = latents.to(self.vae.dtype) + latents = latents.to(self.vae.device, dtype=self.vae.dtype) latents_mean = ( torch.tensor(self.vae.config.latents_mean) .view(1, self.vae.config.z_dim, 1, 1, 1) diff --git a/src/diffusers/pipelines/wan/pipeline_wan_animate.py b/src/diffusers/pipelines/wan/pipeline_wan_animate.py index 5806032c0142..91960fad7d36 100644 --- a/src/diffusers/pipelines/wan/pipeline_wan_animate.py +++ b/src/diffusers/pipelines/wan/pipeline_wan_animate.py @@ -259,7 +259,8 @@ def _get_t5_prompt_embeds( text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask seq_lens = mask.gt(0).sum(dim=1).long() - prompt_embeds = self.text_encoder(text_input_ids.to(device), mask.to(device)).last_hidden_state + model_device = self.text_encoder.device + prompt_embeds = self.text_encoder(text_input_ids.to(model_device), mask.to(model_device)).last_hidden_state prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] prompt_embeds = torch.stack( diff --git a/src/diffusers/pipelines/wan/pipeline_wan_i2v.py b/src/diffusers/pipelines/wan/pipeline_wan_i2v.py index 8061f67ab6b9..59a844e088ae 100644 --- a/src/diffusers/pipelines/wan/pipeline_wan_i2v.py +++ b/src/diffusers/pipelines/wan/pipeline_wan_i2v.py @@ -223,7 +223,8 @@ def _get_t5_prompt_embeds( text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask seq_lens = mask.gt(0).sum(dim=1).long() - prompt_embeds = self.text_encoder(text_input_ids.to(device), mask.to(device)).last_hidden_state + model_device = self.text_encoder.device + prompt_embeds = self.text_encoder(text_input_ids.to(model_device), mask.to(model_device)).last_hidden_state prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] prompt_embeds = torch.stack( diff --git a/src/diffusers/pipelines/wan/pipeline_wan_vace.py b/src/diffusers/pipelines/wan/pipeline_wan_vace.py index b0896d382d67..8c72adf09d6a 100644 --- a/src/diffusers/pipelines/wan/pipeline_wan_vace.py +++ b/src/diffusers/pipelines/wan/pipeline_wan_vace.py @@ -228,7 +228,8 @@ def _get_t5_prompt_embeds( text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask seq_lens = mask.gt(0).sum(dim=1).long() - prompt_embeds = self.text_encoder(text_input_ids.to(device), mask.to(device)).last_hidden_state + model_device = self.text_encoder.device + prompt_embeds = self.text_encoder(text_input_ids.to(model_device), mask.to(model_device)).last_hidden_state prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] prompt_embeds = torch.stack( diff --git a/src/diffusers/pipelines/wan/pipeline_wan_video2video.py b/src/diffusers/pipelines/wan/pipeline_wan_video2video.py index 7780fc712227..f64f334c039e 100644 --- a/src/diffusers/pipelines/wan/pipeline_wan_video2video.py +++ b/src/diffusers/pipelines/wan/pipeline_wan_video2video.py @@ -246,7 +246,8 @@ def _get_t5_prompt_embeds( text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask seq_lens = mask.gt(0).sum(dim=1).long() - prompt_embeds = self.text_encoder(text_input_ids.to(device), mask.to(device)).last_hidden_state + model_device = self.text_encoder.device + prompt_embeds = self.text_encoder(text_input_ids.to(model_device), mask.to(model_device)).last_hidden_state prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] prompt_embeds = torch.stack( diff --git a/src/diffusers/utils/__init__.py b/src/diffusers/utils/__init__.py index 5c63a4bc7661..0b1362f769d4 100644 --- a/src/diffusers/utils/__init__.py +++ b/src/diffusers/utils/__init__.py @@ -115,6 +115,7 @@ is_torch_mlu_available, is_torch_neuronx_available, is_torch_npu_available, + is_torch_tpu_available, is_torch_version, is_torch_xla_available, is_torch_xla_version, diff --git a/src/diffusers/utils/import_utils.py b/src/diffusers/utils/import_utils.py index d2cf394cd9a7..e9c8a9e1858d 100644 --- a/src/diffusers/utils/import_utils.py +++ b/src/diffusers/utils/import_utils.py @@ -178,6 +178,7 @@ def _is_package_available(pkg_name: str, get_dist_name: bool = False) -> tuple[b _torch_xla_available, _torch_xla_version = _is_package_available("torch_xla") _torch_npu_available, _torch_npu_version = _is_package_available("torch_npu") _torch_mlu_available, _torch_mlu_version = _is_package_available("torch_mlu") +_torch_tpu_available, _torch_tpu_version = _is_package_available("torch_tpu") _torch_neuronx_available, _torch_neuronx_version = _is_package_available("torch_neuronx") _transformers_available, _transformers_version = _is_package_available("transformers") _hf_hub_available, _hf_hub_version = _is_package_available("huggingface_hub") @@ -238,6 +239,10 @@ def is_torch_mlu_available(): return _torch_mlu_available +def is_torch_tpu_available(): + return _torch_tpu_available + + def is_torch_neuronx_available(): return _torch_neuronx_available @@ -553,6 +558,11 @@ def is_av_available(): torchao` """ +TORCH_TPU_IMPORT_ERROR = """ +{0} requires the torch_tpu library but it was not found in your environment. Please follow the installation +instructions at https://github.com/pytorch/tpu +""" + QUANTO_IMPORT_ERROR = """ {0} requires the optimum-quanto library but it was not found in your environment. You can install it with pip: `pip install optimum-quanto` @@ -613,6 +623,7 @@ def is_av_available(): ("pytorch_retinaface", (is_pytorch_retinaface_available, PYTORCH_RETINAFACE_IMPORT_ERROR)), ("better_profanity", (is_better_profanity_available, BETTER_PROFANITY_IMPORT_ERROR)), ("nltk", (is_nltk_available, NLTK_IMPORT_ERROR)), + ("torch_tpu", (is_torch_tpu_available, TORCH_TPU_IMPORT_ERROR)), ("torch_neuronx", (is_torch_neuronx_available, TORCH_NEURONX_IMPORT_ERROR)), ] ) diff --git a/src/diffusers/utils/torch_utils.py b/src/diffusers/utils/torch_utils.py index 9f0877eb1d13..00720a705648 100644 --- a/src/diffusers/utils/torch_utils.py +++ b/src/diffusers/utils/torch_utils.py @@ -45,6 +45,7 @@ "cpu": True, "mps": False, "neuron": False, + "tpu": False, "default": True, } BACKEND_EMPTY_CACHE = { @@ -52,6 +53,7 @@ "xpu": torch.xpu.empty_cache, "cpu": None, "mps": torch.mps.empty_cache, + "tpu": getattr(getattr(torch, "tpu", None), "empty_cache", None), "neuron": None, "default": None, } @@ -60,6 +62,7 @@ "xpu": torch.xpu.device_count, "cpu": lambda: 0, "mps": lambda: 0, + "tpu": lambda: getattr(getattr(torch, "tpu", None), "device_count", lambda: 0)(), "neuron": lambda: getattr(getattr(torch, "neuron", None), "device_count", lambda: 0)(), "default": 0, } @@ -68,6 +71,9 @@ "xpu": torch.xpu.manual_seed, "cpu": torch.manual_seed, "mps": torch.mps.manual_seed, + # TPU latents are always generated on CPU (TPU RNG has unaligned DUS bug), + # so CPU seeding is the correct behaviour here. + "tpu": torch.manual_seed, "neuron": torch.manual_seed, "default": torch.manual_seed, } @@ -76,6 +82,7 @@ "xpu": getattr(torch.xpu, "reset_peak_memory_stats", None), "cpu": None, "mps": None, + "tpu": None, "neuron": None, "default": None, } @@ -84,6 +91,7 @@ "xpu": getattr(torch.xpu, "reset_peak_memory_stats", None), "cpu": None, "mps": None, + "tpu": None, "neuron": None, "default": None, } @@ -92,6 +100,7 @@ "xpu": getattr(torch.xpu, "max_memory_allocated", None), "cpu": 0, "mps": 0, + "tpu": 0, "neuron": 0, "default": 0, } @@ -100,6 +109,7 @@ "xpu": getattr(torch.xpu, "synchronize", None), "cpu": None, "mps": None, + "tpu": getattr(getattr(torch, "tpu", None), "synchronize", None), "neuron": getattr(getattr(torch, "neuron", None), "synchronize", None), "default": None, } @@ -197,6 +207,11 @@ def randn_tensor( rand_device = device batch_size = shape[0] + # TPU RNG has an unaligned DUS (dynamic-update-slice) bug — generate on CPU + # and move to TPU via the existing .to(device) call at the end. + if device is not None and device.type == "tpu": + rand_device = torch.device("cpu") + layout = layout or torch.strided device = device or torch.device("cpu") diff --git a/tests/models/transformers/_neuron_tp_worker.py b/tests/models/transformers/_neuron_tp_worker.py index 681f5686bee9..1306335e2928 100644 --- a/tests/models/transformers/_neuron_tp_worker.py +++ b/tests/models/transformers/_neuron_tp_worker.py @@ -99,7 +99,6 @@ def main(): print("[rank0] PASS: Neuron tensor-parallel output matches single-device reference.") dist.barrier() - dist.destroy_process_group() if __name__ == "__main__": diff --git a/tests/models/transformers/_tp_worker_common.py b/tests/models/transformers/_tp_worker_common.py new file mode 100644 index 000000000000..750da2a94e7c --- /dev/null +++ b/tests/models/transformers/_tp_worker_common.py @@ -0,0 +1,109 @@ +# 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. + +"""Shared logic for the per-backend TP-correctness `torchrun` workers (`_tpu_tp_worker.py`, `_neuron_tp_worker.py`). + +Both workers follow the same recipe — build an identical model on every rank, compute a single-device reference, +shard with `enable_parallelism`, run on-device, and compare — differing only in backend-specific details (device +string, sync call, whether the reference itself needs to run on-device, and numerical tolerance). This module holds +that shared recipe; each `__tp_worker.py` is a thin wrapper supplying those details. +""" + +import copy +import importlib +from typing import Callable + +import torch +import torch.distributed as dist +from torch.distributed.device_mesh import DeviceMesh + +from diffusers import TensorParallelConfig + + +def run_tp_correctness_worker( + spec: str, + *, + mesh_device_type: str, + to_device, + backend_label: str, + synchronize: Callable[[], None], + reference_on_device: bool, + atol: float, + rtol: float, +) -> None: + """Assert a model's tensor-parallel output matches its single-device reference, on the given backend. + + Args: + spec: `module:function` reference returning `(model_class, init_dict, cpu_inputs)` for the model under test. + mesh_device_type: The `DeviceMesh` device type (e.g. `"tpu"`, `"neuron"`). + to_device: The value passed to `.to(...)` to move the model/inputs onto the accelerator. Usually the same as + `mesh_device_type`, but some backends (e.g. Neuron) need a more specific device handle here. + backend_label: Human-readable backend name for log messages (e.g. `"TPU"`, `"Neuron"`). + synchronize: Callable that blocks until pending device work completes. + reference_on_device: If `True`, the unsharded reference forward pass also runs on the accelerator (before TP + mutates the weights in place), so it uses the same kernels as the TP forward pass and only sharding + differs. If `False`, the reference runs on CPU. + atol: Absolute tolerance for the final `torch.testing.assert_close` comparison. + rtol: Relative tolerance for the final `torch.testing.assert_close` comparison. + """ + module_name, _, fn_name = spec.partition(":") + model_class, init_dict, inputs = getattr(importlib.import_module(module_name), fn_name)() + + rank = dist.get_rank() + tp_size = dist.get_world_size() + tp_mesh = DeviceMesh(mesh_device_type, list(range(tp_size))) + + # Identical weights on every rank (same seed), kept on CPU as the pre-shard backends require. + torch.manual_seed(0) + model = model_class(**init_dict).eval() + + if reference_on_device: + # Single-device (unsharded) reference on the accelerator, computed before TP mutates the weights in place. + ref_model = copy.deepcopy(model).to(to_device) + synchronize() + inputs_on_device = {k: v.to(to_device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()} + with torch.no_grad(): + ref_output = ref_model(**inputs_on_device, return_dict=False)[0] + synchronize() + ref_output = ref_output.float().cpu() + del ref_model + else: + with torch.no_grad(): + ref_output = model(**inputs, return_dict=False)[0].float().cpu() + inputs_on_device = {k: v.to(to_device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()} + + # Shard across all ranks; the backend is auto-selected from the mesh device type. + model.enable_parallelism(config=TensorParallelConfig(mesh=tp_mesh)) + model = model.to(to_device) + synchronize() + + with torch.no_grad(): + tp_output = model(**inputs_on_device, return_dict=False)[0] + synchronize() + tp_output = tp_output.float().cpu() + + if rank == 0: + assert tp_output.shape == ref_output.shape, f"shape mismatch: {tp_output.shape} vs {ref_output.shape}" + assert torch.isfinite(tp_output).all(), "TP output contains non-finite values" + max_abs = (tp_output - ref_output).abs().max().item() + denom = ref_output.abs().max().item() + 1e-6 + print( + f"[rank0] tp_size={tp_size} output_shape={tuple(tp_output.shape)} " + f"max_abs_diff={max_abs:.4e} max_rel_diff={max_abs / denom:.4e}" + ) + torch.testing.assert_close(tp_output, ref_output, atol=atol, rtol=rtol) + print(f"[rank0] PASS: {backend_label} tensor-parallel output matches single-device reference.") + + dist.barrier() diff --git a/tests/models/transformers/_tp_worker_launch.py b/tests/models/transformers/_tp_worker_launch.py new file mode 100644 index 000000000000..20f989964815 --- /dev/null +++ b/tests/models/transformers/_tp_worker_launch.py @@ -0,0 +1,113 @@ +# 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. + +"""Shared `pytest`-side launcher for the per-backend TP-correctness `torchrun` workers. + +Used by `TensorParallelTPUTesterMixin` subclasses (and any future accelerator's TP test) to launch their +`__tp_worker.py` under `torchrun` and assert it exits cleanly. +""" + +import os +import subprocess +import sys + +import pytest + + +def run_tp_worker_subprocess(worker_filename: str, spec: str, world_size: int, timeout_s: int = 900) -> None: + """Launch a `torchrun` TP-correctness worker subprocess and assert it exits cleanly. + + Args: + worker_filename: Name of the worker script, resolved relative to this file's directory (e.g. + `"_tpu_tp_worker.py"`). + spec: `module:function` reference forwarded to the worker, see `_tp_worker_common.run_tp_correctness_worker`. + world_size: Number of ranks to launch (`torchrun --nproc_per_node`). + timeout_s: Seconds to wait for the subprocess before failing the test. The worker itself only needs a couple + of minutes even from a cold compile; this generously bounds it so a real hang (e.g. a distributed-runtime + barrier timeout) fails the test loudly instead of stalling the run. + """ + worker = os.path.join(os.path.dirname(__file__), worker_filename) + cmd = [sys.executable, "-m", "torch.distributed.run", f"--nproc_per_node={world_size}", worker, spec] + try: + result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout_s) + except subprocess.TimeoutExpired as e: + raise AssertionError( + f"TP worker did not finish within {timeout_s}s (likely stuck on a distributed-runtime barrier).\n" + f"--- stdout ---\n{e.stdout}\n--- stderr ---\n{e.stderr}" + ) from e + assert result.returncode == 0, ( + f"TP worker failed (exit {result.returncode}).\n--- stdout ---\n{result.stdout}\n--- stderr ---\n{result.stderr}" + ) + + +class TensorParallelTPUTesterMixin: + """Mixin for a `@require_torch_tpu` tensor-parallel correctness test, run via `_tpu_tp_worker.py`. + + Subclasses set `TP_SPEC` to a `module:function` reference (see + `_tp_worker_common.run_tp_correctness_worker`'s `spec` argument) and, only if the model spec's head count + doesn't divide 4, override `WORLD_SIZE`. + + `WORLD_SIZE` defaults to 4 rather than an arbitrary rank count: `torch_tpu`'s per-generation topology table + (`torch_tpu._internal.utils.hardware`) only enumerates whole-pod-slice chip counts (1/4/8 for v6e, for example), + not arbitrary sub-slices of a larger single host. A rank count with no matching whole-slice topology has + nothing to advertise and the PJRT client never completes its start-session barrier — the test would hang for + the barrier's full multi-minute timeout instead of failing. 4 is the smallest whole-slice count every current + TPU generation defines (see `_V4_TOPOLOGY` / `_V5E_TOPOLOGY` / `_V6E_TOPOLOGY` / `_V7_TOPOLOGY` in + `torch_tpu._internal.utils.hardware`). `skip_if_unsupported` below still checks the actual host up front and + skips fast instead of hanging when it doesn't have exactly that many chips. + + Requires `TORCH_TPU_TOPOLOGY` and `TORCH_TPU_SLICEBUILDER_ADDRESSES` to be set. Source them via:: + + eval $(python -m torch_tpu._internal.distributed.launchers.singlehost_wrapper | sed 's/^/export /') + """ + + WORLD_SIZE = 4 + # The worker itself only needs a couple of minutes even from a cold XLA compile; this generously bounds the + # subprocess so a real hang (e.g. a barrier timeout this skip failed to catch) fails the test loudly instead of + # stalling the run. + TIMEOUT_S = 900 + TP_SPEC: str = "" + + def skip_if_unsupported(self): + """Skip unless the host has exactly `WORLD_SIZE` TPU chips. + + A topology *string* existing for a chip count (`hardware.get_tpu_topology`) isn't enough to guarantee the + PJRT client can actually form that session: a sub-slice of a larger single host (e.g. claiming 2 of a + 4-chip v6e-4's chips via `TORCH_TPU_TOPOLOGY`/`TORCH_TPU_SLICEBUILDER_ADDRESSES`) can still fail with a + low-level `START_SESSION` GRPC error, since the runtime's session setup is tied to the host's actual + provisioned slice, not just a topology label. The only combination verified to work is running with exactly + as many ranks as the host has chips. + """ + from torch_tpu._internal.utils import hardware + + try: + device_count = hardware.get_tpu_device_count() + except Exception as e: # pragma: no cover - defensive, hardware detection is best-effort + pytest.skip(f"Could not determine local TPU chip count: {e}") + return + + if device_count != self.WORLD_SIZE: + pytest.skip( + f"This host exposes {device_count} TPU chip(s), but this test requires exactly " + f"{self.WORLD_SIZE} (a TPU single-host tensor-parallel job must use all chips on the host; " + f"sub-slicing a larger host is not reliably supported by the runtime). Run this test on a host " + f"with exactly {self.WORLD_SIZE} TPU chips." + ) + + def test_tensor_parallel_tpu_inference(self): + self.skip_if_unsupported() + run_tp_worker_subprocess( + "_tpu_tp_worker.py", self.TP_SPEC, world_size=self.WORLD_SIZE, timeout_s=self.TIMEOUT_S + ) diff --git a/tests/models/transformers/_tpu_tp_worker.py b/tests/models/transformers/_tpu_tp_worker.py new file mode 100644 index 000000000000..e219fb1a6f45 --- /dev/null +++ b/tests/models/transformers/_tpu_tp_worker.py @@ -0,0 +1,80 @@ +# 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. + +"""TPU entry point for the generic TP-correctness worker (see `_tp_worker_common.py`). + +Model-agnostic. The model under test is supplied as a ``module:function`` spec reference on the command line; the +referenced factory returns ``(model_class, init_dict, inputs)`` with CPU tensors, so all model-specific test data lives +with the launching test rather than here. + +Launched as a subprocess by a ``@require_torch_tpu`` test (and runnable directly for debugging):: + + eval $(python -m torch_tpu._internal.distributed.launchers.singlehost_wrapper | sed 's/^/export /') + torchrun --nproc_per_node=4 _tpu_tp_worker.py \\ + tests.models.transformers.test_models_transformer_flux2:make_tpu_tp_spec + +Exit code 0 means the TP path is numerically equivalent to the unsharded model; non-zero means failure. +""" + +import argparse +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.distributed as dist +import torch_tpu # noqa: F401 — registers "tpu" device and "tpu_dist" backend +from torch_tpu._internal import sync as tpu_sync + +from tests.models.transformers._tp_worker_common import run_tp_correctness_worker + + +def main(): + parser = argparse.ArgumentParser(description="TPU tensor-parallel correctness worker.") + parser.add_argument( + "spec", + help="`module:function` reference returning (model_class, init_dict, tpu_inputs) for the model under test.", + ) + args = parser.parse_args() + + dist.init_process_group(backend="tpu_dist") + # The reference runs on the TPU (not CPU) so both the reference and the TP pass use the same Flash Attention + # kernel; the only difference between them is sharding, not numerical implementation. TPU Flash Attention has + # bf16-level numerics, so the tolerance is wider than fp32 — but a wrong shard plan produces grossly different + # output and is caught comfortably within this bound. + run_tp_correctness_worker( + args.spec, + mesh_device_type="tpu", + to_device="tpu", + backend_label="TPU", + synchronize=lambda: tpu_sync.synchronize(None, wait=True), + reference_on_device=True, + atol=0.1, + rtol=0.1, + ) + 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/run_flux2_tp_tpu.py b/tests/models/transformers/run_flux2_tp_tpu.py new file mode 100644 index 000000000000..08a0420a65e5 --- /dev/null +++ b/tests/models/transformers/run_flux2_tp_tpu.py @@ -0,0 +1,333 @@ +#!/usr/bin/env python3 +"""Verify Flux2 transformer forward pass with tensor parallelism on TPU. + +Without --model-id (default): + Builds a small Flux2 model with random weights, runs a TP forward pass, and compares + it against a single-device TPU reference (same model, same device, no sharding). Both + PASS/FAIL and a max_abs_diff are printed. + +With --model-id (e.g. black-forest-labs/FLUX.2-dev): + Loads the real model from the Hub (CPU), applies TP, runs one forward pass on TPU, and + checks the output is finite and has the expected shape. No reference comparison (too slow). + +The script self-relaunches under torchrun when it is not already a distributed worker, so a +single ``python run_flux2_tp_tpu.py`` invocation is enough. The TPU topology env-vars must be +set before the torchrun relaunch; pass them via --topology / --addresses or export them first: + + eval $(python -m torch_tpu._internal.distributed.launchers.singlehost_wrapper | sed 's/^/export /') + python run_flux2_tp_tpu.py --tp-degree 4 + + # or, to test against real weights: + python run_flux2_tp_tpu.py --tp-degree 4 --model-id black-forest-labs/FLUX.2-dev + +To run end-to-end image generation, use run_flux2_tp_tpu_pipeline.py instead. + +The default sequence lengths (latent_h=16, latent_w=16, txt_len=256) give a joint sequence of +512, which satisfies the TPU Flash Attention requirement of seq_len divisible by 512. +""" + +import argparse +import copy +import os +import sys +import time +import traceback + + +# Make in-repo packages importable when run from an arbitrary CWD via torchrun. +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_tpu # noqa: F401 — registers "tpu" device and "tpu_dist" backend +from torch.distributed.device_mesh import DeviceMesh +from torch_tpu._internal import sync as tpu_sync + +from diffusers import Flux2Transformer2DModel, TensorParallelConfig + + +# ── helpers ────────────────────────────────────────────────────────────────── + + +def log(rank: int, msg: str) -> None: + if rank == 0: + print(f"[flux2-tp-tpu] {msg}", flush=True) + + +def relaunch_via_torchrun(tp_degree: int, topology: str | None, addresses: str | None) -> None: + """Re-invoke this script under torch.distributed.run if not already a worker.""" + if os.environ.get("LOCAL_RANK") is not None: + return # already running inside torchrun — nothing to do + + if tp_degree == 1: + return # single-process, torchrun not needed + + if topology: + os.environ["TORCH_TPU_TOPOLOGY"] = topology + if addresses: + os.environ["TORCH_TPU_SLICEBUILDER_ADDRESSES"] = addresses + + cmd = [ + sys.executable, + "-m", + "torch.distributed.run", + f"--nproc-per-node={tp_degree}", + os.path.abspath(__file__), + ] + sys.argv[1:] # forward all original flags to workers + raise SystemExit(os.execvp(sys.executable, cmd)) + + +def _make_inputs( + in_channels: int, + joint_attention_dim: int, + latent_height: int, + latent_width: int, + txt_len: int, + batch_size: int, + device: str | torch.device, + dtype: torch.dtype, +) -> dict: + """Build a forward-pass input dict from model config dimensions.""" + seq_len = latent_height * latent_width + + hidden_states = torch.randn(batch_size, seq_len, in_channels, dtype=dtype, device=device) + encoder_hidden_states = torch.randn(batch_size, txt_len, joint_attention_dim, dtype=dtype, device=device) + + t_c = torch.arange(1) + h_c = torch.arange(latent_height) + w_c = torch.arange(latent_width) + l_c = torch.arange(1) + img_ids = torch.cartesian_prod(t_c, h_c, w_c, l_c).unsqueeze(0).expand(batch_size, -1, -1).to(device) + + txt_ids = ( + torch.cartesian_prod(torch.arange(1), torch.arange(1), torch.arange(1), torch.arange(txt_len)) + .unsqueeze(0) + .expand(batch_size, -1, -1) + .to(device) + ) + + timestep = torch.tensor([500.0], dtype=dtype, device=device).expand(batch_size) + guidance = torch.tensor([3.5], dtype=dtype, device=device).expand(batch_size) + + return { + "hidden_states": hidden_states, + "encoder_hidden_states": encoder_hidden_states, + "img_ids": img_ids, + "txt_ids": txt_ids, + "timestep": timestep, + "guidance": guidance, + } + + +# ── main worker ─────────────────────────────────────────────────────────────── + + +def run(args: argparse.Namespace) -> int: + dist.init_process_group(backend="tpu_dist") + rank = dist.get_rank() + tp_size = dist.get_world_size() + tp_mesh = DeviceMesh("tpu", list(range(tp_size))) + + log(rank, f"tp_size={tp_size} dtype=bfloat16") + + # ── load model ──────────────────────────────────────────────────────────── + t0 = time.perf_counter() + try: + if args.model_id: + log(rank, f"loading from Hub: {args.model_id}") + model = Flux2Transformer2DModel.from_pretrained( + args.model_id, + subfolder="transformer", + torch_dtype=torch.bfloat16, + ) + else: + log(rank, "building small dummy model (random weights)") + model = Flux2Transformer2DModel( + patch_size=1, + in_channels=4, + num_layers=1, + num_single_layers=1, + attention_head_dim=16, + num_attention_heads=4, # must be divisible by tp_degree + joint_attention_dim=32, + timestep_guidance_channels=256, + axes_dims_rope=[4, 4, 4, 4], + ).to(torch.bfloat16) + except Exception: + log(rank, "LOAD FAILED") + if rank == 0: + traceback.print_exc() + return 1 + log(rank, f"load OK ({time.perf_counter() - t0:.1f}s)") + + cfg = model.config + latent_h = args.latent_height + latent_w = args.latent_width + txt_len = args.txt_len + + # ── single-device reference (dummy mode only, before TP mutates weights) ── + # The reference runs on the TPU device (not CPU) so both the reference and the TP forward + # use the same kernels (e.g. Flash Attention). The only difference between them is sharding. + ref_output = None + if not args.model_id: + ref_model = copy.deepcopy(model).to("tpu") + tpu_sync.synchronize(None, wait=True) + torch.manual_seed(0) + ref_inputs = _make_inputs( + cfg.in_channels, + cfg.joint_attention_dim, + latent_h, + latent_w, + txt_len, + batch_size=1, + device="tpu", + dtype=torch.bfloat16, + ) + ref_model.eval() + with torch.no_grad(): + ref_output_tpu = ref_model(**ref_inputs, return_dict=False)[0] + tpu_sync.synchronize(None, wait=True) + ref_output = ref_output_tpu.float().cpu() + del ref_model, ref_inputs, ref_output_tpu + tpu_sync.synchronize(None, wait=True) + log(rank, f"TPU ref computed max_abs={ref_output.abs().max():.4f}") + + # ── apply tensor parallelism ────────────────────────────────────────────── + try: + model.enable_parallelism(config=TensorParallelConfig(mesh=tp_mesh)) + except Exception: + log(rank, "enable_parallelism FAILED") + if rank == 0: + traceback.print_exc() + return 1 + log(rank, "TP applied") + + model = model.to("tpu") + tpu_sync.synchronize(None, wait=True) + log(rank, "model on TPU, triggering compilation ...") + + # ── build inputs ────────────────────────────────────────────────────────── + torch.manual_seed(0) + tpu_inputs = _make_inputs( + cfg.in_channels, + cfg.joint_attention_dim, + latent_h, + latent_w, + txt_len, + batch_size=1, + device="tpu", + dtype=torch.bfloat16, + ) + + # ── warm-up forward (triggers XLA compilation) ──────────────────────────── + model.eval() + try: + with torch.no_grad(): + _ = model(**tpu_inputs, return_dict=False)[0] + tpu_sync.synchronize(None, wait=True) + log(rank, "warm-up OK (graph compiled)") + except Exception: + log(rank, "warm-up FAILED") + if rank == 0: + traceback.print_exc() + return 1 + + # ── timed forward ───────────────────────────────────────────────────────── + try: + t0 = time.perf_counter() + with torch.no_grad(): + tp_out = model(**tpu_inputs, return_dict=False)[0] + tpu_sync.synchronize(None, wait=True) + elapsed = time.perf_counter() - t0 + except Exception: + log(rank, "timed forward FAILED") + if rank == 0: + traceback.print_exc() + return 1 + + tp_out_cpu = tp_out.float().cpu() + + # ── verify ──────────────────────────────────────────────────────────────── + expected_shape = (1, latent_h * latent_w, cfg.in_channels) + if tp_out_cpu.shape != torch.Size(expected_shape): + log(rank, f"FAIL: shape {tuple(tp_out_cpu.shape)} != expected {expected_shape}") + return 1 + + if not torch.isfinite(tp_out_cpu).all(): + log(rank, "FAIL: output contains non-finite values") + return 1 + + if rank == 0: + stats = f"shape={tuple(tp_out_cpu.shape)} max_abs={tp_out_cpu.abs().max():.4f} time={elapsed * 1000:.1f}ms" + + if ref_output is not None: + # dummy-model mode: compare TP output against single-device TPU reference + max_abs_diff = (tp_out_cpu - ref_output).abs().max().item() + denom = ref_output.abs().max().item() + 1e-6 + max_rel_diff = max_abs_diff / denom + stats += f" max_abs_diff={max_abs_diff:.4e} max_rel_diff={max_rel_diff:.4e}" + + # TPU Flash Attention introduces bf16-level rounding vs standard SDPA. + # A wrong shard plan produces grossly different output (off by ~10x) and is caught + # easily within this tolerance; the bound is wider than the MLP case because + # Flash Attention rounds differently than sequential matmul+softmax. + if max_abs_diff > 0.1: + log(rank, f"FAIL: max_abs_diff={max_abs_diff:.4e} exceeds tolerance 0.1") + return 1 + + log(rank, f"PASS {stats}") + + dist.barrier() + dist.destroy_process_group() + return 0 + + +# ── entry point ─────────────────────────────────────────────────────────────── + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description="Flux2 TP-on-TPU verification script.") + p.add_argument("--tp-degree", type=int, default=4, help="number of TPU chips to shard across") + p.add_argument("--model-id", type=str, default="", help="HuggingFace model ID (empty = random weights)") + p.add_argument( + "--latent-height", + type=int, + default=16, + help="latent grid height (default 16; combined with txt-len=256 gives joint-seq=512)", + ) + p.add_argument( + "--latent-width", + type=int, + default=16, + help="latent grid width (default 16; combined with txt-len=256 gives joint-seq=512)", + ) + p.add_argument( + "--txt-len", + type=int, + default=256, + help="text sequence length (default 256; combined with 16x16 image gives joint-seq=512)", + ) + p.add_argument("--topology", type=str, default="", help="TORCH_TPU_TOPOLOGY (e.g. '2,2,1')") + p.add_argument("--addresses", type=str, default="", help="TORCH_TPU_SLICEBUILDER_ADDRESSES") + return p.parse_args() + + +def main() -> None: + args = parse_args() + relaunch_via_torchrun( + args.tp_degree, + args.topology or None, + args.addresses or None, + ) + raise SystemExit(run(args)) + + +if __name__ == "__main__": + try: + main() + except SystemExit: + raise + except Exception: + traceback.print_exc() + sys.exit(1) diff --git a/tests/models/transformers/test_models_transformer_flux.py b/tests/models/transformers/test_models_transformer_flux.py index 53af9eedc50c..714240634b1f 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_tensor_parallel, + require_torch_neuron, + require_torch_tpu, + torch_device, +) from ..testing_utils import ( AttentionBackendTesterMixin, AttentionTesterMixin, @@ -58,6 +64,7 @@ TorchCompileTesterMixin, TrainingTesterMixin, ) +from ._tp_worker_launch import TensorParallelTPUTesterMixin enable_full_determinism() @@ -268,6 +275,37 @@ class TestFluxTransformerTensorParallel(FluxTransformerTesterConfig, TensorParal """Tensor Parallel inference tests for Flux Transformer (CUDA/XPU multi-accelerator).""" +def make_tpu_tp_spec(): + """Model spec consumed by the generic TPU TP worker (``_tpu_tp_worker.py``). + + Returns ``(model_class, init_dict, cpu_inputs)``. Defined here so all Flux-specific test data lives in this file + while the worker stays model-agnostic. + + Overrides ``num_attention_heads`` to 4 (instead of reusing the shared tester config's 2) so + ``TensorParallelTPUTesterMixin``'s default 4-rank ``WORLD_SIZE`` divides the head count — see + ``make_tpu_tp_spec`` in ``test_models_transformer_flux2.py`` for the full rationale (TPU can't shard across an + arbitrary rank count the way CUDA/XPU can). Every other field still comes from the shared config so the rest of + the spec doesn't drift from the other Flux tests. + """ + config = FluxTransformerTesterConfig() + init_dict = {**config.get_init_dict(), "num_attention_heads": 4} + return FluxTransformer2DModel, init_dict, config.get_dummy_inputs(device="cpu") + + +@is_tensor_parallel +@require_torch_tpu +class TestFluxTransformerTensorParallelTPU(TensorParallelTPUTesterMixin): + """Tensor Parallel inference test for Flux Transformer on TPU. + + TPU TP runs through ``torchrun`` with the ``"tpu_dist"`` distributed backend, so it cannot use the + ``torch.multiprocessing``/NCCL spawn path of ``TensorParallelTesterMixin``. This launches the generic worker + with the Flux model spec (``make_tpu_tp_spec``) via ``TensorParallelTPUTesterMixin``; the worker asserts the + sharded output matches a single-device reference, and the test checks its exit code. + """ + + TP_SPEC = "tests.models.transformers.test_models_transformer_flux:make_tpu_tp_spec" + + def make_neuron_tp_spec(): """Model spec consumed by the generic Neuron TP worker (`_neuron_tp_worker.py`). diff --git a/tests/models/transformers/test_models_transformer_flux2.py b/tests/models/transformers/test_models_transformer_flux2.py index 3263ce68202c..48e7bd605ee7 100644 --- a/tests/models/transformers/test_models_transformer_flux2.py +++ b/tests/models/transformers/test_models_transformer_flux2.py @@ -28,7 +28,13 @@ ) 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_tensor_parallel, + require_torch_neuron, + require_torch_tpu, + torch_device, +) from ..testing_utils import ( AttentionTesterMixin, BaseModelTesterConfig, @@ -47,6 +53,7 @@ TorchCompileTesterMixin, TrainingTesterMixin, ) +from ._tp_worker_launch import TensorParallelTPUTesterMixin enable_full_determinism() @@ -176,6 +183,42 @@ def make_neuron_tp_spec(): return Flux2Transformer2DModel, config.get_init_dict(), config.get_dummy_inputs(device="cpu") +def make_tpu_tp_spec(): + """Model spec consumed by the generic TPU TP worker (``_tpu_tp_worker.py``). + + Returns ``(model_class, init_dict, cpu_inputs)``. Defined here so all Flux2-specific test data lives in this file + while the worker stays model-agnostic. + + Overrides ``num_attention_heads`` to 4 (instead of reusing the shared tester config's 2) so the TP degree divides + the head count on a whole-pod-slice TPU host: `torch_tpu`'s per-generation topology table + (``torch_tpu._internal.utils.hardware``) only enumerates whole-slice chip counts (1/4/8 for v6e, for example), not + arbitrary sub-slices of a larger single host, and ``TestFlux2TransformerTensorParallelTPU`` shards across + ``WORLD_SIZE`` ranks to match. Every other field still comes from the shared config so the rest of the spec + doesn't drift from the other Flux2 tests. + """ + config = Flux2TransformerTesterConfig() + init_dict = {**config.get_init_dict(), "num_attention_heads": 4} + return Flux2Transformer2DModel, init_dict, config.get_dummy_inputs(device="cpu") + + +@is_tensor_parallel +@require_torch_tpu +class TestFlux2TransformerTensorParallelTPU(TensorParallelTPUTesterMixin): + """Tensor Parallel inference test for Flux2 Transformer on TPU. + + TPU TP runs through ``torchrun`` with the ``"tpu_dist"`` distributed backend, so it cannot use the + ``torch.multiprocessing``/NCCL spawn path of ``TensorParallelTesterMixin``. This launches the generic worker with + the Flux2 model spec (``make_tpu_tp_spec``) via ``TensorParallelTPUTesterMixin``; the worker asserts the sharded + output matches a single-device reference, and the test checks its exit code. + + ``make_tpu_tp_spec`` overrides ``num_attention_heads`` to 4 so that ``TensorParallelTPUTesterMixin``'s default + 4-rank ``WORLD_SIZE`` divides the head count — unlike the CUDA/XPU ``TensorParallelTesterMixin``, which hardcodes + ``world_size = 2`` to match `Flux2TransformerTesterConfig`'s 2 heads. + """ + + TP_SPEC = "tests.models.transformers.test_models_transformer_flux2:make_tpu_tp_spec" + + @is_tensor_parallel @require_torch_neuron class TestFlux2TransformerTensorParallelNeuron: diff --git a/tests/models/transformers/test_models_transformer_qwenimage.py b/tests/models/transformers/test_models_transformer_qwenimage.py index 7a03a8fe2353..61ec6d493409 100644 --- a/tests/models/transformers/test_models_transformer_qwenimage.py +++ b/tests/models/transformers/test_models_transformer_qwenimage.py @@ -24,7 +24,13 @@ from diffusers.models.transformers.transformer_qwenimage import compute_text_seq_len_from_mask from diffusers.utils.torch_utils import randn_tensor -from ...testing_utils import enable_full_determinism, is_tensor_parallel, require_torch_neuron, torch_device +from ...testing_utils import ( + enable_full_determinism, + is_tensor_parallel, + require_torch_neuron, + require_torch_tpu, + torch_device, +) from ..testing_utils import ( AttentionBackendTesterMixin, AttentionTesterMixin, @@ -41,6 +47,7 @@ TorchCompileTesterMixin, TrainingTesterMixin, ) +from ._tp_worker_launch import TensorParallelTPUTesterMixin enable_full_determinism() @@ -307,6 +314,32 @@ class TestQwenImageTransformerTensorParallel(QwenImageTransformerTesterConfig, T """Tensor Parallel inference tests for QwenImage Transformer (CUDA/XPU multi-accelerator).""" +def make_tpu_tp_spec(): + """Model spec consumed by the generic TPU TP worker (``_tpu_tp_worker.py``). + + Returns ``(model_class, init_dict, cpu_inputs)``. Defined here so all QwenImage-specific test data lives in this + file while the worker stays model-agnostic. ``QwenImageTransformerTesterConfig``'s default ``num_attention_heads`` + (4) already divides ``TensorParallelTPUTesterMixin``'s default 4-rank ``WORLD_SIZE``, so no override is needed + here (contrast Flux/Flux2, whose shared config defaults to 2 heads and does need one). + """ + config = QwenImageTransformerTesterConfig() + return QwenImageTransformer2DModel, config.get_init_dict(), config.get_dummy_inputs(device="cpu") + + +@is_tensor_parallel +@require_torch_tpu +class TestQwenImageTransformerTensorParallelTPU(TensorParallelTPUTesterMixin): + """Tensor Parallel inference test for QwenImage Transformer on TPU. + + TPU TP runs through ``torchrun`` with the ``"tpu_dist"`` distributed backend, so it cannot use the + ``torch.multiprocessing``/NCCL spawn path of ``TensorParallelTesterMixin``. This launches the generic worker + with the QwenImage model spec (``make_tpu_tp_spec``) via ``TensorParallelTPUTesterMixin``; the worker asserts + the sharded output matches a single-device reference, and the test checks its exit code. + """ + + TP_SPEC = "tests.models.transformers.test_models_transformer_qwenimage:make_tpu_tp_spec" + + def make_neuron_tp_spec(): """Model spec consumed by the generic Neuron TP worker (``_neuron_tp_worker.py``). diff --git a/tests/testing_utils.py b/tests/testing_utils.py index 9da89e626198..d9671828314c 100644 --- a/tests/testing_utils.py +++ b/tests/testing_utils.py @@ -46,6 +46,7 @@ is_timm_available, is_torch_available, is_torch_neuronx_available, + is_torch_tpu_available, is_torch_version, is_torchao_available, is_torchsde_available, @@ -565,6 +566,14 @@ def require_torch_neuron(test_case): )(test_case) +def require_torch_tpu(test_case): + """Decorator marking a test that requires a TPU device (torch_tpu).""" + return pytest.mark.skipif( + not is_torch_tpu_available(), + reason="test requires TPU device (torch_tpu)", + )(test_case) + + def require_torch_multi_gpu(test_case): """ Decorator marking a test that requires a multi-GPU setup (in PyTorch). These tests are skipped on a machine without