GPU Acceleration, Mixed Precision & torch.compile

1. Precision Formats: FP32 vs FP16 vs BF16

Standard deep learning computes in 32-bit floating point (FP32). Modern NVIDIA Tensor Cores (Ampere, Hopper, Blackwell) process 16-bit representations at up to 4x higher mathematical throughput:

2. Automatic Mixed Precision with torch.amp & GradScaler

Automatic Mixed Precision evaluates compute-intensive layers (matrix multiplications, convolutions) in 16-bit while keeping numerically sensitive reductions (Softmax, LayerNorm, loss aggregations) in FP32. When using FP16, GradScaler multiplies loss by a scale factor to prevent gradient underflow:

import torch
import torch.nn as nn
from torch.amp import autocast, GradScaler

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MyNeuralNet().to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()

# Initialize FP16 gradient scaler
scaler = GradScaler('cuda')

for inputs, targets in dataloader:
    inputs, targets = inputs.to(device), targets.to(device)
    optimizer.zero_grad(set_to_none=True)

    # 1. Forward pass under autocast context
    # (Automatically casts Conv/Linear layers to FP16/BF16)
    with autocast(device_type="cuda", dtype=torch.float16):
        outputs = model(inputs)
        loss = criterion(outputs, targets)

    # 2. Scaled backward pass
    # Multiplies loss by scale factor so gradients don't flush to zero
    scaler.scale(loss).backward()

    # 3. Unscales gradients and invokes optimizer.step()
    # (Skips step if NaN/Inf gradients are detected and reduces scale factor)
    scaler.step(optimizer)

    # 4. Update scale factor for next iteration
    scaler.update()

3. Next-Gen JIT Compilation: torch.compile()

Traditionally, PyTorch runs in eager mode: each line of Python sends an individual CUDA kernel launch across the driver. Introduced in PyTorch 2.0, torch.compile() captures the computational graph usingTorchDynamo and compiles it into fused Triton CUDA kernels with TorchInductor:

# Compiling your model with one line of code:
compiled_model = torch.compile(
    model,
    mode="reduce-overhead",  # Best for latency & small batches using CUDA graphs
    fullgraph=False          # Falls back cleanly to Python eager mode on unsupported ops
)

# Benchmark inference / training
with torch.inference_mode():
    y = compiled_model(sample_tensor)

4. Multi-GPU Scaling: Why DDP Replaces DataParallel

Instead, use DistributedDataParallel (DDP). DDP spawns one independent OS process per GPU, completely bypassing the Python GIL and performing all-reduce gradient synchronizations via high-speed NCCL ring interconnects:

# Multi-GPU Launch Command (e.g., 4 GPUs on 1 node):
# torchrun --nproc_per_node=4 train_ddp.py

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler

def main_ddp():
    # Initialize process group via NCCL backend
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)

    # Wrap model in DDP
    model = MyModel().to(local_rank)
    model = DDP(model, device_ids=[local_rank])

    # DistributedSampler ensures each GPU receives non-overlapping data slices
    sampler = DistributedSampler(dataset, shuffle=True)
    loader = DataLoader(dataset, batch_size=32, sampler=sampler)

    for epoch in range(10):
        sampler.set_epoch(epoch)  # Critical for proper multi-GPU shuffling
        # ... standard training loop ...