0Pricing
Learn AI with Python · Lesson

Mixed Precision Training with AMP

torch.cuda.amp.autocast(), GradScaler, FP16 vs BF16, memory savings, speed gains.

Mixed Precision Training with AMP is a free Learn AI with Python lesson on CoddyKit — lesson 3 of 4. You can read the complete lesson below for free — then practise it hands-on in the browser with a built-in code editor and a 24/7 AI tutor. It is part of the Learn AI with Python learning path, one of 4 lessons in the course, and your progress syncs across the web and the CoddyKit app.

What is Mixed Precision

Mixed precision training runs most operations in 16-bit floating point (FP16) instead of 32-bit (FP32). FP16 uses half the memory and runs faster on modern GPU tensor cores, while a few sensitive operations stay in FP32 for stability.

The Benefits

Mixed precision gives you:

  • ~2x memory savings, enabling larger batches or models
  • Faster matrix multiplications on tensor cores
  • Little to no loss in final model accuracy when done correctly

The FP16 Underflow Problem

FP16 has a narrow range. Small gradient values can underflow to zero, silently stopping learning. This is the central challenge mixed precision must solve, and the reason we cannot just cast everything to FP16.

torch.cuda.amp

PyTorch Automatic Mixed Precision (AMP) handles the details with two tools: autocast chooses precision per operation, and GradScaler prevents gradient underflow.

import torch
from torch.cuda.amp import autocast, GradScaler

The autocast Context

Wrap the forward pass in autocast(). Inside it, PyTorch automatically runs each op in the safest precision: matmuls in FP16, reductions like softmax and loss in FP32.

with autocast():
    output = model(x)
    loss = criterion(output, y)

Introducing GradScaler

GradScaler fights underflow by multiplying the loss by a large scale factor before backprop. This shifts tiny gradients into FP16 representable range; the scale is removed before the optimizer step.

scaler = GradScaler()

Scaling the Loss

scaler.scale(loss).backward() scales the loss up, then runs backprop. The resulting gradients are also scaled, keeping small values from vanishing.

scaler.scale(loss).backward()

scaler.step

scaler.step(optimizer) first unscales the gradients back to true magnitude, then calls optimizer.step(), but only if no infs or NaNs appeared. If the gradients overflowed, the step is skipped.

scaler.step(optimizer)

scaler.update

scaler.update() adjusts the scale factor for next iteration: it increases the scale when steps succeed and decreases it after an overflow. This adaptive scaling keeps training stable automatically.

scaler.update()

The Full AMP Loop

Putting the pieces together gives a complete mixed-precision training step.

scaler = GradScaler()
for x, y in loader:
    x, y = x.cuda(), y.cuda()
    optimizer.zero_grad()
    with autocast():
        out = model(x)
        loss = criterion(out, y)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

Practical Tips

Keep these in mind:

  • Only call backward() on the scaled loss
  • autocast wraps the forward pass only, not the backward or optimizer step
  • AMP shines on tensor-core GPUs; gains are smaller on older hardware

Quick Check

Test your AMP knowledge.

Recap

You learned mixed precision training with AMP:

  • autocast() picks FP16 or FP32 per operation in the forward pass
  • GradScaler scales the loss to avoid gradient underflow
  • The cycle is scaler.scale(loss).backward(), scaler.step(optimizer), scaler.update()
  • Result: roughly 2x memory savings and faster training

Frequently asked questions

Is the “Mixed Precision Training with AMP” lesson free?

Yes — the full text of “Mixed Precision Training with AMP” is free to read here on the web, and the Learn AI with Python course includes 4 lessons in total. To practise it interactively (a built-in code editor and a 24/7 AI tutor) and unlock the rest of the Learn AI with Python course, upgrade to CoddyKit PRO.

What will I learn in “Mixed Precision Training with AMP”?

torch.cuda.amp.autocast(), GradScaler, FP16 vs BF16, memory savings, speed gains. You practise Learn AI with Python with hands-on code you run directly in the browser, and a 24/7 AI tutor answers your questions as you work through the lesson.

Do I need any experience to start Learn AI with Python?

No prior experience is required. Learn AI with Python on CoddyKit is structured for beginners through advanced learners; this is — lesson 3 of 4, so you can start here or from the beginning and move at your own pace.

How long does the “Mixed Precision Training with AMP” lesson take?

Most CoddyKit lessons take about 5–10 minutes. Each one is bite-sized and interactive, so you make steady progress and pick up exactly where you left off across the web and the app.

Can I write and run code in this Learn AI with Python lesson?

Yes. Every Learn AI with Python lesson includes a built-in code editor, so you write and run real code right in your browser and get instant AI feedback — no local setup required.

All lessons in this course

  1. Multi-GPU Training with DataParallel
  2. DistributedDataParallel (DDP)
  3. Mixed Precision Training with AMP
  4. Efficient Training with Hugging Face Accelerate
← Back to Learn AI with Python