A from-scratch latent Diffusion Transformer trained with rectified flow on 2M LAION-Aesthetics images — with measured, profiler-driven performance engineering.
The attention, AdaLN-Zero conditioning, patchify/unpatchify, rectified-flow objective, Euler sampler, and classifier-free guidance are all hand-written (~600 lines of model code). The only pretrained components are a frozen Stable Diffusion VAE and a frozen CLIP ViT-L/14 text encoder. Training streams WebDataset tar shards under DeepSpeed ZeRO-2 on Slurm.
- 254.5M-param MM-DiT-style backbone — dual-stream image/text tokens fused by joint attention (not U-Net cross-attention), AdaLN-Zero time conditioning, QK-RMSNorm, SwiGLU.
- Rectified flow, not DDPM — training reduces to velocity regression
‖v_θ(z_t, c, t) − (z₁ − z₀)‖²; sampling is a plain Euler ODE from noise to data. - 2.2× faster / 2× less memory from a profiler-driven attention fix (measured, see below).
- Streaming data pipeline over 1,976 WebDataset shards (2.0M captioned images, ~200 GB) — measured at 757 imgs/s, 29× the model's consumption.
- Profiling harness included (
profile_laion.py): forward/backward split, peak memory, NCCL comms share, and dataloader throughput, on 1–N GPUs.
torch.profiler traces showed the original hand-written attention materialized every layer's
B×12×1101×1101 score matrix and ran softmax in fp32 — the memory wall and ~21% of GPU time in
dtype casts. Every OOM crash landed on the same softmax(scores.float()) line. Routing the same
QKV layout through F.scaled_dot_product_attention (FlashAttention kernels):
| batch 16, 12 layers, 512px | manual attention | SDPA / Flash | gain |
|---|---|---|---|
| throughput | 26.1 imgs/s | 56.7 imgs/s | 2.2× |
| forward / backward | 212 / 368 ms | 75 / 174 ms | 2.8× / 2.1× |
| peak memory | 35.4 GB | 17.2 GB | 2.05× |
| max batch per 44 GB | 16 (OOM @ 24) | ≥ 32 (OOM @ 64) | ~2× |
Two further profiler findings:
-
Compute-bound, not data-bound. The WebDataset loader sustains 756.8 imgs/s against a model consumption of 26–57 imgs/s; the GPU waits 0.2 ms per batch for data.
-
42% of every step re-encodes frozen models. In the full pipeline, VAE image encoding (18%) plus CLIP text encoding (10%) plus the remaining load path (14%) consume as much CUDA time as the entire DiT forward pass. Precomputing latents + embeddings offline is a ~2× end-to-end win and the top roadmap item.
Multi-GPU (DDP, 2× L40S): 86% scaling efficiency; NCCL all-reduce is 13% of CUDA time for the 254M-param model.
Raw JSON results and Chrome traces land in checkpoints/<run>/profile/; every number above is
reproducible with scripts/run_profile.sh and scripts/run_flash_compare.sh.
Instead of a β noise schedule, define a linear interpolation between data latent z₀ and noise
z₁ ~ N(0, I):
z_t = z₀ + t · (z₁ − z₀), t ~ U(0, 1)
v* = z₁ − z₀
L = E ‖ v_θ(z_t, c, t) − v* ‖²
The velocity target is constant along each path, so diffusion training becomes direct regression
(diffusion.py). Sampling integrates dx/dt = v(x, c, t) from
t=1 → 0 with Euler steps and classifier-free guidance
(sampler.py); captions are dropped with p_uncond = 0.1 during
training so the zero-embedding serves as the unconditional branch.
Images (512²) are encoded by a frozen stabilityai/sd-vae-ft-mse into 4×64×64 latents (scaled by the
VAE's scaling_factor) — 16× less spatial area for the backbone to model.
flowchart LR
subgraph inputs
Z["noisy latent z_t
4×64×64"]
C["caption tokens
CLIP ViT-L 77×768"]
T["timestep t"]
end
Z --> PE["patchify Conv2d(4→768, s=2)
+ learned pos-emb → 1024 tokens"]
T --> SIN["sinusoidal emb → MLP"]
C --> POOL["mean-pool → proj"]
SIN --> TT(("t̃"))
POOL --> TT
C --> CP["context proj 768→768"]
PE --> L1
CP --> L1
subgraph DiT["12 × DiT layer (dual-stream, joint attention)"]
L1["adaLN(x,t̃) & adaLN(c,t̃) → QKV each
concat [c;x] → joint MHSA → split
per-stream: fc → adaLN → SwiGLU MLP
single residual per stream"]
end
TT -.conditions every layer.-> DiT
L1 --> FN["LayerNorm → proj 768→16
(zero-init) → unpatchify"]
FN --> V["velocity v̂
4×64×64"]
Design notes (and deviations from SD3's MM-DiT):
| Choice | This repo | SD3 |
|---|---|---|
| Text encoders | CLIP-L only | CLIP-L + CLIP-G + T5-XXL |
| adaLN | scale + shift, zero-init | scale + shift + gate ×2 |
| Residuals per block | one, ungated | two, gated |
| MLP | SwiGLU | GELU |
| Positional embedding | learned | 2D sin-cos |
| QK norm | RMSNorm on Q and K | RMSNorm |
Zero-initializing both the adaLN modulation and the final output projection makes the network start as an identity/zero-velocity map, which keeps early training stable without gates.
Training uses a captioned, aspect-bucketed export of LAION-Aesthetics at 512 base resolution: 1,999,908 images in 1,976 uncompressed WebDataset tar shards (~200 GB) across 40 SDXL-style aspect-ratio buckets. Images were cover-resized (bicubic, antialias, never upsampled), deterministically cropped, re-encoded at JPEG q95, and re-captioned with VLMs; a manifest records per-shard sha256 checksums.
The loader (webdataset_loader.py) splits
shards across DDP ranks (shards[rank::world_size]) and DataLoader workers, with a shard shuffle +
10k-sample shuffle buffer and warn_and_continue fault tolerance — no global index, no full
download, sequential tar reads.
The current training transform center-crops to square 512²; exploiting the aspect buckets with multi-resolution batching is future work.
git clone https://github.com/Aaronaferns/re-diffusion.git
cd re-diffusion
uv sync # or: pip install -e .Pretrained encoders (stabilityai/sd-vae-ft-mse, openai/clip-vit-large-patch14) download from
Hugging Face on first use; pre-cache them if your compute nodes are offline.
# single node, N GPUs
python -m torch.distributed.run --standalone --nproc_per_node=N \
src/re_diffusion/train_laion.py --config configs/laion_h100.yaml
# Slurm
sbatch scripts/train_laion.slurm # 2× H100
sbatch scripts/train_laion_l40s.slurm # 8× L40SConfigs (configs/*.yaml) control batch size, LR schedule (linear warmup → cosine), bf16, ZeRO
stage, EMA decay, CFG drop rate, mid-training sampling (image grids + CLIPScore to TensorBoard), and
checkpointing. Training is step-based over the effectively infinite stream — no epochs.
# FID / Inception Score / CLIPScore against real shards
python src/re_diffusion/evaluate.py --config configs/eval.yaml --ckpt <ckpt.pt>
# profiling sweep: compute, memory, comms, dataloader
bash scripts/run_profile.sh
# manual-vs-SDPA attention comparison
bash scripts/run_flash_compare.sh
# regenerate the figures in assets/figures/
python scripts/make_figures.pysrc/re_diffusion/
model.py # DiT backbone: patchify, joint attention, adaLN-Zero, SwiGLU
diffusion.py # rectified-flow forward process, train steps, EMA
sampler.py # Euler ODE sampler with classifier-free guidance
train_laion.py # DeepSpeed ZeRO-2 streaming training loop
evaluate.py # FID / IS / CLIPScore
profile_laion.py # torch.profiler harness (compute / dataloader / full modes)
config.py # dataclass configs (YAML/JSON serialisable)
datasetProcessing/
webdataset_loader.py # rank- and worker-sharded WebDataset pipeline
configs/ # per-hardware training configs
scripts/ # Slurm launchers, profiling sweeps, figure generation
assets/figures/ # measured-results charts (committed)
- Fused SDPA / FlashAttention (2.2× measured)
- Euler sampler + classifier-free guidance
- Profiling harness with committed traces
- Offline VAE-latent + CLIP-embedding precompute (~2× end-to-end, measured headroom)
- Multi-resolution training on the native aspect buckets
- Gradient checkpointing / FSDP for larger backbones
- Optimizer state in checkpoints for exact resume (model + EMA are already saved)
- Peebles & Xie, Scalable Diffusion Models with Transformers (DiT), 2023
- Liu et al., Flow Straight and Fast: Rectified Flow, 2022
- Esser et al., Scaling Rectified Flow Transformers (SD3 / MM-DiT), 2024
- Rombach et al., High-Resolution Image Synthesis with Latent Diffusion Models, 2022