Skip to content

Commit 68ff594

Browse files
jfacevedo-googlejfacevedo-google
committed
flux
1 parent 2809859 commit 68ff594

File tree

3 files changed

+1094
-0
lines changed

3 files changed

+1094
-0
lines changed

‎examples/research_projects/pytorch_xla/training/text_to_image/README_flux.md‎

Lines changed: 181 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,181 @@
1+
# Stable Diffusion XL text-to-image fine-tuning using PyTorch/XLA
2+
3+
The `train_text_to_image_xla.py` script shows how to fine-tune stable diffusion model on TPU devices using PyTorch/XLA.
4+
5+
It has been tested on v4 and v5p TPU versions. Training code has been tested on multi-host.
6+
7+
This script implements Distributed Data Parallel using GSPMD feature in XLA compiler
8+
where we shard the input batches over the TPU devices.
9+
10+
As of 10-31-2024, these are some expected step times.
11+
12+
| accelerator | global batch size | step time (seconds) |
13+
| ----------- | ----------------- | --------- |
14+
| v5p-512 | 16384 | 1.01 |
15+
| v5p-256 | 8192 | 1.01 |
16+
| v5p-128 | 4096 | 1.0 |
17+
| v5p-64 | 2048 | 1.01 |
18+
19+
## Create TPU
20+
21+
To create a TPU on Google Cloud first set these environment variables:
22+
23+
```bash
24+
export TPU_NAME=<tpu-name>
25+
export PROJECT_ID=<project-id>
26+
export ZONE=<google-cloud-zone>
27+
export ACCELERATOR_TYPE=<accelerator type like v5p-8>
28+
export RUNTIME_VERSION=<runtime version like v2-alpha-tpuv5 for v5p>
29+
```
30+
31+
Then run the create TPU command:
32+
```bash
33+
gcloud alpha compute tpus tpu-vm create ${TPU_NAME} --project ${PROJECT_ID}
34+
--zone ${ZONE} --accelerator-type ${ACCELERATOR_TYPE} --version ${RUNTIME_VERSION}
35+
--reserved
36+
```
37+
38+
You can also use other ways to reserve TPUs like GKE or queued resources.
39+
40+
## Setup TPU environment
41+
42+
Install PyTorch and PyTorch/XLA nightly versions:
43+
```bash
44+
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
45+
--project=${PROJECT_ID} --zone=${ZONE} --worker=all \
46+
--command='
47+
pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu
48+
pip install 'torch_xla[tpu] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.8.0.dev-cp310-cp310-linux_x86_64.whl' \
49+
-f https://storage.googleapis.com/libtpu-releases/index.html \
50+
-f https://storage.googleapis.com/libtpu-wheels/index.html
51+
52+
# Optional: if you're using custom kernels, install pallas dependencies
53+
pip install 'torch_xla[pallas]' \
54+
-f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html \
55+
-f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
56+
'
57+
```
58+
59+
Verify that PyTorch and PyTorch/XLA were installed correctly:
60+
61+
```bash
62+
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
63+
--project ${PROJECT_ID} --zone ${ZONE} --worker=all \
64+
--command='python3 -c "import torch; import torch_xla;"'
65+
```
66+
67+
Install dependencies:
68+
```bash
69+
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
70+
--project=${PROJECT_ID} --zone=${ZONE} --worker=all \
71+
--command='
72+
git clone https://github.com/huggingface/diffusers.git
73+
cd diffusers
74+
git checkout main
75+
cd examples/research_projects/pytorch_xla/training/text_to_image/
76+
pip3 install -r requirements_sdxl.txt
77+
cd ../../../../../
78+
pip3 install .'
79+
```
80+
81+
## Run the training job
82+
83+
### Authenticate
84+
85+
Run the following command to authenticate your token.
86+
87+
```bash
88+
huggingface-cli login
89+
```
90+
91+
This script only trains the unet part of the network. The VAE and text encoder
92+
are fixed.
93+
94+
```bash
95+
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
96+
--project=${PROJECT_ID} --zone=${ZONE} --worker=all \
97+
--command='
98+
export XLA_DISABLE_FUNCTIONALIZATION=1
99+
export TORCH_DISABLE_FUNCTIONALIZATION_META_REFERENCE=1
100+
export PROFILE_DIR=/tmp/
101+
export CACHE_DIR=/tmp/
102+
export DATASET_NAME=lambdalabs/naruto-blip-captions
103+
export PER_HOST_BATCH_SIZE=4
104+
export TRAIN_STEPS=50
105+
export OUTPUT_DIR=/tmp/trained-model/
106+
python examples/research_projects/pytorch_xla/training/text_to_image/train_text_to_image_flux.py --pretrained_model_name_or_path=black-forest-labs/FLUX.1-dev --dataset_name=$DATASET_NAME --resolution=1024 --center_crop --random_flip --train_batch_size=$PER_HOST_BATCH_SIZE --max_train_steps=$TRAIN_STEPS --learning_rate=1e-06 --mixed_precision=bf16 --profile_duration=80000 --output_dir=$OUTPUT_DIR --dataloader_num_workers=4 --loader_prefetch_size=4 --device_prefetch_size=4'
107+
```
108+
109+
Pass `--print_loss` if you would like to see the loss printed at every step. Be aware that printing the loss at every step disrupts the optimized flow execution, thus the step time will be longer.
110+
111+
### Environment Envs Explained
112+
113+
* `XLA_DISABLE_FUNCTIONALIZATION`: To optimize the performance for AdamW optimizer.
114+
* `PROFILE_DIR`: Specify where to put the profiling results.
115+
* `CACHE_DIR`: Directory to store XLA compiled graphs for persistent caching.
116+
* `DATASET_NAME`: Dataset to train the model.
117+
* `PER_HOST_BATCH_SIZE`: Size of the batch to load per CPU host. For e.g. for a v5p-16 with 2 CPU hosts, the global batch size will be 2xPER_HOST_BATCH_SIZE. The input batch is sharded along the batch axis.
118+
* `TRAIN_STEPS`: Total number of training steps to run the training for.
119+
* `OUTPUT_DIR`: Directory to store the fine-tuned model.
120+
121+
## Run inference using the output model
122+
123+
To run inference using the output, you can simply load the model and pass it
124+
input prompts. The first pass will compile the graph and takes longer with the following passes running much faster.
125+
126+
```bash
127+
export CACHE_DIR=/tmp/
128+
```
129+
130+
```python
131+
import torch
132+
import os
133+
import sys
134+
import numpy as np
135+
136+
import torch_xla.core.xla_model as xm
137+
from time import time
138+
from diffusers import StableDiffusionXLPipeline
139+
import torch_xla.runtime as xr
140+
141+
MODEL_PATH = os.environ.get("OUTPUT_DIR", None)
142+
143+
CACHE_DIR = os.environ.get("CACHE_DIR", None)
144+
if CACHE_DIR:
145+
xr.initialize_cache(CACHE_DIR, readonly=False)
146+
147+
def main():
148+
device = xm.xla_device()
149+
pipe = StableDiffusionXLPipeline.from_pretrained(
150+
MODEL_PATH,
151+
torch_dtype=torch.bfloat16
152+
)
153+
pipe.to(device)
154+
pipe.unet.enable_xla_flash_attention(partition_spec=("data", None, None, None))
155+
prompt = ["A naruto with green eyes and red legs."]
156+
start = time()
157+
print("compiling...")
158+
image = pipe(prompt, num_inference_steps=30, guidance_scale=7.5).images[0]
159+
print(f"compile time: {time() - start}")
160+
print("generate...")
161+
start = time()
162+
image = pipe(prompt, num_inference_steps=30, guidance_scale=7.5).images[0]
163+
print(f"generation time (after compile) : {time() - start}")
164+
image.save("naruto.png")
165+
166+
if __name__ == '__main__':
167+
main()
168+
```
169+
170+
Expected Results:
171+
172+
```bash
173+
WARNING:root:libtpu.so and TPU device found. Setting PJRT_DEVICE=TPU.
174+
Loading pipeline components...: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 7.93it/s]
175+
compiling...
176+
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 30/30 [01:35<00:00, 3.19s/it]
177+
compile time: 241.23492813110352
178+
generate...
179+
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 30/30 [00:04<00:00, 6.72it/s]
180+
generation time (after compile) : 5.266263246536255
181+
```

‎examples/research_projects/pytorch_xla/training/text_to_image/requirements_flux.txt‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
accelerate>=0.16.0
2+
transformers>=4.25.1
3+
datasets>=2.19.1
4+
ftfy
5+
tensorboard
6+
Jinja2
7+
peft==0.7.0
8+
sentencepiece==0.2.0

0 commit comments

Comments
 (0)