|
| 1 | +# SPDX-FileCopyrightText: (c) 2026 Tenstorrent AI ULC |
| 2 | +# |
| 3 | +# SPDX-License-Identifier: Apache-2.0 |
| 4 | + |
| 5 | +"""FLUX.2-dev — benchmark-side pipeline for the imagegen harness (multichip). |
| 6 | +
|
| 7 | +FLUX.2 is the first image-gen benchmark that runs tensor-parallel across |
| 8 | +several Tenstorrent devices: the ~24B Mistral3 text encoder and the ~32B Flux2 |
| 9 | +transformer are SPMD-sharded over the mesh's ``model`` axis (always degree 4 — |
| 10 | +the contraction-parallel degree the shard specs are written for), while the VAE |
| 11 | +decoder is replicated. The shared diffusion harness in ``imagegen_benchmark.py`` |
| 12 | +is otherwise unchanged: it builds this pipeline once and runs a warmup pass |
| 13 | +(1 step, compiles every component) followed by a steady-state pass. |
| 14 | +
|
| 15 | +The standard diffusers ``Flux2Pipeline`` orchestrates the run (tokenizer + |
| 16 | +scheduler stay on CPU); every compute module is routed onto TT via |
| 17 | +``torch.compile(backend="tt")`` wrappers, mirroring the nightly e2e test in |
| 18 | +``tests/torch/models/flux2/test_flux2_pipeline.py``. |
| 19 | +
|
| 20 | +Memory strategy (peak ≈ max(component) rather than the sum): the text encoder |
| 21 | +is placed, used and evicted before the transformer is placed, and the VAE is |
| 22 | +placed lazily at first decode. Per-component forward+sync times are collected |
| 23 | +into ``self._perf`` for the harness to read after each ``generate()`` call. |
| 24 | +""" |
| 25 | + |
| 26 | +import os |
| 27 | +import time |
| 28 | +from typing import Optional |
| 29 | + |
| 30 | +import numpy as np |
| 31 | +import torch |
| 32 | +import torch_xla |
| 33 | +import torch_xla.distributed.spmd as xs |
| 34 | +import torch_xla.runtime as xr |
| 35 | +from diffusers import Flux2Pipeline |
| 36 | +from loguru import logger |
| 37 | +from torch_xla.distributed.spmd import Mesh |
| 38 | + |
| 39 | +from third_party.tt_forge_models.flux2.pytorch.src.model_utils import ( |
| 40 | + DTYPE, |
| 41 | + GUIDANCE_SCALE, |
| 42 | + HEIGHT, |
| 43 | + MESH_NAMES, |
| 44 | + MESH_SHAPES, |
| 45 | + PROMPT, |
| 46 | + REPO_ID, |
| 47 | + SEED, |
| 48 | + WIDTH, |
| 49 | + Mistral3TextEncoderWrapper, |
| 50 | + shard_text_encoder_specs, |
| 51 | + shard_transformer_specs, |
| 52 | + tokenize_prompt, |
| 53 | +) |
| 54 | + |
| 55 | + |
| 56 | +class _DeviceDenoiser: |
| 57 | + """Routes Flux2Pipeline's transformer calls to the TP-sharded model on TT. |
| 58 | +
|
| 59 | + Each ``__call__`` is one denoising step; its forward+sync time is appended |
| 60 | + to ``perf["unet_steps"]`` (the harness's generic per-step metric). |
| 61 | + """ |
| 62 | + |
| 63 | + def __init__(self, transformer, mesh, perf): |
| 64 | + self._dev = torch_xla.device() |
| 65 | + self._perf = perf |
| 66 | + self.config = transformer.config |
| 67 | + self.dtype = next(transformer.parameters()).dtype |
| 68 | + |
| 69 | + transformer = transformer.to(self._dev) |
| 70 | + if hasattr(transformer, "tie_weights"): |
| 71 | + transformer.tie_weights() |
| 72 | + for tensor, spec in shard_transformer_specs(transformer).items(): |
| 73 | + xs.mark_sharding(tensor, mesh, spec) |
| 74 | + self._compiled = torch.compile(transformer, backend="tt") |
| 75 | + |
| 76 | + def __call__(self, **kwargs): |
| 77 | + moved = { |
| 78 | + k: (v.to(self._dev) if torch.is_tensor(v) else v) for k, v in kwargs.items() |
| 79 | + } |
| 80 | + t0 = time.perf_counter() |
| 81 | + out = self._compiled(**moved) |
| 82 | + torch_xla.sync() |
| 83 | + # cpu cast forces a sync — the timer ends after the result lands on host. |
| 84 | + if isinstance(out, (tuple, list)): |
| 85 | + result = type(out)(o.cpu() if torch.is_tensor(o) else o for o in out) |
| 86 | + else: |
| 87 | + result = out.cpu() |
| 88 | + self._perf["unet_steps"].append(time.perf_counter() - t0) |
| 89 | + return result |
| 90 | + |
| 91 | + |
| 92 | +class _DeviceVAEDecoder: |
| 93 | + """Routes Flux2Pipeline's vae.decode() to TT (replicated), placed lazily. |
| 94 | +
|
| 95 | + The decode forward+sync time is recorded into ``perf["vae"]`` and the raw |
| 96 | + decoder output (pixels in [-1, 1]) is stashed on ``last_pixels`` so the |
| 97 | + harness can save it without depending on the pipeline's PIL postprocess. |
| 98 | + """ |
| 99 | + |
| 100 | + def __init__(self, vae, mesh, perf): |
| 101 | + self._dev = torch_xla.device() |
| 102 | + self._perf = perf |
| 103 | + self.config = vae.config |
| 104 | + self.dtype = next(vae.parameters()).dtype |
| 105 | + self.bn = vae.bn # stays on CPU; pipeline reads it host-side for denorm |
| 106 | + self._vae = vae |
| 107 | + self._compiled = None |
| 108 | + self.last_pixels = None |
| 109 | + |
| 110 | + def decode(self, latents, return_dict=False): |
| 111 | + # Lazy device placement: keep the VAE off-device during the denoise loop |
| 112 | + # so it does not inflate the denoiser's peak DRAM; place it only now. |
| 113 | + if self._compiled is None: |
| 114 | + vae = self._vae.to(self._dev) |
| 115 | + self._compiled = torch.compile( |
| 116 | + lambda z: vae.decode(z, return_dict=False)[0], backend="tt" |
| 117 | + ) |
| 118 | + t0 = time.perf_counter() |
| 119 | + out = self._compiled(latents.to(self._dev)) |
| 120 | + torch_xla.sync() |
| 121 | + image = out.cpu() if torch.is_tensor(out) else out |
| 122 | + self._perf["vae"] = time.perf_counter() - t0 |
| 123 | + self.last_pixels = image |
| 124 | + return (image,) |
| 125 | + |
| 126 | + |
| 127 | +class Flux2Config: |
| 128 | + def __init__( |
| 129 | + self, |
| 130 | + height: int = HEIGHT, |
| 131 | + width: int = WIDTH, |
| 132 | + compile_options: Optional[dict] = None, |
| 133 | + ): |
| 134 | + self.height = height |
| 135 | + self.width = width |
| 136 | + # Harness-set compile options (forwarded for parity with the other |
| 137 | + # imagegen pipelines; FLUX.2 does not switch options inline). |
| 138 | + self.compile_options = compile_options or {} |
| 139 | + |
| 140 | + |
| 141 | +class Flux2Pipeline_TT: |
| 142 | + """diffusers Flux2Pipeline with every module on TT, tensor-parallel sharded. |
| 143 | +
|
| 144 | + Built once by the harness; ``generate()`` is called twice (warmup + |
| 145 | + steady-state). The raw transformer/VAE modules are kept so the TT wrappers |
| 146 | + can be rebuilt fresh on every call (the diffusers pipe's attributes get |
| 147 | + replaced by wrappers during a run). |
| 148 | + """ |
| 149 | + |
| 150 | + def __init__(self, config: Flux2Config): |
| 151 | + self.config = config |
| 152 | + self._perf = {} |
| 153 | + |
| 154 | + def setup(self): |
| 155 | + # Enable SPMD. CONVERT_SHLO_TO_SHARDY=1 turns on the pytorch/xla |
| 156 | + # ConvertStableHloToSdy pass so the StableHLO handed to tt-mlir carries |
| 157 | + # shardy annotations — without it, presharded args lose their @Sharding |
| 158 | + # custom call and compilation fails. Mirrors infra.enable_spmd(). |
| 159 | + os.environ["CONVERT_SHLO_TO_SHARDY"] = "1" |
| 160 | + xr.use_spmd() |
| 161 | + self.num_devices = xr.global_runtime_device_count() |
| 162 | + self.mesh_shape = MESH_SHAPES[self.num_devices] |
| 163 | + device_ids = np.array(range(self.num_devices)) |
| 164 | + self.mesh = Mesh(device_ids, self.mesh_shape, MESH_NAMES) |
| 165 | + logger.info( |
| 166 | + f"FLUX.2 mesh {self.mesh_shape} (names={MESH_NAMES}) " |
| 167 | + f"over {self.num_devices} devices" |
| 168 | + ) |
| 169 | + self.pipe = Flux2Pipeline.from_pretrained(REPO_ID, torch_dtype=DTYPE) |
| 170 | + # Keep the raw modules so wrappers can be rebuilt on each generate(). |
| 171 | + self._raw_transformer = self.pipe.transformer |
| 172 | + self._raw_vae = self.pipe.vae |
| 173 | + |
| 174 | + def generate( |
| 175 | + self, |
| 176 | + prompt: str, |
| 177 | + num_inference_steps: int, |
| 178 | + seed: Optional[int] = SEED, |
| 179 | + ): |
| 180 | + dev = torch_xla.device() |
| 181 | + self._perf = { |
| 182 | + "te1": None, |
| 183 | + "unet_steps": [], |
| 184 | + "vae": None, |
| 185 | + "total": None, |
| 186 | + "step_label": "Transformer", |
| 187 | + } |
| 188 | + t_total_start = time.perf_counter() |
| 189 | + |
| 190 | + # ── Stage 1: text encoder (sharded, compiled) → prompt embeds, evict ── |
| 191 | + logger.info("[STAGE] Text encoder: start") |
| 192 | + text_encoder = self.pipe.text_encoder |
| 193 | + encoder_wrapper = Mistral3TextEncoderWrapper(text_encoder).eval() |
| 194 | + input_ids, attention_mask = tokenize_prompt(prompt) |
| 195 | + |
| 196 | + text_encoder = text_encoder.to(dev) |
| 197 | + if hasattr(text_encoder, "tie_weights"): |
| 198 | + text_encoder.tie_weights() |
| 199 | + te_specs = shard_text_encoder_specs(text_encoder) |
| 200 | + assert te_specs, "text-encoder shard spec is empty — descent failed (would OOM)" |
| 201 | + for tensor, spec in te_specs.items(): |
| 202 | + xs.mark_sharding(tensor, self.mesh, spec) |
| 203 | + te_compiled = torch.compile(encoder_wrapper, backend="tt") |
| 204 | + |
| 205 | + t0 = time.perf_counter() |
| 206 | + with torch.no_grad(): |
| 207 | + prompt_embeds = te_compiled(input_ids.to(dev), attention_mask.to(dev)) |
| 208 | + torch_xla.sync() |
| 209 | + prompt_embeds = prompt_embeds.cpu() |
| 210 | + self._perf["te1"] = time.perf_counter() - t0 |
| 211 | + |
| 212 | + # Free the 24B encoder from device before placing the 32B denoiser. |
| 213 | + self.pipe.text_encoder = text_encoder.to("cpu") |
| 214 | + del te_compiled, encoder_wrapper |
| 215 | + import gc |
| 216 | + |
| 217 | + gc.collect() |
| 218 | + torch_xla.sync() |
| 219 | + logger.info("[STAGE] Text encoder: done") |
| 220 | + |
| 221 | + # ── Stage 2: denoiser (sharded) + VAE (replicated, lazy) → image ───── |
| 222 | + logger.info("[STAGE] Transformer + VAE: start") |
| 223 | + self.pipe.transformer = _DeviceDenoiser( |
| 224 | + self._raw_transformer, self.mesh, self._perf |
| 225 | + ) |
| 226 | + vae_wrapper = _DeviceVAEDecoder(self._raw_vae, self.mesh, self._perf) |
| 227 | + self.pipe.vae = vae_wrapper |
| 228 | + |
| 229 | + generator = torch.Generator().manual_seed(seed) if seed is not None else None |
| 230 | + self.pipe( |
| 231 | + prompt=None, |
| 232 | + prompt_embeds=prompt_embeds, |
| 233 | + height=self.config.height, |
| 234 | + width=self.config.width, |
| 235 | + num_inference_steps=num_inference_steps, |
| 236 | + guidance_scale=GUIDANCE_SCALE, |
| 237 | + generator=generator, |
| 238 | + ) |
| 239 | + logger.info("[STAGE] Transformer + VAE: done") |
| 240 | + |
| 241 | + self._perf["total"] = time.perf_counter() - t_total_start |
| 242 | + # Raw VAE pixels in [-1, 1], shape (1, 3, H, W) — the harness's |
| 243 | + # save_image() expects this range, so return it instead of the PIL. |
| 244 | + return vae_wrapper.last_pixels |
0 commit comments