The Core Pretraining Step: Forward, Loss, Backward, and Gradient Clipping
Agent-Assisted Engineering Check
Prerequisites
Dataloader from Day 49 and AdamW optimizer from Day 50.
Learning Objective
Execute the canonical 5-step pretraining iteration, understand memory deallocation via set_to_none=True, and apply L2 gradient norm clipping.
Instructions
Study source: within this activity's existing reading/practice time, use CS336 Assignment 1: Basics (Spring 2025 tagged release), tests/adapters.py: run_gradient_clipping; tests/test_nn_utils.py: test_gradient_clipping, then explain one connection to the local task in your own words. Keep the local exercise and acceptance checks; use its original explanation if the external source is unavailable. In every pretraining step: (1) optimizer.zero_grad(set_to_none=True) deallocates gradient buffers; (2) logits, loss = model(x, y) computes forward activations and cross-entropy; (3) loss.backward() accumulates partial derivatives; (4) torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) clamps the total gradient vector norm to 1.0; (5) optimizer.step() applies AdamW updates. Trace why gradient clipping prevents catastrophic weight corruption.
Practice & Prediction Log
1. Write the Python code executing one full training step.
2. If the unclipped global gradient norm is ||g||_2 = 5.0 and max_norm = 1.0, what scaling multiplier is applied to every parameter's gradient?
3. Why is optimizer.zero_grad(set_to_none=True) strictly superior to optimizer.zero_grad()?
The scaling factor is min(1.0, max_norm / ||g||_2). Memory write of zero bytes vs free pointer.
1. Training step code:
optimizer.zero_grad(set_to_none=True)
logits, loss = model(x, y)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()2. Scaling multiplier: 1.0 / 5.0 = 0.2. Every gradient tensor is multiplied by 0.2, capping the total L2 norm at 1.0.
3.
set_to_none=True frees gradient tensor memory rather than writing zeros into allocated buffers, reducing GPU memory bandwidth consumption and enabling faster autograd allocation.
Learner implements the canonical 5-step training step, calculates gradient scaling factors, and articulates memory benefits of set_to_none=True.
Authoritative Source Resources & Provenance
CS336 Assignment 1: Basics (Spring 2025 tagged release)
Relevant Section / Range: tests/adapters.py: run_gradient_clipping; tests/test_nn_utils.py: test_gradient_clipping
Verified Version & Date: spring2025 tag; README, tests and MIT LICENSE checked 2026-11-01
Lead assignment-driven source for pretraining component practice and tests. nanoGPT remains a compact historical code-reading companion, not the current recommended training curriculum.
Fluent Python, PyTorch, and linear algebra
Local self-contained pretraining scripts run on standard local CPUs or budget GPUs without Stanford cluster dependency.
Compared Candidate Sources & Selection Analysis (2)
Simpler code layout, useful complementary reference for minimal PyTorch Transformer.
High-level API focus rather than building from raw tensor fundamentals.
nanoGPT: train.py pretraining loop
Relevant Section / Range: train.py: training loop (while True): forward under ctx, scaler.scale(loss).backward(), clip_grad_norm_(grad_clip), scaler.step(optimizer)
Verified Version & Date: Git master commit 71c3905 (Verified: 2026-10-30)
Compact historical code reference for the named implementation functions; CS336 spring2025 provides the lead component exercises. Do not infer top-p, atomic persistence or cloud budget enforcement from nanoGPT.
Python, PyTorch, and Day 46-48 Transformer model architecture.
Runnable pretraining scripts and step-by-step tensor mechanics are fully self-contained in the curriculum lessons.
Compared Candidate Sources & Selection Analysis (2)
Heavy abstraction layer wrapping distributed training; nanoGPT exposes raw PyTorch loops without framework overhead.
High-performance enterprise distributed framework; unsuitable for local educational first-principles comprehension.