Exercises

# ============================================================
# Drill 1 — zero_grad
# ============================================================
# Fix this loop:
# for xb, yb in loader:
#     loss = loss_fn(model(xb), yb)
#     loss.backward()
#     opt.step()
#
# ANSWER: forgot opt.zero_grad(). Gradients accumulate across steps.
# Best practice:
#   opt.zero_grad(set_to_none=True)    # faster than .zero_grad()
#   loss.backward()
#   opt.step()

# ============================================================
# Drill 2 — Eval mode
# ============================================================
# Your predictions on the validation set are noisy. Why?
#
# ANSWER: forgot model.eval(). Dropout still active, BatchNorm uses
# batch stats. Call model.eval() AND wrap with torch.inference_mode():
#   model.eval()
#   with torch.inference_mode():
#       preds = model(X)

# ============================================================
# Drill 3 — Device + non_blocking
# ============================================================
# Your GPU sits at 40% during training. Loader is fine. Cause?
#
# ANSWER: synchronous CPU→GPU copy.
# Use pin_memory + non_blocking:
#   loader = DataLoader(ds, batch_size=128, num_workers=4, pin_memory=True)
#   xb = xb.to(device, non_blocking=True)
#   yb = yb.to(device, non_blocking=True)

# ============================================================
# Drill 4 — Mixed precision
# ============================================================
# Training a transformer is too slow. You have an A100. Fix it.
#
# ANSWER: autocast + GradScaler.
# scaler = torch.amp.GradScaler('cuda')
# for xb, yb in loader:
#   opt.zero_grad(set_to_none=True)
#   with torch.amp.autocast(device_type='cuda', dtype=torch.float16):
#     loss = loss_fn(model(xb), yb)
#   scaler.scale(loss).backward()
#   scaler.unscale_(opt)
#   torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
#   scaler.step(opt)
#   scaler.update()

# ============================================================
# Drill 5 — Multi-GPU
# ============================================================
# You have 4 GPUs in one box. DataParallel or DistributedDataParallel?
#
# ANSWER: DDP, ALWAYS.
# DataParallel is single-process, suffers from GIL, and gradients funnel
# through the master GPU. DDP runs one process per GPU and uses NCCL all-reduce.
# Launch:
#   torchrun --standalone --nproc-per-node=4 train.py
#
# Inside train.py:
#   dist.init_process_group(backend='nccl')
#   local_rank = int(os.environ['LOCAL_RANK'])
#   torch.cuda.set_device(local_rank)
#   model = DDP(model.to(local_rank), device_ids=[local_rank])

# ============================================================
# Bonus — Why does loss become NaN after a few epochs?
# ============================================================
# ANSWER:
# - Exploding gradients -> gradient clipping (clip_grad_norm_)
# - Log of zero / sqrt of negative in a custom loss
# - Mixed precision instability -> use bf16 on A100/H100; or scale loss

# ============================================================
# Scoring
# ============================================================
# 5 / 5 -> review training loops at speed
# 3 / 5 -> bookmark pytorch/cheatsheet
# < 3   -> read PyTorch's 'Performance Tuning Guide'

# Fill in: opt.____() before loss.backward()