Scaling up a DiT: replicating Peebles & Xie, then optimizing the training loop
Went from a toy 2D spiral DDPM to a 130M-parameter class-conditional Diffusion Transformer on full ImageNet-1k, found an 8.3x training speedup with TF32 + bf16 + torch.compile, completed the paper's full 400K-step schedule, then scaled to the 675.8M-parameter DiT-XL/2 on 2 GPUs — where a late-training divergence turned into a lesson about gradient clipping and checkpoint hygiene.