Trained a 210M text-to-image diffusion transformer from scratch on one GPU in 3.5 days

IvanMikhnenkov · reddit · 2026-09-10

To gain hands-on experience training a flow model from scratch, the author built and trained a 210M-parameter diffusion transformer: 4.2M curated images at 256², rectified flow on the FLUX.2 VAE, flan-t5-base for text (128 tokens max). Only the VAE and text encoder are frozen pretrained pieces; the transformer, recipe, data pipeline, and evaluation are original. The full run took 3.5 days on a single RTX PRO 6000.

Key lessons, in order of importance:

Captions that actually fit the images mattered most: a web crawl made the model worse; curated photos with good captions fixed it.
A timestep shift for the 32-channel latent, plus aspect-ratio buckets from step one instead of square crops.
Register tokens with learned null attention slots — the null slots ended up absorbing 90% of cross-attention.
torch.compile for training, not just inference: 2.4× faster.
Training loss stopped being informative after day one while images kept improving, so they track FID, detector-based object accuracy, and human-preference models instead.

Everything is open: weights (CC BY-NC), code, every decision with sources, a training dashboard, and an attention playground, plus a browser demo. Next step: RL fine-tuning with Flow-GRPO using the failure grid as the target list.

Original post →

Trained a 210M text-to-image diffusion transformer from scratch on one GPU in 3.5 days

More from Multimodal