🎓 A1Learner
🔍

PyTorch Autocast on CPU

Autocast is a context manager that runs eligible operations in a lower precision dtype. On CPU that dtype is bfloat16, a 16-bit float with less precision and the same range as float32.

Plain evaluation runs in float32

A Linear layer without autocast takes float32 in and returns float32. Every operation pays full precision cost.

import torch

torch.manual_seed(0)
layer = torch.nn.Linear(8, 8)
x = torch.randn(2, 8)
print(layer(x).dtype)

Autocast switches the Linear output

Inside the autocast block, the same layer returns bfloat16. The layer itself stays stored in float32.

with torch.autocast("cpu"):
    out = layer(x)
print(out.dtype)

Some operations stay in float32

Autocast keeps a short list of operations in float32 for accuracy. Softmax and mean are on that list.

with torch.autocast("cpu"):
    probs = torch.softmax(x, dim=1)
    centre = x.mean()
print(probs.dtype, centre.dtype)

The answers still agree

Lower precision changes the last digits of the scores. The predicted class for each row stays the same here.

with torch.no_grad():
    plain = layer(x)
    with torch.autocast("cpu"):
        auto = layer(x)
print(torch.equal(plain.argmax(dim=1), auto.argmax(dim=1)))

Example

import time

import torch

torch.manual_seed(0)
model = torch.nn.Sequential(
    torch.nn.Linear(256, 256),
    torch.nn.ReLU(),
    torch.nn.Linear(256, 256),
    torch.nn.ReLU(),
    torch.nn.Linear(256, 10),
)
model.eval()
x = torch.randn(128, 256)


def timed(use_autocast, reps=20):
    with torch.no_grad():
        for _ in range(3):
            if use_autocast:
                with torch.autocast("cpu"):
                    model(x)
            else:
                model(x)
        start = time.perf_counter()
        out = None
        for _ in range(reps):
            if use_autocast:
                with torch.autocast("cpu"):
                    out = model(x)
            else:
                out = model(x)
        ms = (time.perf_counter() - start) * 1000 / reps
    return out, ms


plain, plain_ms = timed(False)
auto, auto_ms = timed(True)
print("Plain output shape:", plain.shape)
print("Autocast output dtype:", auto.dtype)
print("Predictions agree:", torch.equal(plain.argmax(dim=1), auto.argmax(dim=1)))
print("Outputs close:", torch.allclose(plain, auto.float(), atol=0.05))
print(f"Plain time per run: {plain_ms:.0f} ms")
print(f"Autocast time per run: {auto_ms:.0f} ms")

Key points

  • Autocast: a block that picks a cheaper dtype for eligible operations
  • bfloat16: a 16-bit float with float32 range and fewer fraction bits
  • Use torch.autocast("cpu") around a forward pass to enable it
  • Linear and matmul operations run in bfloat16 under CPU autocast
  • Softmax and mean stay in float32 under CPU autocast
  • Predictions under autocast matched the plain predictions in the closing example
  • In short, autocast lets heavy operations run cheaper while sensitive ones stay precise

In this example, a three-layer MLP scores a batch of 128 rows. The timed helper runs each version three times to warm up, then averages 20 runs. The plain version returns float32 scores with shape 128 by 10. The autocast version returns bfloat16 scores. The predicted classes agree on every row, and the scores compare close within a small tolerance. The two times vary on every run, so treat them as observations, not promises. On the recorded run, autocast was the slower of the two. This small model gives bfloat16 too little work to repay its conversion cost on this machine.

Note

On this machine the small model in the closing example ran slower under autocast. A gain is not guaranteed.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Under torch.autocast("cpu"), which dtype does a Linear layer use for its maths?