GPU and Mixed Precision

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 18 / 80 · Module M03 · ~2 hours

0 · Precision is an operation policy

Mixed precision is a policy for operations, not an instruction to halve every tensor. A narrower exponent range can erase gradients or overflow activations; fewer significand bits can hide small updates even when every value remains finite. Device placement, numerical correctness and throughput are three different questions.

BF16 did not make overflow impossible, CUDA autocast did not universally switch its default to BF16, and a datatype alone does not establish an end-to-end speedup. The historical claims about one-third of pretraining runs failing and a named executive's debugging quote had no supporting primary passage and are not used here. The technical lesson survives without them.

Prerequisites

Read the training loop and understand tensors, registered parameters/buffers and a scalar loss. The CPU demonstrations need PyTorch 2.9 only. GPU references below are complete synthetic programs, explicitly unexecuted, and are not instructions to acquire a GPU or download models.

1 · Getting to CUDA — the two-line ritual

CPU and GPU are separate memory/execution domains. module.to(device) modifies registered parameters and buffers and returns the same module; an ordinary tensor attribute is not a registered buffer. tensor.to(...) returns a converted tensor if conversion is needed and may return the identical object when device/dtype already match. Assign its result rather than assuming mutation. Move a model before constructing its optimizer so the optimizer references the intended parameters.

Inspect the operands of the failing operation, not just next(model.parameters()).device. Inputs, targets, masks and temporary tensors must be compatible with that operation. torch.arange(T, device=x.device) normally remains an integer index; boolean masks should stay boolean. The old advice to give every tensor dtype=x.dtype is wrong for both.

import torch
from torch import nn
 
class WithBuffer(nn.Module):
    def __init__(self):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(2))
        self.register_buffer("scale", torch.ones(2))
        self.unregistered = torch.ones(2)
    def forward(self, x):
        return x*self.weight*self.scale
 
m = WithBuffer()
assert m.to(dtype=torch.float64) is m
assert m.weight.dtype == m.scale.dtype == torch.float64
assert m.unregistered.dtype == torch.float32
x = torch.ones(3)
assert x.to("cpu") is x
assert x.to(torch.float64) is not x and x.dtype == torch.float32
indices = torch.arange(3, device=x.device)
assert indices.dtype == torch.int64
print("module/tensor conversion and buffer contract passed on CPU")

2 · The pin-memory and nonblocking transfer contract

Page-locked host memory permits DMA without paging the source out. DataLoader's pin_memory=True can stage supported tensor batches in pinned memory; custom batch classes need their own pinning method. non_blocking=True requests a copy asynchronous with respect to the host where possible. The Tensor API documents attempts for pageable as well as pinned sources; the claim that it can only help pinned input is too strong.

Host asynchrony is not GPU copy/compute overlap. Commands on one stream are ordered. To overlap the next batch's host-to-device copy with current GPU work requires pinned input, a free DMA engine and a separate stream plus correct event/stream dependencies and tensor lifetimes. Do not mutate a pinned source until its copy finishes. For a device-to-host nonblocking copy, synchronize before reading the CPU result. Pinning immediately before every copy can cost more than it saves; measure the complete input pipeline.

This diagram describes dependencies, not a measured overlap trace. The GPU reference uses safe same-stream copies; it does not claim that the two flags alone create overlap or recover a fixed 20–40% throughput loss.

3 · fp16 vs bf16 — range and resolution

One more bit is the sign. Normalized numbers also have an implicit leading significand bit. FP16's smallest subnormal is $2^{-24}$, about 5.96e-8; values below its smallest normal are not all immediately zero. Backend arithmetic may flush subnormals. BF16 shares FP32's exponent-bit count, not every exact maximum/rounding behavior.

Near one, FP16 spacing is $2^{-10}$ and BF16 spacing is $2^{-7}$. Conversion normally rounds; it does not simply chop bits without ever changing an exponent. BF16 still overflows and underflows, and finite rounding errors can be consequential. Prefer supported BF16 for a workload that passes numerical/task-quality checks, use FP16 plus scaling when appropriate, and keep an FP32 reference while debugging. Query actual backend support rather than treating a marketing family name as the full API contract.

import torch
for dtype in [torch.float32, torch.float16, torch.bfloat16]:
    f = torch.finfo(dtype)
    print(str(dtype), "max", f.max, "tiny", f.tiny, "eps", f.eps)
assert torch.tensor(1e-8).half().item() == 0
assert torch.tensor(2.**-24).half().item() > 0  # subnormal, CPU cast
assert torch.isinf(torch.tensor(70000.).half())
assert torch.isfinite(torch.tensor(70000.).bfloat16())
assert torch.isinf(torch.tensor(1e39, dtype=torch.float64).bfloat16())
assert torch.tensor(1.001).bfloat16().float().item() == 1.
assert torch.tensor(1.001).half().float().item() != 1.
# A small additive epsilon can disappear; disabling autocast doesn't recast inputs.
var = torch.tensor(0., dtype=torch.float16)
assert torch.isinf(1/(var+1e-8))
safe = 1/(var.float()+1e-8)
assert torch.isfinite(safe)
print("range, subnormal, rounding and FP32-island controls passed")

4 · AMP with bf16 — dispatch, not a global cast

Autocast chooses dtypes for eligible operations; with ordinary FP32 model parameters, these remain FP32. Do not call model.half() or cast all input batches to BF16 merely because autocast is enabled. FP32 inputs are expected and eligible matrix multiplications can still be autocast. The old explanation that FP32 input necessarily forced the first autocast matmul onto a slow fallback was incorrect.

Autocast belongs around forward and loss computation, not backward. Backward kernels use the relevant forward-operation dtypes; leaf gradients for FP32 parameters are generally FP32, although intermediate low-precision gradients can still underflow. Not every reduction, BatchNorm call or unlisted operation is automatically FP32. Consult the device-specific op list. In-place operations, explicit out= and explicit dtype arguments have different eligibility.

A CPU BF16 test can demonstrate dispatch and finite gradients, not CUDA Tensor Core speed or model-quality equivalence:

import torch
from torch import nn
import torch.nn.functional as F
 
torch.manual_seed(8)
model = nn.Linear(4, 3)
x = torch.randn(8, 4); y = torch.arange(8) % 3
with torch.autocast("cpu", dtype=torch.bfloat16):
    logits = model(x)
    loss = F.cross_entropy(logits, y)
assert logits.dtype == torch.bfloat16 and loss.dtype == torch.float32
loss.backward()
assert model.weight.dtype == model.weight.grad.dtype == torch.float32
assert torch.isfinite(model.weight.grad).all()
# The float32 reference is compared, not assumed identical.
ref = F.cross_entropy(model(x), y)
assert abs(loss.item()-ref.item()) < .02
print("CPU BF16 autocast", loss.item(), "FP32 reference", ref.item())

5 · AMP with fp16 and GradScaler

Loss scaling primarily protects small intermediate gradients from underflow. It does not fix an already overflowing forward pass. By the chain rule, $\nabla(SL)=S\nabla L$ in real arithmetic. Powers-of-two scales are convenient, but range/rounding still matter; arbitrary multiplication is not guaranteed exact.

scaler.scale(loss).backward() scales the backward signal. scaler.step(optimizer) unscales if needed, checks for nonfinite gradients and skips that optimizer's update if they were detected. scaler.update() adjusts the scale after the attempted update. Default growth/backoff are configurable; the scale can fall below one. No automatic warning at every scale decrease is promised. Repeated decreases are evidence to inspect, not proof the model is broken solely because the scale reached one.

Use torch.amp.GradScaler("cuda") and torch.autocast("cuda", dtype=...) under the 2.9 API rather than deprecated torch.cuda.amp aliases. BF16 training usually needs no scaler because of its wider range, not because all BF16 numerics are safe.

6 · Gradient clipping under AMP — the order matters

Accumulate at a constant scale for the whole effective batch, unscale once, then inspect/clip, attempt the optimizer update and update the scaler. Clipping a scaled vector $Sg$ to threshold $c$ and then dividing by $S$ produces a norm near $c/S$, not $c$, whenever the clip triggers.

import torch
from torch import nn
 
g = torch.tensor([3., 4.], dtype=torch.float64)
c, scale = 1., 1024.
p = nn.Parameter(torch.zeros(2, dtype=torch.float64)); p.grad = g.clone()
reported = torch.nn.utils.clip_grad_norm_([p], c)
assert abs(reported.item()-5.) < 1e-12
assert torch.allclose(p.grad, torch.tensor([.6,.8], dtype=torch.float64), atol=1e-6)
q = nn.Parameter(torch.zeros(2, dtype=torch.float64)); q.grad = scale*g
_ = torch.nn.utils.clip_grad_norm_([q], c); q.grad /= scale
assert q.grad.norm() < .001 and p.grad.norm() > .99
# Execute scaler skip semantics on CPU, without pretending this is FP16 CUDA.
r = nn.Parameter(torch.tensor(1.)); opt = torch.optim.SGD([r], lr=.1)
s = torch.amp.GradScaler("cpu", init_scale=8.)
opt.zero_grad(); s.scale(r*r).backward(); s.unscale_(opt)
assert torch.allclose(r.grad, torch.tensor(2.))
s.step(opt); s.update(); assert torch.allclose(r, torch.tensor(.8))
opt.zero_grad(); s.scale(r*r).backward(); r.grad.fill_(float("inf"))
before = r.detach().clone(); old_scale = s.get_scale()
s.step(opt); s.update()
assert torch.equal(r, before) and s.get_scale() < old_scale
print("unscale-before-clip and CPU scaler nonfinite-skip controls passed")

6.1 Complete GPU reference — unexecuted

This single-GPU synthetic program compares FP32, FP16/scaler, supported BF16 and BF16+compile with matching initial state and data. It warms each path, resets initial state and RNG, then times synchronized full training steps and reports peak allocated/reserved memory plus clean FP32 validation loss. Compile/cold-start time is deliberately excluded from the warmed metric; report it separately for short jobs. torch.compile is a choice, not a guarantee that all operations fuse.

The reference avoids network/dataset downloads. No output from this program is claimed in this review. The CPU program above is the executed precision evidence.

# GPU reference — unexecuted
import copy
import time
import torch
from torch import nn
import torch.nn.functional as F
 
def benchmark(mode, steps=100, float8=False):
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA hardware/runtime required; not a CPU benchmark")
    dtype = {"fp32":None, "fp16":torch.float16,
             "bf16":torch.bfloat16, "bf16_compile":torch.bfloat16}[mode]
    if dtype == torch.bfloat16 and not torch.cuda.is_bf16_supported():
        raise RuntimeError("BF16 unsupported on this device/runtime")
    torch.manual_seed(12)
    base = nn.Sequential(nn.Linear(256, 512), nn.GELU(), nn.Linear(512, 256)).cuda()
    if float8:
        # Optional torchao v0.14.1 reference; not installed by this lesson.
        from torchao.float8 import convert_to_float8_training
        base = convert_to_float8_training(base)
    initial = copy.deepcopy(base.state_dict())
    run = torch.compile(base) if mode.endswith("compile") or float8 else base
    cpu_x = torch.randn(64,256).pin_memory()
    cpu_y = (torch.arange(64) % 256).pin_memory()
    vx, vy = torch.randn(64,256,device="cuda"), torch.arange(64,device="cuda") % 256
    def execute(n):
        opt = torch.optim.AdamW(base.parameters(), lr=1e-3)
        sched = torch.optim.lr_scheduler.StepLR(opt, step_size=50, gamma=.5)
        scaler = torch.amp.GradScaler("cuda", enabled=(dtype == torch.float16))
        accepted = 0; base.train()
        for _ in range(n):
            opt.zero_grad(set_to_none=True)
            x = cpu_x.to("cuda", non_blocking=True)
            y = cpu_y.to("cuda", non_blocking=True)
            with torch.autocast("cuda", dtype=dtype or torch.float16,
                                enabled=(dtype is not None)):
                loss = F.cross_entropy(run(x), y)
            if not torch.isfinite(loss): raise FloatingPointError("forward loss")
            scaler.scale(loss).backward(); scaler.unscale_(opt)
            finite = all(p.grad is None or torch.isfinite(p.grad).all()
                         for p in base.parameters())
            if finite: torch.nn.utils.clip_grad_norm_(base.parameters(), 1.)
            elif not scaler.is_enabled(): raise FloatingPointError("gradient")
            previous = scaler.get_scale(); scaler.step(opt); scaler.update()
            # Single optimizer, ordinary GradScaler growth/backoff policy only.
            if scaler.get_scale() >= previous:
                sched.step(); accepted += 1
        return accepted
    execute(5); torch.cuda.synchronize()
    base.load_state_dict(initial); torch.manual_seed(12)
    torch.cuda.reset_peak_memory_stats()
    start = time.perf_counter(); accepted = execute(steps)
    torch.cuda.synchronize(); seconds = time.perf_counter()-start
    peak = torch.cuda.max_memory_allocated(); reserved = torch.cuda.max_memory_reserved()
    base.eval()
    with torch.no_grad(): val = F.cross_entropy(base(vx), vy).item()
    return dict(mode=mode, float8=float8, seconds=seconds, updates=accepted,
                peak_allocated=peak, peak_reserved=reserved, validation=val,
                torch=torch.__version__, gpu=torch.cuda.get_device_name())
 
if __name__ == "__main__":
    for mode in ["fp32", "fp16", "bf16", "bf16_compile"]:
        if mode.startswith("bf16") and not torch.cuda.is_bf16_supported(): continue
        print(benchmark(mode))
    # Optional on a verified supported torchao/FP8 stack:
    # print(benchmark("bf16_compile", float8=True))

Validation here is a numerical comparison on synthetic labels, not evidence about downstream quality. Peak memory includes optimizer state and workspace, not just activations. Reinitialize/snapshot stochastic data order and all states for real paired comparisons; report repetitions, shape, batch size, precision flags, device, software and compile amortization. FP32 hardware may itself use TF32 according to backend settings, which also belong in the report.

6.2 Compiler boundaries and FP32 islands

Compile the callable once; enter autocast for its invocations. Shape/control-flow guards can cause graph breaks or recompilation. There is no universal PyTorch-version threshold that makes every model correct under every compiler/backend/AMP combination. Debug eager FP32 first, compare outputs and gradients, then restore precision and compilation separately.

For a sensitive reciprocal with a low-precision input, autocast(enabled=False) alone leaves the input low precision. Explicitly convert: safe = 1/(var.float()+eps). Choose epsilon from the intended mathematical scale; increasing it changes the function. Casting an infinite result to FP32 after overflow does not recover the lost number.

6.3 FP8 — scaled tensors and specialized kernels

E4M3 and E5M2 allocate different exponent/fraction budgets; concrete FP8 variants also differ in infinity/NaN conventions. A common recipe uses E4M3 for weights/activations and E5M2 for gradients, but it is not mandatory for every library/model. Scaled quantization chooses a scale from tensor/block/tile statistics, rounds/clips to a representable FP8 value, and uses higher-precision accumulation or sensitive paths. Merely storing torch.float8_* tensors does not supply a complete training stack.

The pinned torchao v0.14.1 convert_to_float8_training returns a module with eligible Linear layers replaced; it is not torchao.convert_to_float8_training at the package root. Preserve the returned root module, check dimensional/backend requirements and use the library's compatible recipe. Compilation can be important for efficient kernels but does not follow from the datatype definition. Device support is broader and more specific than 'only H100' or 'no consumer GPU ever has FP8'.

DeepSeek-V3's report describes mixed FP8 training, fine-grained activation/weight scaling and higher-precision exceptions. 'The whole model was trained entirely in FP8' omits those exceptions. The report is not proof of identical Llama-family recipes, exact 2× end-to-end speed or half of total optimizer memory. Read the actual model report before attributing a precision policy.

6.4 DeviceMesh and FSDP2 preview

A DeviceMesh names groups in a logical multidimensional device topology. Data-parallel, tensor-parallel and pipeline dimensions partition work differently; a dimension name is not a promise of physical bandwidth. FSDP2 fully_shard shards parameters/state and materializes needed parameters for computation; tensor parallelism partitions layer operations. All ranks must construct consistent meshes and use the right submesh for each operation. A mesh alone does not train a giant model or eliminate communication. This is architectural vocabulary, not an executed distributed recipe.

7 · Failure cases: locate the first nonfinite value

• Device mismatch: inspect input, labels, parameters, registered buffers and temporaries. Use a device-aware integer arange for positions, not a floating index cast.
• Tiny epsilon: the CPU control shows 1e-8 disappearing in FP16 at zero variance. A representable epsilon may still yield an overflowing reciprocal; use a justified FP32 island and verify both forward and backward.
• BF16 still produces NaN: check division by zero, invalid logs, squared overflow and data. Its wide range is not a mathematical-domain repair.
• No speed gain: profile shape, arithmetic intensity, data loading, synchronization, kernel support and compile overhead. Do not blame FP32 input merely because autocast receives it.
• Frozen updates with finite loss: inspect underflow and unscale/clip order. The wrong-order control loses a factor of 1024 in the example; a finite logged loss cannot detect that alone.

9 · Retention scaffold — worked answers

Q1 — What mutates in .to()? A module converts registered state and returns itself. A tensor conversion returns a result; no-op conversions can return the same tensor. Unregistered attributes do not follow module conversion.

Q2 — What makes transfer asynchronous? non_blocking=True requests host-asynchronous behavior, not overlap by itself. Reliable copy/compute overlap additionally needs pinned source, available DMA hardware, a separate stream and correct dependencies/lifetimes. Same-stream ordering remains.

Q3 — Compare FP16/BF16 and scaling. FP16 has 5 exponent/10 fraction bits; BF16 has 8/7. Maxima are 65504 and about 3.390e38. Scaling primarily rescues FP16's tiny intermediate gradients; BF16 usually avoids that need but can still overflow/underflow and has coarser resolution.

Q4 — What does scaler.step add? It unscales if not already done and skips that optimizer's update when nonfinite gradients were detected. Update the scale and do not advance an optimizer-update clock for a skipped step. An optimizer's return value is not a reliable accepted-step boolean.

Q5 — Why unscale before clipping? For $g=(3,4)$, $c=1$ and $S=1024$, correct clipping returns a gradient near $(.6,.8)$; wrong-order clipping then unscaling yields that divided by 1024. Both can look finite while only one represents the intended threshold.

Q6 — How should the five-way GPU benchmark be done? Use the complete reference with matching initial state, fixed data, warmup, synchronization and recorded peak memory. Compare FP32, FP16/scaler, supported BF16, BF16+compile and optional supported FP8. Report skipped updates and both cold/warmed costs; do not prefill a speedup or equal-quality answer. GPU/FP8 runs remain unexecuted here.

Sources and execution boundary

PyTorch 2.9 AMP, AMP examples, official transfer tutorial, FP8 format paper, torchao v0.14.1 conversion source, and DeepSeek-V3 v2 support these bounded mechanisms. CPU casts/autocast/scaler controls are executed. CUDA transfer overlap, Tensor Core/FP8/compiled throughput, distributed operation and pretrained-model quality are not reproduced.

Optional pre-read resources

These inherited resource links are optional, not the primary evidence for this review. Video identity, duration and current availability are not certified.

• Mixed Precision Training in PyTorch
• How GPUs Are Made and Why They're So Fast
• Automatic Mixed Precision examples
• NVIDIA Mixed Precision Training Guide

Previous: ← DL S017 · Training Loop · Next: DL S019 · Debugging →