Skip to content

Commit 2816222

Browse files
jfacevedo-googlejfacevedo-google
committed
implementing flux on TPUs with ptxla
1 parent a26d570 commit 2816222

File tree

6 files changed

+214
-4
lines changed

6 files changed

+214
-4
lines changed

‎examples/research_projects/pytorch_xla/inference/flux/README.md‎

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
# Generating images using Flux and PyTorch/XLA
2+
3+
The `flux_inference` script shows how to do image generation using Flux on TPU devices using PyTorch/XLA. It uses the pallas kernel for flash attention for faster generation.
4+
5+
It has been tested on [Trillium](https://cloud.google.com/blog/products/compute/introducing-trillium-6th-gen-tpus) TPU versions. No other TPU types have been tested.
6+
7+
## Create TPU
8+
9+
To create a TPU on Google Cloud, follow [this guide](https://cloud.google.com/tpu/docs/v6e)
10+
11+
## Setup TPU environment
12+
13+
SSH into the VM and install Pytorch, Pytorch/XLA
14+
15+
```bash
16+
pip install torch~=2.5.0 torch_xla[tpu]~=2.5.0 -f https://storage.googleapis.com/libtpu-releases/index.html -f https://storage.googleapis.com/libtpu-wheels/index.html
17+
pip install torch_xla[pallas] -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
18+
```
19+
20+
Verify that PyTorch and PyTorch/XLA were installed correctly:
21+
22+
```bash
23+
python3 -c "import torch; import torch_xla;"
24+
```
25+
26+
Install dependencies
27+
28+
```bash
29+
pip install transformers accelerate sentencepiece structlog
30+
pushd ../../..
31+
pip install .
32+
popd
33+
```
34+
35+
## Run the inference job
36+
37+
### Authenticate
38+
39+
Run the following command to authenticate your token in order to download Flux weights.
40+
41+
```bash
42+
huggingface-cli login
43+
```
44+
45+
Then run:
46+
47+
```bash
48+
python flux_inference.py
49+
```
50+
51+
The script loads the text encoders onto the CPU and the Flux transformer and VAE models onto the TPU. The first time the script runs, the compilation time is longer, while the cache stores the compiled programs. On subsequent runs, compilation is much faster and the subsequent passes being the fastest.
52+
53+
On a Trillium v6e-4, you should expect ~9 sec / 4 images or 2.25 sec / image (as devices run generation in parallel):
54+
55+
```bash
56+
WARNING:root:libtpu.so and TPU device found. Setting PJRT_DEVICE=TPU.
57+
Loading checkpoint shards: 100%|███████████████████████████████| 2/2 [00:00<00:00, 7.01it/s]
58+
Loading pipeline components...: 40%|██████████▍ | 2/5 [00:00<00:00, 3.78it/s]You set `add_prefix_space`. The tokenizer needs to be converted from the slow tokenizers
59+
Loading pipeline components...: 100%|██████████████████████████| 5/5 [00:00<00:00, 6.72it/s]
60+
2025-01-10 00:51:25 [info ] loading flux from black-forest-labs/FLUX.1-dev
61+
2025-01-10 00:51:25 [info ] loading flux from black-forest-labs/FLUX.1-dev
62+
2025-01-10 00:51:26 [info ] loading flux from black-forest-labs/FLUX.1-dev
63+
2025-01-10 00:51:26 [info ] loading flux from black-forest-labs/FLUX.1-dev
64+
Loading pipeline components...: 100%|██████████████████████████| 3/3 [00:00<00:00, 4.29it/s]
65+
Loading pipeline components...: 100%|██████████████████████████| 3/3 [00:00<00:00, 3.26it/s]
66+
Loading pipeline components...: 100%|██████████████████████████| 3/3 [00:00<00:00, 3.27it/s]
67+
Loading pipeline components...: 100%|██████████████████████████| 3/3 [00:00<00:00, 3.25it/s]
68+
2025-01-10 00:51:34 [info ] starting compilation run...
69+
2025-01-10 00:51:35 [info ] starting compilation run...
70+
2025-01-10 00:51:37 [info ] starting compilation run...
71+
2025-01-10 00:51:37 [info ] starting compilation run...
72+
2025-01-10 00:52:52 [info ] compilation took 78.5155531649998 sec.
73+
2025-01-10 00:52:53 [info ] starting inference run...
74+
2025-01-10 00:52:57 [info ] compilation took 79.52986721400157 sec.
75+
2025-01-10 00:52:57 [info ] compilation took 81.91776501700042 sec.
76+
2025-01-10 00:52:57 [info ] compilation took 80.24951512600092 sec.
77+
2025-01-10 00:52:57 [info ] starting inference run...
78+
2025-01-10 00:52:57 [info ] starting inference run...
79+
2025-01-10 00:52:58 [info ] starting inference run...
80+
2025-01-10 00:53:22 [info ] inference time: 25.112665320000815
81+
2025-01-10 00:53:30 [info ] inference time: 7.7019307739992655
82+
2025-01-10 00:53:38 [info ] inference time: 7.693858365000779
83+
2025-01-10 00:53:46 [info ] inference time: 7.690621814001133
84+
2025-01-10 00:53:53 [info ] inference time: 7.679490454000188
85+
2025-01-10 00:54:01 [info ] inference time: 7.68949568500102
86+
2025-01-10 00:54:09 [info ] inference time: 7.686633744000574
87+
2025-01-10 00:54:16 [info ] inference time: 7.696786873999372
88+
2025-01-10 00:54:24 [info ] inference time: 7.691988694999964
89+
2025-01-10 00:54:32 [info ] inference time: 7.700649563999832
90+
2025-01-10 00:54:39 [info ] inference time: 7.684993574001055
91+
2025-01-10 00:54:47 [info ] inference time: 7.68343457499941
92+
2025-01-10 00:54:55 [info ] inference time: 7.667921153999487
93+
2025-01-10 00:55:02 [info ] inference time: 7.683585194001353
94+
2025-01-10 00:55:06 [info ] avg. inference over 15 iterations took 8.61202360273334 sec.
95+
2025-01-10 00:55:07 [info ] avg. inference over 15 iterations took 8.952725123600006 sec.
96+
2025-01-10 00:55:10 [info ] inference time: 7.673799695001435
97+
2025-01-10 00:55:10 [info ] avg. inference over 15 iterations took 8.849190365400379 sec.
98+
2025-01-10 00:55:10 [info ] saved metric information as /tmp/metrics_report.txt
99+
2025-01-10 00:55:12 [info ] avg. inference over 15 iterations took 8.940161458400205 sec.
100+
```

‎examples/research_projects/pytorch_xla/inference/flux/flux_inference.py‎

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
from time import perf_counter
2+
from pathlib import Path
3+
from argparse import ArgumentParser
4+
5+
import structlog
6+
7+
import torch
8+
import torch_xla.core.xla_model as xm
9+
import torch_xla.runtime as xr
10+
import torch_xla.debug.profiler as xp
11+
import torch_xla.debug.metrics as met
12+
from diffusers import FluxPipeline
13+
import torch_xla.distributed.xla_multiprocessing as xmp
14+
15+
logger = structlog.get_logger()
16+
metrics_filepath = '/tmp/metrics_report.txt'
17+
18+
def _main(index, args, text_pipe, ckpt_id):
19+
20+
cache_path = Path('/tmp/data/compiler_cache_tRiLlium_eXp')
21+
cache_path.mkdir(parents=True, exist_ok=True)
22+
xr.initialize_cache(str(cache_path), readonly=False)
23+
24+
profile_path = Path('/tmp/data/profiler_out_tRiLlium_eXp')
25+
profile_path.mkdir(parents=True, exist_ok=True)
26+
profiler_port = 9012
27+
profile_duration = args.profile_duration
28+
if args.profile:
29+
logger.info(f'starting profiler on port {profiler_port}')
30+
_ = xp.start_server(profiler_port)
31+
device0 = xm.xla_device()
32+
33+
logger.info(f'loading flux from {ckpt_id}')
34+
flux_pipe = FluxPipeline.from_pretrained(ckpt_id, text_encoder=None, tokenizer=None,
35+
text_encoder_2=None, tokenizer_2=None, torch_dtype=torch.bfloat16).to(device0)
36+
37+
prompt = 'photograph of an electronics chip in the shape of a race car with trillium written on its side'
38+
width = args.width
39+
height = args.height
40+
guidance = args.guidance
41+
n_steps = 4 if args.schnell else 28
42+
43+
logger.info('starting compilation run...')
44+
ts = perf_counter()
45+
with torch.no_grad():
46+
prompt_embeds, pooled_prompt_embeds, text_ids = text_pipe.encode_prompt(
47+
prompt=prompt, prompt_2=None, max_sequence_length=512)
48+
prompt_embeds = prompt_embeds.to(device0)
49+
pooled_prompt_embeds = pooled_prompt_embeds.to(device0)
50+
51+
image = flux_pipe(prompt_embeds=prompt_embeds, pooled_prompt_embeds=pooled_prompt_embeds,
52+
num_inference_steps=28, guidance_scale=guidance, height=height, width=width).images[0]
53+
logger.info(f'compilation took {perf_counter() - ts} sec.')
54+
image.save('/tmp/compile_out.png')
55+
56+
base_seed = 4096 if args.seed is None else args.seed
57+
seed_range = 1000
58+
unique_seed = base_seed + index * seed_range
59+
xm.set_rng_state(seed=unique_seed, device=device0)
60+
times = []
61+
logger.info('starting inference run...')
62+
for _ in range(args.itters):
63+
ts = perf_counter()
64+
with torch.no_grad():
65+
prompt_embeds, pooled_prompt_embeds, text_ids = text_pipe.encode_prompt(
66+
prompt=prompt, prompt_2=None, max_sequence_length=512)
67+
prompt_embeds = prompt_embeds.to(device0)
68+
pooled_prompt_embeds = pooled_prompt_embeds.to(device0)
69+
70+
if args.profile:
71+
xp.trace_detached(f"localhost:{profiler_port}", str(profile_path), duration_ms=profile_duration)
72+
image = flux_pipe(prompt_embeds=prompt_embeds, pooled_prompt_embeds=pooled_prompt_embeds,
73+
num_inference_steps=n_steps, guidance_scale=guidance, height=height, width=width).images[0]
74+
inference_time = perf_counter() - ts
75+
if index == 0:
76+
logger.info(f"inference time: {inference_time}")
77+
times.append(inference_time)
78+
logger.info(f'avg. inference over {args.itters} iterations took {sum(times)/len(times)} sec.')
79+
image.save(f'/tmp/inference_out-{index}.png')
80+
if index == 0:
81+
metrics_report = met.metrics_report()
82+
with open(metrics_filepath, 'w+') as fout:
83+
fout.write(metrics_report)
84+
logger.info(f'saved metric information as {metrics_filepath}')
85+
86+
if __name__ == '__main__':
87+
parser = ArgumentParser()
88+
parser.add_argument('--schnell', action='store_true', help='run flux schnell instead of dev')
89+
parser.add_argument('--width', type=int, default=1024, help='width of the image to generate')
90+
parser.add_argument('--height', type=int, default=1024, help='height of the image to generate')
91+
parser.add_argument('--guidance', type=float, default=3.5, help='gauidance strentgh for dev')
92+
parser.add_argument('--seed', type=int, default=None, help='seed for inference')
93+
parser.add_argument('--profile', action='store_true', help='enable profiling')
94+
parser.add_argument('--profile-duration', type=int, default=10000, help='duration for profiling in msec.')
95+
parser.add_argument('--itters', type=int, default=15, help='tiems to run inference and get avg time in sec.')
96+
args = parser.parse_args()
97+
if args.schnell:
98+
ckpt_id = "black-forest-labs/FLUX.1-schnell"
99+
else:
100+
ckpt_id = "black-forest-labs/FLUX.1-dev"
101+
text_pipe = FluxPipeline.from_pretrained(ckpt_id, transformer=None, vae=None, torch_dtype=torch.bfloat16).to('cpu')
102+
xmp.spawn(_main, args=(args, text_pipe, ckpt_id))

‎examples/research_projects/pytorch_xla/README.md‎ renamed to ‎examples/research_projects/pytorch_xla/training/text_to_image/README.md‎

File renamed without changes.

‎examples/research_projects/pytorch_xla/requirements.txt‎ renamed to ‎examples/research_projects/pytorch_xla/training/text_to_image/requirements.txt‎

File renamed without changes.

‎examples/research_projects/pytorch_xla/train_text_to_image_xla.py‎ renamed to ‎examples/research_projects/pytorch_xla/training/text_to_image/train_text_to_image_xla.py‎

File renamed without changes.

‎src/diffusers/models/attention_processor.py‎

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2318,9 +2318,12 @@ def __call__(
23182318
query = apply_rotary_emb(query, image_rotary_emb)
23192319
key = apply_rotary_emb(key, image_rotary_emb)
23202320

2321-
hidden_states = F.scaled_dot_product_attention(
2322-
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
2323-
)
2321+
if XLA_AVAILABLE:
2322+
query /= math.sqrt(head_dim)
2323+
hidden_states = flash_attention(query, key, value, causal=False)
2324+
else:
2325+
hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)
2326+
23242327
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
23252328
hidden_states = hidden_states.to(query.dtype)
23262329

@@ -2521,7 +2524,12 @@ def __call__(
25212524
query = apply_rotary_emb(query, image_rotary_emb)
25222525
key = apply_rotary_emb(key, image_rotary_emb)
25232526

2524-
hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)
2527+
if XLA_AVAILABLE:
2528+
query /= math.sqrt(head_dim)
2529+
hidden_states = flash_attention(query, key, value)
2530+
else:
2531+
hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)
2532+
25252533
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
25262534
hidden_states = hidden_states.to(query.dtype)
25272535

0 commit comments

Comments
 (0)