Training Faster

A large training run can cost thousands of dollars. So small efficiency gains can be worth real money. Three simple tricks are used in almost every serious project. They can double your speed, or let you train models that wouldn't otherwise fit into memory.

• Runs can cost thousands
• Small gains, real money
• Three simple tricks

Mixed precision

Numbers usually take thirty two bits. But most of a network's work doesn't need that much accuracy. Using sixteen bits needs half the memory, and on a G P U it's much faster. Autocast does this automatically, using smaller numbers where it's safe. Run it. Inside autocast, the output uses sixteen bits, but the weights stay in thirty two bits, so nothing's lost. This often gives a big speedup, with almost no change to your code.

import torch
from torch import nn

x16 = torch.tensor(1.0, dtype=torch.float16)
x32 = torch.tensor(1.0, dtype=torch.float32)
print("Bytes per number:", x16.element_size(), "vs", x32.element_size())

model = nn.Linear(10, 2)
inputs = torch.rand(4, 10)

with torch.autocast(device_type="cpu", dtype=torch.bfloat16):
    out = model(inputs)

print("Inside autocast:", out.dtype)
print("Weights stay:   ", model.weight.dtype)

Bytes per number: 2 vs 4
Inside autocast: torch.bfloat16
Weights stay:    torch.float32

Gradient accumulation

What if your G P U can only fit eight images at once, but you want a batch of sixty four? Run several small batches, adding up the gradients, and only update after all of them. We divide the loss by the number of steps, to keep the scale right. Run it. Same memory as eight, and the same effect as a batch of sixty four. Large language models are often trained in this way.

import torch
from torch import nn

model = nn.Linear(8, 1)
opt = torch.optim.SGD(model.parameters(), lr=0.1)
x, y = torch.rand(64, 8), torch.rand(64, 1)
steps = 8

opt.zero_grad()
for i in range(steps):
    loss = ((model(x[i * 8:(i + 1) * 8]) - y[i * 8:(i + 1) * 8]) ** 2).mean() / steps
    loss.backward()
opt.step()

print("Real batch size:", 8)
print("Effective batch size:", 8 * steps)
print("One weight update done, memory stayed small")

Real batch size: 8
Effective batch size: 64
One weight update done, memory stayed small

Gradient clipping

Sometimes one bad batch produces a huge gradient, and a single step can ruin hours of training. Clipping sets a speed limit. If the gradient is too large, it's scaled down, keeping its direction. Run it. A huge gradient gets reduced to exactly one. Almost every transformer uses this, usually with a limit of about one.

import torch
from torch import nn

model = nn.Linear(4, 1)
x = torch.rand(4, 4) * 1000
loss = (model(x) ** 2).mean()
loss.backward()

before = model.weight.grad.norm().item()
nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
after = model.weight.grad.norm().item()

print(f"Gradient size before: {before:,.1f}")
print(f"Gradient size after:  {after:.3f}")

Gradient size before: 634,175.4
Gradient size after:  1.000

Summary