GPU Training and Mixed Precision

Build the mental model

Neural network training is dominated by matrix multiplications inside forward and backward passes, and GPUs contain thousands of small cores built exactly for that kind of massively parallel arithmetic, so moving computation there can cut training time by an order of magnitude on real workloads. The mechanism is .to(device): it returns a copy of a tensor or model living on the target device, and every tensor participating in an operation must live on the same device — mixing a GPU model with CPU data raises a runtime device-mismatch error, since PyTorch never silently transfers data across devices for you. Beyond just using the GPU, mixed precision training pushes further: most operations run in float16 (half the memory, faster on tensor cores) while numerically fragile steps, like accumulating gradients or certain reductions, are automatically kept in float32 by torch.autocast, with a GradScaler scaling the loss up before backward and back down before the optimizer step to prevent tiny float16 gradients from underflowing to zero. Because these APIs check availability internally, a program written with device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') runs correctly, if not fast, on CPU-only machines.

Connect it to a real scenario

Training the Tutorial Platform's search-ranking model or lesson-recommendation model on the full interaction history — millions of clicks and completions — would take impractically long on CPU alone; renting a GPU instance for the training job and writing the model to be device-agnostic from day one (using device = torch.device(...) everywhere instead of hardcoding .cpu()) means the exact same training script scales from a laptop prototype to a GPU box without code changes. Mixed precision matters most for the largest model here, the recommendation transformer, where it can roughly halve training time and memory, letting the team try more architectures in the same budget — while the lightweight sentiment classifier trains fine either way.

• Detect the device with torch.device('cuda' if torch.cuda.is_available() else 'cpu') and call .to(device) on both the model and every input tensor — code should never hardcode 'cuda'.
• All tensors participating in an operation must be on the same device; mixing GPU and CPU tensors raises a runtime error rather than an automatic transfer.
• Mixed precision (torch.autocast + GradScaler) speeds up GPU training by running most ops in float16 while keeping numerically sensitive steps in float32 automatically.

Try the working example

import torch
import torch.nn as nn

torch.manual_seed(0)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device:", device)

class TinyClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(10, 16)
        self.fc2 = nn.Linear(16, 2)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        return self.fc2(x)

model = TinyClassifier().to(device)
sample_input = torch.randn(4, 10).to(device)

with torch.no_grad():
    output = model(sample_input)

print("Output device:", output.device)
print("Output shape:", output.shape)

# Mixed precision training sketch -- only actually uses float16 on CUDA;
# falls back safely to normal float32 compute on CPU
scaler = torch.cuda.amp.GradScaler(enabled=torch.cuda.is_available())
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
target = torch.randint(0, 2, (4,)).to(device)

optimizer.zero_grad()
with torch.autocast(device_type=device.type, enabled=torch.cuda.is_available()):
    logits = model(sample_input)
    loss = nn.functional.cross_entropy(logits, target)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

print("Training step completed. Loss:", loss.item())

Using device: cpu
Output device: cpu
Output shape: torch.Size([4, 2])
Training step completed. Loss: 0.7749
(on a machine with a CUDA GPU, the first two lines would print 'cuda' instead, and autocast would actually run parts of the computation in float16)

PyTorch Docs — Automatic Mixed Precision — Deep Learning