5 Lines of Code to Make Your LLM Train 3x Faster on GPU

The takeaway first

TL;DR

LLMs are hard to train, primarily due to the compute involved. However we can use some tactics to achieve on-par performance with same compute budget with a few modifications.

These modfications and simple yet powerful exploits to fully utilize the capabilities of modern GPUs can significantly reduce training time. Most of it comes down to learning how GPU works.

We apply five one-line changes to a 124M parameter GPT-2 Model's training loop and reduce training time per step from 250 ms to 80 ms on a single RTX 5070 making training 3x faster

Milliseconds per training step on an RTX 5070, measured after each change

Baseline

Setup

• Model: GPT-2 124M with 12 layers, 12 heads, 768 embedding size and GPT-2 initialization.
• Data: Tiny Shakespeare (1M characters), ~338k tokens
• Hardware: one RTX 5070 (8 GB). Windows, PyTorch 2.x. Windows needs triton -> pip install -U "triton-windows>=3.8,<3.9".
• Metric: wall-clock milliseconds per training step

t0 = time.time()
optimizer.zero_grad()
logits, loss = model(x.to(device), y.to(device))
loss.backward()
optimizer.step()
torch.cuda.synchronize()  # wait for the GPU, not just the CPU's queue
t1 = time.time()
print(f"{(t1 - t0) * 1000:.2f} ms")

Optimization 1 · Precision

TF32: a cheaper multiply, same float32 tensors

Most of a training step is matrix multiplication, and on Ampere and newer GPUs that multiply can run at lower precision without our code ever seeing a lower-precision number.

TF32 is a mode of the tensor core, the hardware unit that multiplies small blocks of matrices. Inside each multiply it keeps float32's sign and 8 exponent bits but cuts the mantissa from 23 bits to 10. The inputs arrive as float32, the running sum stays float32, and the output is float32. The precision is dropped only inside the instruction.

• Nothing in our code changes type. logits.dtype still prints torch.float32.
• Only matrix multiplications get faster, and in GPT-2 they are most of the work. The largest is the classifier that turns 768 channels into 50,257 vocabulary scores.
• The results are slightly less exact, but the training loss looks no different.

On paper the payoff is big. The A100 datasheet lists 19.5 TFLOPS for float32 and 156 for TF32, an 8× ceiling.

# 'highest' = float32, 'high' = TF32, 'medium' = bfloat16
torch.set_float32_matmul_precision('high')

If moving bytes is the bottleneck, the next step is to move fewer of them.

Optimization 2 · Precision

bfloat16 autocast: half the bytes per activation

Mixed precision goes after the bottleneck TF32 exposed. The forward pass stores activations in 16 bits instead of 32, so every matmul input is half as much data to move. The weights stay float32 so the optimizer updates them at full precision.

bfloat16 and float16 both use 16 bits, but they spend them differently:

• bfloat16 keeps all 8 exponent bits, so it covers the same range of values as float32 and gives up precision instead (7 mantissa bits). Values don't overflow or round to zero, so nothing extra is needed.
• float16 keeps more precision but shrinks the range, so small gradients round to zero. Training in float16 needs a GradScaler, one more piece of state to manage. bfloat16 avoids it.
• torch.autocast chooses op by op. Matrix multiplies run in bfloat16. Softmax, layer norm and the loss stay in float32 because small errors there distort the result.

with torch.autocast(device_type=device, dtype=torch.bfloat16):
    logits, loss = model(x.to(device), y.to(device))
loss.backward()      # outside
optimizer.step()     # outside

Inside the block, logits.dtype is torch.bfloat16 while model.transformer.wte.weight.dtype is still torch.float32. Activations in low precision and weights in full precision is what "mixed" refers to.

Precision is now as low as it can safely go. What's left is how many times the data travels.

Optimization 3 · Memory traffic

torch.compile: remove the interpreter, fuse the kernels

torch.compile reads the whole network before running it and rewrites it into fewer, larger GPU programs, so data is read from memory once and used many times.

It is a compiler for neural networks, much as GCC is a compiler for C. It removes two costs:

1. The Python interpreter. In eager mode Python walks the model line by line and launches each operation without knowing what comes next. Compiled, the network runs as a single unit with no Python in between.
2. Repeated trips to memory. The GELU, 0.5·x·(1 + tanh(√(2/π)·(x + 0.044715·x³))), is about eight small operations. Eager PyTorch launches eight kernels, and each one reads from GPU memory and writes its result back. The compiled version reads x once, does all eight operations while the data is on the chip, and writes once. This is called kernel fusion.

Kernel fusion keeps the data on the chip. The arithmetic is the same; what disappears is the round trips to memory.

Fusion matters this much because of where the data lives. HBM, the GPU's main memory, holds tens of gigabytes but sits off the compute chip, so every read and write is slow. On-chip memory (caches and registers) holds only tens of megabytes but is very fast. A step that goes back to HBM after every operation spends most of its time moving data.

model = GPT(config).to(device)
model = torch.compile(model)

Optimization 4 · Memory traffic

Flash attention: a fusion the compiler cannot find

The attention matrix is the largest temporary tensor in the step. Flash attention computes the same output without ever writing that T×T matrix to memory.

Our attention was four lines:

att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
att = att.masked_fill(self.bias[:, :, :T, :T] == 0, float('-inf'))
att = F.softmax(att, dim=-1)
y = att @ v

At T = 1024, att holds about a million numbers for every head of every sequence in the batch. Each of the four lines writes that matrix to HBM, and the next line reads it back. Flash attention does all four steps in one kernel, a block at a time, so the full matrix never exists in memory.

This works because of the online softmax. A normal softmax needs the whole row before it can normalize. The online version keeps a running maximum and a running sum and corrects them as each new block arrives, so the row can be handled in pieces. Flash attention does more arithmetic than the four-line version and is still faster (the paper reports up to 7.6×), because arithmetic was never the expensive part.

torch.compile can't find this. It fuses the operations we wrote; it doesn't know that softmax can be reorganized.

# is_causal=True applies the autoregressive mask
y = F.scaled_dot_product_attention(q, k, v, is_causal=True)

Optimization 5 · Nice Numbers

Padding the vocabulary: Nice numbers for CUDA

GPU kernels cut matrices into tiles whose sides are powers of two: 16, 32, 64, 128. When a dimension doesn't divide evenly, the leftover goes through a separate, slower pass. The cheapest remaining speedup is to stop handing the GPU awkward numbers.

• 50,257, GPT-2's vocabulary size, is odd, so it doesn't split into power-of-two tiles. It also sets the size of the two largest matrices in the model: the token embedding and the output classifier.
• 50,304 is the nearest multiple of 128, so it divides by 8, 16, 32, 64 and 128.
• The tokenizer never produces the 47 extra tokens, so their embedding rows are never used. The classifier learns to give them near-zero probability, as it already does for the thousands of real tokens that never appear in Shakespeare. Nothing breaks.

@dataclass
class GPTconfig:
    vocab_size: int = 50304   # was 50257; nearest multiple of 128
    block_size: int = 4096
    n_embd: int = 768
    n_layers: int = 12
    n_head: int = 12

Results

Where the 250 ms went

Per-step latency after each optimization, RTX 5070. Green bars are the memory-traffic changes.

The 170 ms saved splits evenly: 85 ms from the two precision changes and 85 ms from the three memory-traffic changes. They behave differently, though. Precision gave one big jump with TF32 and then stalled. The memory-traffic changes kept paying off, with torch.compile alone saving 70 ms.

Karpathy's A100 went from 1000 ms to 93 ms, about 11× against our 3.1×. His float32 baseline started much slower, and the A100's roughly 2 TB/s of memory bandwidth leaves more room for each change to help. The pattern is the same on both cards: TF32 falls well short of its datasheet number, autocast adds little on top, and torch.compile is a big jump.

Wrapping up

Conclusions

1. Is the GPU Ampere or newer? If yes, turn on TF32 and bfloat16 autocast. Together they are two lines and cost no accuracy we can measure. If not, go straight to compile.
2. Is the model running in eager PyTorch? If yes, torch.compile is the biggest single gain available, and flash attention is the one fusion it won't find on its own.

Related reading

GPT-2 · Paper Notes

GPT-2: When Next-Token Prediction Becomes Multitask Learning

The architecture, data and tokenizer behind the model this post trains.