Training diffusion models with PyTorch

---

This tutorial shows how to train diffusion models on TPUs using PyTorch Lightning and Pytorch XLA.

Warning: This model uses a third-party dataset. Google provides no representation, warranty, or other guarantees about the validity, or any other aspects of this dataset.

Objectives

• Create a Cloud TPU
• Install PyTorch Lightning
• Clone the diffusion repo
• Prepare the Imagenette dataset
• Run the training script

Costs

In this document, you use the following billable components of Google Cloud:

• Compute Engine
• Cloud TPU

To generate a cost estimate based on your projected usage, use the pricing calculator. New Google Cloud users might be eligible for a free trial.

Before you begin

Before starting this tutorial, check that your Google Cloud project is correctly set up.

1. Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.

2. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

   Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

   Go to project selector

3. Make sure that billing is enabled for your Google Cloud project.

4. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

   Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

   Go to project selector

5. Make sure that billing is enabled for your Google Cloud project.

This walkthrough uses billable components of Google Cloud. Check the Cloud TPU pricing page to estimate your costs. Be sure to clean up resources you create when you've finished with them to avoid unnecessary charges.

Create a Cloud TPU

This tutorial uses a v4-8, but it works similarly on all accelerator sizes in a single host.

Set up some environment variables to make the commands easier to use.

export ZONE=us-central2-b
export PROJECT_ID=your-project-id
export ACCELERATOR_TYPE=v4-8
export RUNTIME_VERSION=tpu-ubuntu2204-base
export TPU_NAME=your_tpu_name

Create a Cloud TPU.

Note: If there is not enough capacity currently available to create the TPU Pod, you can queue your request using queued resources. Queued resources allow you to receive capacity once it becomes available. To request your Cloud TPU resources as queued resources, use the gcloud compute tpus queued-resources create command instead. For more information, see Manage Queued Resources.

gcloud compute tpus tpu-vm create ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--accelerator-type=${ACCELERATOR_TYPE} \
--version=${RUNTIME_VERSION} \
--subnetwork=tpusubnet

Note: The following commands run on all TPU VMs concurrently.

Install required software

1. Install required packages along with PyTorch/XLA latest release v2.5.0.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --command="sudo apt-get update -y && sudo apt-get install libgl1 -y
   git clone https://github.com/pytorch-tpu/stable-diffusion.git
   cd stable-diffusion
   pip install -r requirements.txt
   pip install -e .
   pip install https://github.com/Lightning-AI/lightning/archive/refs/heads/master.zip -U
   pip install -e git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers
   pip install clip
   pip install torch~=2.5.0 torch_xla[tpu]~=2.5.0 torchvision -f https://storage.googleapis.com/libtpu-releases/index.html"
   

2. Fix source files to be compatible with torch 2.2 and newer.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --command="cd stable-diffusion/
   sed -i 's/from torch._six import string_classes/string_classes = (str, bytes)/g' src/taming-transformers/taming/data/utils.py
   sed -i 's/trainer_kwargs\\[\"callbacks\"\\]/# trainer_kwargs\\[\"callbacks\"\\]/g' main_tpu.py"
   

3. Download Imagenette (a smaller version of Imagenet dataset) and move it to the appropriate directory.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --command="wget -nv https://s3.amazonaws.com/fast-ai-imageclas/imagenette2.tgz
   tar -xf  imagenette2.tgz
   mkdir -p ~/.cache/autoencoders/data/ILSVRC2012_train/data
   mkdir -p ~/.cache/autoencoders/data/ILSVRC2012_validation/data
   mv imagenette2/train/*  ~/.cache/autoencoders/data/ILSVRC2012_train/data
   mv imagenette2/val/* ~/.cache/autoencoders/data/ILSVRC2012_validation/data"
   

4. Download the first stage pretrained model.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --command="cd stable-diffusion/
   wget -nv -O models/first_stage_models/vq-f8/model.zip https://ommer-lab.com/files/latent-diffusion/vq-f8.zip
   cd  models/first_stage_models/vq-f8/
   unzip -o model.zip"
   

Train the model

Run the training with following command. Note that the training process is expected to take around 30 mins on the v4-8.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--command="python3 stable-diffusion/main_tpu.py --train --no-test --base=stable-diffusion/configs/latent-diffusion/cin-ldm-vq-f8-ss.yaml -- data.params.batch_size=32 lightning.trainer.max_epochs=5 model.params.first_stage_config.params.ckpt_path=stable-diffusion/models/first_stage_models/vq-f8/model.ckpt lightning.trainer.enable_checkpointing=False lightning.strategy.sync_module_states=False"

Clean up

Perform a cleanup to avoid incurring unnecessary charges to your account after using the resources you created:

Use Google Cloud CLI to delete the Cloud TPU resource.

  $  gcloud compute tpus tpu-vm delete diffusion-tutorial --zone=us-central2-b
  

What's next

• Training Resnet50 on Cloud TPU with PyTorch
• Troubleshooting Pytorch on TPUs
• Pytorch/XLA documentation