fareedkhan-dev--train-llm-from-scratch/docs/foundations/optimization.md

Optimization & Training Systems

Once the loss is defined, training is an engineering problem: move parameters in the right direction without numerical instability, memory blowups, or throughput collapse.

The main ingredients in this repo are:

• AdamW;
• linear warmup plus cosine learning-rate decay;
• gradient accumulation;
• gradient clipping;
• bf16 autocast;
• DistributedDataParallel for multi-GPU training.

The training step

flowchart LR
    B["batch"] --> F["forward under autocast"]
    F --> L["loss / grad_accum"]
    L --> BW["backward"]
    BW --> GA{"more microsteps?"}
    GA -- yes --> B
    GA -- no --> CLIP["clip grad norm"]
    CLIP --> LR["set scheduled LR"]
    LR --> STEP["AdamW step"]
    STEP --> ZERO["zero gradients"]

The pretraining loop in scripts/pretrain_base.py uses this pattern:

for micro in range(cfg.grad_accum):
    xb, yb = next(batch_iter)
    with amp_autocast(cfg.amp_dtype, ctx.device):
        _, loss = model(xb, yb)
        loss = loss / cfg.grad_accum
    loss.backward()

torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip)
optimizer.step()
optimizer.zero_grad(set_to_none=True)

Dividing the loss by grad_accum keeps the gradient scale the same as if the full effective batch had fit in memory.

AdamW

Adam keeps exponential moving averages of gradients and squared gradients:

[ m_t = \beta_1 m_{t-1} + (1-\beta_1)g_t ]

[ v_t = \beta_2 v_{t-1} + (1-\beta_2)g_t^2 ]

After bias correction, parameters are updated approximately by:

[ \theta_{t+1} = \theta_t - \eta \frac{\hat{m}_t}{\sqrt{\hat{v}_t}+\epsilon} ]

AdamW decouples weight decay from the gradient update:

[ \theta_{t+1} = \theta_t - \eta \left( \frac{\hat{m}_t}{\sqrt{\hat{v}_t}+\epsilon}

• \lambda \theta_t \right) ]

The repo applies weight decay only to matrix-like parameters:

if p.dim() >= 2:
    decay.append(p)
else:
    no_decay.append(p)

This is the standard GPT recipe: decay large weight matrices, but not biases, LayerNorm scales, or one-dimensional parameters.

Learning-rate warmup and cosine decay

The learning rate is small at the start, ramps up, then decays:

[ \eta(s) = \eta_{\max}\frac{s+1}{S_{\text{warmup}}} \quad \text{if } s < S_{\text{warmup}} ]

After warmup:

[ \eta(s) = \eta_{\min}

• \frac{1}{2}(1+\cos(\pi p))(\eta_{\max}-\eta_{\min}) ]

where:

[ p = \frac{s-S_{\text{warmup}}}{S_{\max}-S_{\text{warmup}}} ]

Implementation in src/post_training/optim.py:

if step < warmup_steps:
    return lr * (step + 1) / max(1, warmup_steps)
progress = (step - warmup_steps) / max(1, max_steps - warmup_steps)
coeff = 0.5 * (1.0 + math.cos(math.pi * progress))
return min_lr + coeff * (lr - min_lr)

Warmup prevents early unstable updates while weights are still poorly calibrated. Cosine decay reduces step size as training approaches the end of the budget.

Gradient accumulation

If one batch is too large for GPU memory, split it into microbatches:

[ B_{\text{effective}} = B_{\text{micro}} \times N_{\text{accum}} \times N_{\text{gpus}} ]

Example:

• microbatch size: 8;
• accumulation steps: 12;
• GPUs: 2.

[ B_{\text{effective}} = 8 \times 12 \times 2 = 192 ]

The optimizer steps once after all microbatches have contributed gradients.

Gradient clipping

Gradient clipping limits the global norm:

[ g \leftarrow g \cdot \min\left(1, \frac{c}{|g|_2}\right) ]

If the gradient norm is below the threshold (c), nothing changes. If it is too large, the whole gradient vector is scaled down. This is a stability guard, especially useful in RL and long-sequence training.

bf16 autocast

bf16 uses fewer bits than fp32, but keeps an 8-bit exponent like fp32. That makes it much more forgiving than fp16 for deep learning training.

The repo uses autocast for forward computation:

with amp_autocast(cfg.amp_dtype, ctx.device):
    logits, _ = model(tokens)
    loss = sft_loss(logits, tokens, mask)

Model parameters usually remain fp32. Many matrix multiplications run in bf16, improving memory and throughput on supported GPUs.

DistributedDataParallel

DDP creates one process per GPU. Each process:

1. owns a full copy of the model;
2. receives a different shard or random stream of data;
3. computes gradients locally;
4. synchronizes gradients across processes before the optimizer step.

With gradient accumulation, synchronization is needed only on the last microstep. The repo uses model.no_sync() for earlier microsteps to avoid unnecessary communication.

flowchart LR
    R0["rank 0 GPU"] --> G["gradient all-reduce"]
    R1["rank 1 GPU"] --> G
    G --> U0["rank 0 optimizer step"]
    G --> U1["rank 1 optimizer step"]

What to watch during training

Metric	Healthy behavior	Problem signal
train loss	falls steadily	flat near random baseline
dev loss	falls, then stabilizes	rises while train loss falls
grad norm	finite, bounded after clipping	NaN or repeated huge spikes
tokens/sec	stable for same config	sudden drop or dataloader stall
KL in RL stages	bounded	runaway drift from reference
reward in RL stages	rises with variance	zero signal for many iterations

Memory levers

If a config does not fit, reduce in this order:

1. batch_size;
2. context_length;
3. n_blocks;
4. n_embed;
5. n_head only if it still divides n_embed.

Context length is especially expensive because attention uses a (T \times T) score matrix.

Next

After training, the model still only emits logits. Generation turns those logits into text. Continue to Generation & Sampling.