Mixed Precision

Your model doesn't quite fit on your graphics card. You could rent a bigger one, at several times the price. Or you could use smaller numbers for most of the maths. Training gets faster, memory use falls, and the final accuracy usually stays the same.

• The model doesn't quite fit
• Rent a bigger card?
• Or use smaller numbers

The arithmetic

Half a million parameters take about two megabytes at thirty two bits, and half that at sixteen. Autocast chooses the right precision for each operation automatically. On a C P U, we use a format called B float sixteen. Run it. The memory is halved, and the output type confirms that it's using the smaller format. On a graphics card, you'd pair that with a gradient scaler, which stops very small gradients from rounding down to zero.

import time, torch

model = torch.nn.Sequential(torch.nn.Linear(512, 512), torch.nn.ReLU(), torch.nn.Linear(512, 512))
x = torch.randn(256, 512)
params = sum(p.numel() for p in model.parameters())

print(f"Parameters: {params:,}")
print(f"32-bit memory: {params * 4 / 1e6:.1f} MB")
print(f"16-bit memory: {params * 2 / 1e6:.1f} MB")

with torch.autocast("cpu", dtype=torch.bfloat16):
    print("Output type:", model(x).dtype)

Parameters: 525,312
32-bit memory: 2.1 MB
16-bit memory: 1.1 MB
Output type: torch.bfloat16

Summary