Attempting to Train Stable Diffusion From Scratch For Free

Introduction

Is free of cost
Gives reasonable image outputs for a text input
Trains fast

TPUs have been very difficult for me to utilize effectively. While a TPU at 10% FLOPs utilization is still 2.24x as fast as two T4s, I literally have never seen TPU utilization go above 40% in any of my previous projects. However, part of this project was to understand how to get the monster FLOPs utilization seen in other large models (e.g. 54.9% model FLOPs utilization over 1024 TPU-v4s in ViT-22B [4, 5]), so I welcomed this challenge.
Kaggle TPU availability varies on the time of day and during peak usage hours (generally 2-8 PM UTC), you may have to wait up to two hours to get a hold of one.
TPUs require models to be compiled via XLA to be run. For the LDM that I trained (508 million parameter UNet), it takes around 30 minutes before the model actually trains on the TPU which makes profiling the model slow and painful.
Although I'm using torch_xla, XLA runs via TensorFlow which is notorious for having errors that make zero sense and generally being hard to debug.
Kaggle can't seem to control TPU sessions very well. Sometimes I run all notebook cells but some cells choose not to run so I have to rerun all notebook cells. Often I cancel a cell's execution and it doesn't stop unless I cancel it twice. Nearly every time I restart a session, Kaggle doesn't let me run any cells unless I restart it again. Most infuriating of them all, sometimes XLA runs into an error (usually an OOM) and the notebook just hangs. I cancel the cell execution, and nothing happens. I restart the session, and nothing happens. I restart it again, and nothing happens. The ONLY way I found to get around this situation is to factory reset the notebook, which sucks when you waited an hour to get a TPU session, wait 30 minutes for the model to compile, run into an error which hangs the notebook, and then now you have to wait another hour to have the factory reset-ted session to start running because you're working during TPU peak hours so you're #27 in the TPU queue. All of this happens exclusively on Kaggle TPU sessions.

Optimization Tricks

Logging loss every 32 steps (instead of every step)

Writing my own multi-head attention PyTorch module

Using SiLU instead of GELU

Splitting heads smarter

def slower_split_heads(x):
   splitted = x.reshape(B, L, D, H) # B L (H*D) -> B L D H
   return splitted.permute(0, 3, 1, 2) # B L D H -> B H L D; this can be fed into torch.nn.functional.scaled_dot_product_attention

def faster_split_heads(x):
   splitted = x.reshape(B, L, H, D) # B L (H*D) -> B L H D
   return splitted.permute(0, 2, 1, 3) # B L H D -> B H L D; this can be fed into torch.nn.functional.scaled_dot_product_attention

Avoiding torch.split like the plague

def slower_self_attention(x):
    qkv_proj = nn.Linear(embed_dim, 3*embed_dim)(x)
    q, k, v = torch.split(qkv_proj, embed_dim, dim=1)
    # scaled dot product attention and linear projection

def faster_self_attention(x):
    q = nn.Linear(embed_dim, embed_dim)(x)
    k = nn.Linear(embed_dim, embed_dim)(x)
    v = nn.Linear(embed_dim, embed_dim)(x)
    # scaled dot product attention and linear projection

def slower_swiglu(x):
    lin_gate = nn.Linear(channels, 2*channels)(x)
    lin, gate = torch.split(lin_gate, channels)
    return lin * torch.nn.functional.silu(gate)

def faster_swiglu(x):
    lin = nn.Linear(channels, channels)(x)
    gate = nn.Linear(channels, channels)(x)
    return lin * torch.nn.functional.silu(gate)

Removing biases in transformer blocks

Modifying VAE and DM architecture

Using larger batch sizes to hide CPU bottleneck

Using torch.compile(..., backend='torch_xla_trace_once') for the CLIP embedder

If this model is tracing bound, I think maybe we should give dynamo a try

FLOPs utilization goes from 55.3% without compiling to 68.4% with compiling

Trace viewer when compiling the VAE and CLIP embedder. Look, no idle time!





Strange Bugs

XLA_USE_BF16=1 destroys training

Using bfloat16 results in strange noise scheduler behavior when compared to float32.

self.betas = 1 - self.alphas
...
beta_ratio = self.betas[t-1] / self.betas[t]

def norm(x): return (x - x.mean()) / x.std()
x = torch.randn((10,), dtype=torch.bfloat16)
x1000 = x + 1000 # norm(x1000) should equal norm(x) which is roughly equal to x

class StableNorm(nn.Module): # runs GroupNorm in FP32 because of bfloat16 stability issues when x is large but with small variance (i.e. x = 100) 
    def __init__(self, num_groups: int, num_channels: int):
        super().__init__()
        self.norm = nn.GroupNorm(num_groups, num_channels).double() # under XLA_DOWNCAST_BF16=1, double is casted down to fp32
    
    def forward(self, x):
        return self.norm(x.double()).float() # under XLA_DOWNCAST_BF16=1, float is casted down to bf16

Using TFRecordDataset causes train function to hang before training begins

Einops sometimes doesn't work

Setting scale_factor argument when upsampling images slows training

def forward(self, x):
    # BEFORE
    #return self.conv(F.interpolate(x, scale_factor=2.0, mode='nearest'))

    # AFTER
    _B, _C, H, W = x.shape
    return self.conv(F.interpolate(x, size=(H*2, W*2), mode='nearest'))

Not logging loss with loss.item() causes OOMs

More CPU memory is used up every epoch

# BEFORE
def train_loop(self, epochs=1, steps_per_epoch=-1, save_every_n_epochs=1):
    for epoch in range(epochs):
        loader = pl.MpDeviceLoader(self.loader, self.device)
        for step, (imgs, _captions) in enumerate(loader):
            # training step
            if step == steps_per_epoch - 1:
                # save models, log images, etc.
                break

# AFTER
def train_loop(self, epochs=1, steps_per_epoch=-1, save_every_n_epochs=1):
    loader = pl.MpDeviceLoader(self.loader, self.device)
    for step, (imgs, _captions) in enumerate(loader):
        # training step
        if step == steps_per_epoch - 1:
            # save models, log images, etc.
            epoch += 1
            if epoch == epochs:
                break

Training from a model checkpoint is slower than training from random initialization

BEFORE FIX: Notice that the TPUs are idling and its cause: opt_step (correlating to an opt.step() call) has really big Python overhead. Something is up with the optimizer! (Note: this actually runs at around 1.45 steps/s, I don't know why this is slower than the 1.13 steps/s run, but it highlights the same problem)

# OLD
def load_models(device='cpu', lr=1e-4, ckpt_path=None):
    vae = VAE().to(device)
    disc = Discriminator().to(device)    
    vae_opt = Adam(vae.parameters(), lr, (0.5, 0.9))
    disc_opt = Adam(disc.parameters(), lr, (0.5, 0.9))

    if ckpt_path is not None:
        ckpt = torch.load(ckpt_path, map_location=torch.device('cpu'))
        vae.load_state_dict(ckpt['vae_state_dict'])
        disc.load_state_dict(ckpt['disc_state_dict'])
        vae_opt.load_state_dict(ckpt['vae_opt_state_dict'])
        disc_opt.load_state_dict(ckpt['disc_opt_state_dict'])
    ...

# NEW - I was thinking of finding a nicer way to write this fix in code, but I didn't care enough
def load_models(device='cpu', lr=1e-4, ckpt_path=None):
    vae = VAE().to(device)
    disc = Discriminator().to(device)

    if ckpt_path is None:
        vae_opt = Adam(vae.parameters(), lr, (0.5, 0.9))
        disc_opt = Adam(disc.parameters(), lr, (0.5, 0.9))
    else:
        ckpt = torch.load(ckpt_path)
        vae.load_state_dict(ckpt['vae_state_dict'])
        disc.load_state_dict(ckpt['disc_state_dict'])
        vae = vae.to(device)
        disc = disc.to(device)

	# now the optimizers will be initialized with on-device parameters
        vae_opt = Adam(vae.parameters(), lr, (0.5, 0.9))
        disc_opt = Adam(disc.parameters(), lr, (0.5, 0.9))
        vae_opt.load_state_dict(ckpt['vae_opt_state_dict'])
        disc_opt.load_state_dict(ckpt['disc_opt_state_dict'])
    ...

Calling loss.item() on only the master ordinal causes training to hang

def log_step(self,
    log_metrics: dict, step: int,
    epoch: int, epochs: int, steps_per_epoch: int, pbar
):
    log_items = {k: v.item() for k, v in log_metrics.items()} # AFTER - no training hang
    if xm.is_master_ordinal():
        # log_items = {k: v.item() for k, v in log_metrics.items()} # BEFORE - causes training hang
        if self.wandb_run is not None:
            wandb.log(log_items, step=self.global_step) 

The training hang would always happen when I was saving my model, so I thought my problem had to do with saving the model. But as you can see above, the root of this model is unrelated to model saving, so I was spending a lot of time focusing on the wrong part of my code.
At this point, I had little experience debugging programs that would stop and idle. I've worked with programs that would freeze because it would run into an infinite loop, but that wouldn't make the CPUs start idling. I didn't know why a program would start doing nothing, so to me, the root of this problem could've been anywhere and thus I looked everywhere in my code for potential problems, mostly in areas with no issues. Also because of the lack of errors, I spent a lot of time on runs where I just printed out extra messages before freezing to help isolate the problem.
Since XLA takes a while to compile big models, I tried debugging with small models. However, when I did this, I couldn't reproduce the training hang.

The same model in different parts of training use different amounts of CPU RAM

Things That Didn't Work

Tokenizing captions in DataLoader

Training fast with FSDP

Non-sinusoidal timestep embeddings

def sinusoidal_timestep_embedding(timestep: int): # timestep between 0-999 inclusive
  sinusoid_emb = sinusoidal_embedding_fn(timestep) # takes in int, returns length 256 vector
  x1 = Linear(256, 1024)(sinusoid_emb)
  x2 = SiLU(x1)
  x3 = Linear(1024, 1024)(x2)
  return x3

I was effectively reducing my batch size on the embedding by a factor of 1000. To clarify, with a batch size of 256, and assuming no repeated timesteps, I would only have gradients for 256 rows of the embedding per batch. But of the embeddings that were used in the forward pass, since each embedding was used for only one element within the batch (effective batch size of one if a gradient was received and a batch size of zero otherwise), its gradients would be very noisy. For the original sinusoidal timestep embedding, the learnable parameters rest in the linear layers which means that all of the parameters would be used for each element within the batch (effective batch size of 256).
Because the different timesteps represent slowly changing levels of noise strength in the input image, it makes sense that rows of timestep embedding that are spatially close to each other (e.g. row 10 and row 11) should have similar element contents (i.e. the vector distance between the rows should be close to 0) since spatially close rows represent similar levels of noise strength. A vanilla learnable embedding does not have this inductive bias whereas the sinusoidal embedding does, which means that training should be much easier with the sinusoidal embedding.

The k-th row on both images represents the activation values of each embedding for t=k. Notice how the different rows in the vanilla embedding are completely independent of each other and how this is not the case for the sinusoidal embedding.

Training Progress

VAE

DM

Conclusion

Read the tutorials carefully. Many of the issues I faced were things that that I glossed over in tutorials.
Make your debug runs as small and focused as possible. I used up so much time debugging as I would spend 30 minutes of waiting, get an error message, change one or two line of code, and repeat. For example, if you can replicate the issue with a smaller model, use a smaller model. If you can replicate the issue using only one device rather than eight, use only one device. You get the idea.
Document your work. It's good to do this in case you encounter the same problem in the future so you can refer to how you fixed the problem. I documented my work by writing descriptions on each W&B run of what I changed and of any new issues that happened in the run. I also saved my training code and uploaded it on W&B for every run. Making this report would've been impossible if I didn't do both of these.
Use built-in debugging tools and profilers.

Sources