DataArena
Week 6 • Day 3Week 6

Day 3:Building Neural Networks with `torch.nn.Module`, `AdamW` & The 5-Step Training Loop

Battle-tested daily lesson note • 100% runnable code, interview Q&A, test suite & interactive quiz

Week 6 — Day 3: Building Neural Networks with torch.nn.Module, AdamW & The 5-Step Training Loop

Mentor's Note (10+ YOE Staff AI Engineer POV): Every PyTorch model in the world—from a 2-layer MLP to a 405-Billion parameter Llama-3 model—inherits from torch.nn.Module and is trained using the Exact Same 5-Line Training Loop!

And what is the #1 bug in PyTorch interviews? Forgetting optimizer.zero_grad() before loss.backward()! Remember from Day 1 (self.grad += ...) that PyTorch accumulates (+=) gradients by default! If you forget optimizer.zero_grad(), Batch #2's gradients get added on top of Batch #1's old gradients, sending your weights flying into the stratosphere!

Part 1: The Canonical 5-Step PyTorch Training Loop

python
model.train() # Enables training behaviors like Dropout! for x_batch, y_batch in dataloader: optimizer.zero_grad() # Step 1: Clear old gradients from previous batch! logits = model(x_batch) # Step 2: Forward pass (calls model.forward()) loss = criterion(logits, y_batch) # Step 3: Compute scalar loss loss.backward() # Step 4: Backpropagate gradients through the graph optimizer.step() # Step 5: Update weights (w <- w - lr * grad)

Part 2: Why Modern LLMs Use AdamW Instead of Standard SGD or Adam

  • Standard SGD: Uses the same scalar learning rate η\eta for all 8 billion parameters, getting stuck in ravines.
  • Adam (Adaptive Moment Estimation): Tracks an exponentially decaying moving average of the 1st moment (Mean of gradients mtm_t, momentum) and the 2nd moment (Uncentered Variance of gradients vtv_t): mt=β1mt−1+(1−β1)gt,vt=β2vt−1+(1−β2)gt2m_t = \beta_1 m_{t-1} + (1 - \beta_1) g_t, \qquad v_t = \beta_2 v_{t-1} + (1 - \beta_2) g_t^2
  • AdamW (Decoupled Weight Decay — Loshchilov & Hutter 2019): In original Adam, L2L_2 regularization was added to the gradient gtg_t, which meant parameters with large historical gradients (vtv_t) got less weight decay! AdamW decouples weight decay directly from the adaptive gradient step: θt←θt−1−η(m^tv^t+ϵ+λθt−1)\theta_t \leftarrow \theta_{t-1} - \eta \left( \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} + \lambda \theta_{t-1} \right) Every modern LLM (GPT-4, Llama-3, Qwen, Mistral) is trained with AdamW!

Part 3: Hands-On Lab — Multi-Layer Perceptron (nn.Module) & AdamW Trainer (mlp_training_loop.py)

Check out src/mlp_training_loop.py (in Tab 2: Codes — click Run in Browser to execute) and run the script and tests:

powershell
python week06_deep_learning_and_pytorch_from_scratch/day03_nn_module_layers_losses_and_adamw_optimizer/src/mlp_training_loop.py python -m unittest discover -s week06_deep_learning_and_pytorch_from_scratch/day03_nn_module_layers_losses_and_adamw_optimizer/tests -v

In Simple English: The 5-Step Training Recipe + Why AdamW Beats Plain SGD

Think of training a neural network like practicing free throws in basketball:

  1. optimizer.zero_grad(): Erase yesterday's practice notes before taking today's shot (otherwise gradients pile up infinitely!).
  2. Forward Pass (model(x)): Take the shot and see where the ball lands.
  3. Loss (loss_fn(pred, y)): Measure how many inches you missed the hoop.
  4. Backward Pass (loss.backward()): Figure out whether to push slightly higher or softer next time.
  5. Optimizer Step (optimizer.step()): Adjust your muscles (weights).

Why AdamW? Plain Gradient Descent uses the exact same step size for every weight. AdamW acts like an adaptive coach with momentum (mtm_t): weights with noisy steep hills take smaller cautious steps (vt\sqrt{v_t} scaling), weights on flat plateaus keep rolling smoothly, and Weight Decay gently prevents any single weight from growing huge and memorizing training data!

Week 6 Day 3: Week 6 — Day 3: Building Neural Networks with `torch.nn.Module`, `AdamW` & The 5-Step Training Loop | 6-Month AI Engineer Roadmap · DataArena