🎓 A1Learner
🔍
0 of 71 lessons completed

PyTorch Precision and Loss Scaling

Precision is the number of bits a float dtype uses to store each value. float64 spends 64 bits per value, and float32 spends 32.

One value, two precisions

The same numbers can live in either dtype. Only the storage width changes.

import time

import torch

x32 = torch.tensor([0.1, 1.7, 3.3], dtype=torch.float32)
x64 = x32.to(torch.float64)
print("Same values, two dtypes:", x32.dtype, x64.dtype)

Wider costs more time

A large matmul runs in both precisions here. The float64 version pays for its extra bits on every value.

torch.manual_seed(0)
a32 = torch.randn(1024, 1024, dtype=torch.float32)
b32 = torch.randn(1024, 1024, dtype=torch.float32)
a64 = a32.to(torch.float64)
b64 = b32.to(torch.float64)


def timed(a, b, reps=10):
    for _ in range(2):
        a @ b
    start = time.perf_counter()
    for _ in range(reps):
        a @ b
    return (time.perf_counter() - start) * 1000 / reps


print(f"float32 matmul: {timed(a32, b32):.0f} ms")
print(f"float64 matmul: {timed(a64, b64):.0f} ms")

Wider costs more memory

element_size reports the bytes one value occupies. The wider dtype takes exactly twice the room.

print("Bytes per value:", x32.element_size(), x64.element_size())

Example

import torch


def gradient(dtype):
    x = torch.tensor([0.1, 1.7, 3.3], dtype=dtype, requires_grad=True)
    y = (x ** 7).sum()
    y.backward()
    return [round(v, 4) for v in x.grad.tolist()]


grad32 = gradient(torch.float32)
grad64 = gradient(torch.float64)

print("float32 gradient:", grad32)
print("float64 gradient:", grad64)
print("Gradients exactly equal:", grad32 == grad64)

Key points

  • float32 is PyTorch's default dtype for model maths
  • float64 doubles the memory per value, at 8 bytes against 4
  • Use .to(torch.float64) to cast a tensor to the wider dtype
  • Gradients in float32 can differ from float64 in the last printed digits
  • Careful gradient checks run in float64 for its exactness
  • Loss scaling multiplies the loss by a large factor
  • The factor stops small gradients rounding to zero in float16, a 16-bit float
  • In short, wider precision buys exactness and pays for it in time and memory
  • Loss scaling belongs to GPU float16 training

In this example, one gradient is computed twice, once per precision. The function raises each value to the seventh power and sums the results. Backward fills the gradient, which is seven times each value to the sixth power. Both lists agree until the last value. There float32 prints 9040.2754 and float64 prints 9040.2758. The difference is pure rounding, and it grows with the size of the numbers. float64 is the reference when an exact gradient matters, and float32 is the default because it is cheaper.

Note

CPU autocast uses bfloat16, whose wide range makes a scaler unnecessary on this machine.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Why do gradient checks run in float64?