Training a Diffusion Model From Scratch in 3.5 Days

Training a diffusion model from scratch used to imply a cluster invoice and a research team. A solo build posted by Ivan Mikhnenkov compresses it to one workstation card: a 210M-parameter text-to-image diffusion transformer, a 4.2M-image dataset at 256² resolution, 3.5 days of training on a single RTX PRO 6000, with weights and a write-up attached. What follows treats that run as a measurement problem rather than another code walkthrough. Reconstructed from raw specs, it consumed roughly 100M samples seen (not 4.2M), sustained on the order of 100 TFLOPS, and cost somewhere between about $15 of electricity on owned silicon and the low hundreds of dollars rented. Run the same arithmetic on your own plans and the pretrain-versus-fine-tune decision stops being a budget question and becomes a distribution-distance question, which is the one that actually matters.

The Run and Its Three Raw Numbers

The post states its build directly and derives nothing, which is exactly what makes it auditable. Separate what was reported from what it implies:

Number	Value	Status
Parameters	210M	reported directly
Dataset	4.2M images at 256²	reported, ambiguous (unique vs seen)
Wall clock	3.5 days (84 h) on one card	reported directly
Batch × steps	256 × 400,000	reported directly
Tokens per image	~256 latent tokens, five aspect buckets	reported design choice
Text encoder	flan-t5-base, frozen	reported directly

Three findings in the post are genuinely novel: learned null-attention slots absorbing roughly 90% of cross-attention mass mid-network, a validation loss that barely moved while held-out FID fell from 33.7 to 27.0, and an inference timestep shift worth more than doubling the sample count. Read it for those. This teardown covers the numbers the author did not compute, starting with the load-bearing ambiguity in "4.2M images." Read as samples seen, that figure contradicts the post's own arithmetic, as the next two sections show. Read as dataset size, everything reconciles.

For scale, 4.2M images is about 1% of LAION-400M, the 400-million-pair open dataset that kicked off web-scale image-text scraping (the LAION-400M paper). A run this small is a curated-niche exercise, not a web-prior one, and that distinction drives the decision rule at the end.

Throughput You Can Derive Without Assumptions

Two throughput numbers fall out of division alone, and they answer different questions.

coverage:  4.2M images ÷ 3.5 days   = 1.2M images/day  ≈ 13.9 images/s
training:  256 × 400k samples ÷ 84h = 29M samples/day  ≈ 339 samples/s

Coverage throughput, about 1.2M unique images per day, needs zero architectural assumptions. You can compare it across any card, any model, any resolution, which makes it the only number in the thread that transfers unmodified to your hardware planning.

The training clock runs 24 times faster because it counts every sample the optimizer touched. The gap between the two clocks is the epoch count: 102.4M samples ÷ 4.2M images is about 24 passes, matching the author's remark that train and validation loss tracked each other for 24 epochs. As a single RTX PRO 6000 diffusion training throughput datapoint, roughly 339 samples per second at 256 latent tokens is the figure to quote when sizing your own runs.

Reconstructing the FLOP Budget

Training compute for a transformer follows 6ND: about 2ND for the forward pass, 4ND for the backward. For diffusion, D counts samples seen times tokens per image, because every sample is re-noised and revisited at a new noise level on each visit. The FLOP budget for a small text-to-image model is therefore fixed by three published numbers.

The single-pass floor does not work

Suppose 4.2M meant one training pass. Tokens would be 4.2M × 256, about 1.1B, and total compute 6 × 210M × 1.1B, roughly 1.4 × 10^18 FLOPs. A Blackwell-class card at a realistic sustained 100 to 150 TFLOPS delivers that in 2.5 to 4 hours; even at dense BF16 peak it is under an hour. A compiled, batch-256 run sitting near 1% utilization for 84 hours is not credible, and the author's own 24-epoch remark independently rejects the reading. So 4.2M is dataset size, and the interesting question becomes what 84 hours actually bought.

What 400k steps bought

samples seen = 256 × 400,000      = 102.4M
tokens       = 102.4M × 256       ≈ 26.2B
FLOPs        = 6 × 210M × 26.2B   ≈ 3.3 × 10^19
sustained    = 3.3e19 ÷ 302,400 s ≈ 109 TFLOPS

About 3.3 × 10^19 FLOPs, or 109 TFLOPS sustained for 3.5 days straight. Is that good? The RTX PRO 6000 datasheet headlines AI TOPS, a figure that assumes FP4 precision with sparsity, a mode no diffusion training loop actually runs in. Realizable dense BF16 or FP8 training throughput sits several times lower, so any utilization math that divides by marketing TOPS overstates capacity by roughly 4× or more. Against a dense BF16 ceiling in the several-hundred-TFLOPS range, 109 TFLOPS is about 20 to 25% model FLOPs utilization (model FLOPs utilization primer), which is unglamorous and entirely typical for a small model with short 256-token sequences. The estimate even slightly understates the run, since 6ND ignores register tokens, cross-attention into the frozen encoder, and the EMA copy. This run is reproducible arithmetic, not a miracle.

What Training a Diffusion Model From Scratch Actually Costs

Two ways to own the same 84 GPU-hours. If you already have the card, the marginal cost is electricity: roughly 600 W of board power plus CPU, memory, and cooling lands near 1 kW at the wall, or about 84 kWh over the run. At typical residential rates of 15 to 25 cents per kWh, that is $13 to $21. If you rent, 96 GB-class cards commonly run a few dollars per hour.

Cost line	Assumptions	Total
Electricity	~1 kW wall, 84 kWh at $0.15–0.25/kWh	$13–21
Cloud rental	84 GPU-hours at $2–4/h	$170–340
Per epoch, rented	÷ ~24 epochs	~$7–14
Per million samples seen, rented	÷ 102.4M samples	~$2–3

At this scale, the answer to "how many GPU hours to train a diffusion model" is 84, and the cost to train a 200M-class diffusion transformer is tens of dollars on owned silicon versus low hundreds rented. The card itself is a several-thousand-dollar purchase that amortizes across every future run; electricity is the marginal cost that actually gates iteration. The per-million-samples figure is the one to hold against per-image API pricing once volume gets serious.

Scaling to 1B Parameters and 40M Images

Compute scales with parameters at fixed token count, so a 1B model multiplies the budget by 1000/210, about 4.8×, to roughly 1.6 × 10^20 FLOPs. At the same 109 TFLOPS that is about 17 days, call it two to three additional weeks, plausibly less since larger models often sustain higher utilization on the same card, and the datasheet's 96 GB of memory means it fits.

Dataset scale is the sharper multiplier:

Scenario	Samples seen	FLOPs	At ~109 TFLOPS
This run: 210M, 4.2M images × ~24 epochs	102M	3.3 × 10^19	3.5 days
210M, 40M images, single pass	40M	1.3 × 10^19	~1.4 days (undertrained)
210M, 40M images × ~24 epochs	960M	3.1 × 10^20	~1 month
1B, 40M images × ~24 epochs	960M	1.5 × 10^21	~5 months

Two cautions before you extrapolate. First, the Diffusion Transformers (DiT) paper's FID-versus-Gflops curves are roughly log-linear, so each 4× jump in compute buys a bounded, compressing step in visible quality; 9.4× the compute for 40M images at matched epochs does not buy 9.4× the model. Second, Chinchilla-style tokens-per-parameter rules from language modeling do not transfer: every diffusion token is re-weighted across many noise levels rather than seen once, so there is no settled compute-optimal law, and scaling calls still ride empirical compute-quality curves.

Published efficient models agree this run is small. The PixArt-alpha paper and the Sana paper, both around 0.6B parameters, still report multi-GPU training budgets plus architectural efficiency work, staged training and linear attention, to reach their quality. A 210M run finishing in 3.5 days on one card is an outlier worth auditing, which is what the sections above did, not a template to copy blindly.

Pretrain, Fine-Tune, or Call an API

With the budget question settled, distribution distance rules the choice of when to pretrain vs fine-tune a diffusion model:

Option	Best when	Priors inherited	Iteration cost	Footprint and control
Pretrain from scratch	Data diverges sharply from web images (medical, industrial, game assets) or you need a compact controllable base	None, you build them	Each idea costs days	Small model you fully own, full access to internals
Fine-tune	Target sits near natural images (styles, products, characters)	Orders of magnitude more data, free	Hours per iteration	Larger base, license-bound
API	Quality ceiling dominates and volume is low	Whatever the vendor ships	None, per-image pricing	Zero footprint, zero control

The crossover rule in one line each. Pretrain below a billion parameters when your data diverges from web-scale natural images, when deployment needs a small controllable model, or when iteration economics favor owning the base, since every insight like the attention-sink finding only exists because the author owned the weights. Fine-tune when your target sits near natural images, because a pretrained latent model inherits priors built from vastly more data than any solo run can afford, and hours of LoRA beat weeks of from-scratch compute near the web distribution. Call an API when the quality ceiling matters and volume is low; at high volume, per-image pricing flips the math back toward ownership. Note the asymmetry in tooling too: deciding to train a text-to-image model from scratch means owning dataset, VAE, and recipe, while fine-tuning paths in Hugging Face diffusers are close to one command.

A Reproduction Checklist for Small DiT Runs

1. Fix samples seen before anything else. Batch × steps is your real budget; epochs are just samples ÷ dataset size.
2. Train in latent space through a frozen VAE. Roughly 256 tokens per image at 256² keeps the FLOP multiplier small; pixel space carries hundreds of times more tokens.
3. Treat patch size as a compute dial. Tokens per image set cost per sample; halving patch side length roughly quadruples the count.
4. Pick precision, then measure it. Profile sustained TFLOPS over a fixed step count on your actual build. The source reports torch.compile at 2.4× over eager; treat that as a target, not a promise.
5. Cadence the evals against a fixed prompt set. In the source run, loss moved only 0.805 to 0.754 while FID went 33.7 to 27.0 and FD-DINOv2 improved 570 to 218. The loss is a health signal, not a quality score.
6. Derive wall clock from measured samples per second, not step counts. At ~340 samples/s per card at this scale, 100M samples means about 3.5 days.
7. Publish the ledger: parameters, tokens per image, samples seen, sustained TFLOPS, GPU-hours, kWh. It is the difference between a claim and an auditable run.

Where These Numbers Can Break

Every figure above is only as good as the post it came from. The specific failure modes:

• Unique images versus samples seen. Multiply batch by steps before quoting anything. Here 102.4M ÷ 4.2M ≈ 24 reconciles with the author's 24-epoch remark, but verify against the write-up before citing either figure.
• Approximate token counts. Five aspect buckets mean counts vary around 256, and 6ND ignores register tokens and cross-attention, so true utilization sits slightly above the estimate.
• Precision changes the denominator. BF16 versus FP8 moves the dense ceiling by about 2×, so the same 109 TFLOPS reads as either a respectable or a poor utilization rate.
• Averages hide stalls. The 339 samples/s is a run average; dataloader stalls, bucket reshuffles, and eval passes all live inside it, and 3.5 days of wall clock need not equal 84 hours of pure stepping.
• Frozen overhead is still overhead. VAE and text encoding are cheap only if cached offline; computed online, they quietly eat wall clock outside the 6ND accounting.

The audit itself takes one minute on any published run: 6 × parameters × samples × tokens, divided by wall clock, divided by a dense datasheet ceiling. If implied utilization crosses 100%, some number is wrong. Training a diffusion model from scratch at small scale is now cheap enough that money is no longer the reason to do it or avoid it. Do it for the priors you cannot download, fine-tune for everything else, and rent the ceiling only while your volume stays low.