This project is a from-scratch implementation of diffusion model training in raw C++/CUDA. It is currently in progress, with support for both the classic UNet architecture, based on Diffusion Models Beat GANs on Image Synthesis, and the transformer architecture (DiT), as detailed in Scalable Diffusion Models with Transformers. My work is focused on developing the DiT model from scratch, while also enhancing Chen Lu's unet.cu by adding distributed training support and optimizations such as mixed precision training. The project is inspired by Andrej Karpathy's llm.c.
UNet currently supports training. You can train it on the images from the ImageNet 64x64 dataset via
gunzip unet/data/elephant_train.bin.gz
python unet/train_diffusion.py --init_model_only True
make -C unet train_diffusion
./unet/train_diffusion
This currently supports unconditional diffusion model training, and the end-to-end training loop is currently running at about 42% the speed of PyTorch with torch.compile when run on a single H100. Further detailed benchmarks will have to be done to understand bottlenecks + adjust implementation for better performance. I do think we can incorporate mixed-precision training here, though (FP16 w/ loss scaling).
| Platform | Time on H100 (ms) |
|---|---|
| This repo (CUDA implementation) | 56.98 |
PyTorch (w/ torch.compile) |
23.68 |
In Progress:
- support for distributed training via MPI in UNet
- support for mixed precision training in UNet
- support for DiT full fledged training
I've always been intrigued by diffusion models but found the math and implementation challenging. My interest in ML systems and GPU programming led me to start this project. Inspired by Karpathy's llm.c, I aimed to directly program the GPU for faster, more efficient training.
My goal is to develop a solution that could potentially surpass PyTorch's torch.compile, which optimizes model execution on NVIDIA GPUs using advanced techniques like JIT compilation, operator fusion, and kernel optimizations. These optimizations significantly improve runtime performance by reducing overhead and maximizing hardware resource utilization.
If you're interested in learning more about diffusion models and CUDA programming, here are some resources that I found incredibly helpful:
-
Understanding Diffusion Models:
- https://www.youtube.com/watch?v=W-O7AZNzbzQ - This video provides a great explanation of the research paper.
- https://www.youtube.com/watch?v=HoKDTa5jHvg - If you're struggling with the math behind diffusion models, like I was, this video is a lifesaver.
-
GPU Programming:
- Programming Massively Parallel Processors (Book & Lecture Series): https://www.youtube.com/playlist?list=PLRRuQYjFhpmubuwx-w8X964ofVkW1T8O4 - This is a great starting point for learning the fundamentals of GPU Programming/HPC.
- Getting Started with CUDA for Python Programmers: https://www.youtube.com/watch?v=nOxKexn3iBo - Great introductory YouTube series specifically for Python programmers venturing into CUDA.
- My Optimization Bible: CUDA Matrix Multiplication Optimization Tutorial: https://siboehm.com/articles/22/CUDA-MMM - This tutorial is where I learned the majority of the optimization techniques I used in this project. Highly recommended!
Articles/Blogs
- GPU Programming
- The CUDA Parallel Programming Model
- A HISTORY OF NVIDIA STREAM MULTIPROCESSOR
- Parallel Thread Execution
- How to Optimize a CUDA Matmul Kernel for cuBLAS-like Performance: a Worklog
- Making Deep Learning Go Brrrr From First Principles
- CUDA Matrix Multiplication Optimization
- What Every Developer Should Know About GPU Computing
- A minimal GPU design in Verilog to learn how GPUs work from the ground up
- GPU Programming: When, Why and How?
- Understanding GPU internals
- Understanding the GPU programming model
Tutorials
Notebooks
Videos
- How GPU Computing Works
- Getting Started With CUDA for Python Programmers
- Programming Massively Parallel Processors - Lecture Series by the Book Author
- Programming Massively Parallel Processors: A Hands-on Approach and then this YT series
- Programming Parallel Computers
- GPU Programming Lectures
- From Scratch CUDA
- CUDA Programming
- CUDA MODE Lectures
- Original paper implementation by OpenAI: https://github.com/openai/guided-diffusion
- Got most of complex cuda implementations (attention, groupnorm, etc.) from here: https://github.com/karpathy/llm.c
- For custom unet implementation: https://github.com/clu0/unet.cu