Training Tips and Tricks: LR Scheduling, Gradient Clipping and Mixed Precision

The previous three subtopics (optimizers, regularization, BatchNorm) cover the core training theory. In practice, three more engineering techniques make the difference between a training run that works smoothly and one that stalls, diverges, or takes forever — and all three are used in essentially every serious deep learning project, including LLM training.

Learning rate scheduling changes the step size (learning rate) OVER THE COURSE of training, rather than keeping it fixed. Two common patterns: warmup starts with a small learning rate and gradually increases it over the first several hundred/thousand steps — useful because early in training, weights are random and gradients can be erratic, so a big step right away risks a bad, unstable update; ramping up gives the model a chance to settle into a reasonable region first. Decay does the opposite later in training: gradually SHRINK the learning rate (linearly, or following a cosine curve) as training progresses, so early steps make big exploratory progress while later steps make small, precise refinements near the minimum — like using big shovel-fulls to dig a hole quickly, then a small trowel to smooth the final shape.

Gradient clipping guards against a specific failure mode called the 'exploding gradient': occasionally, a batch produces an unusually huge gradient (common in recurrent architectures and early in training generally), and taking a full step in that direction can catastrophically wreck the weights, sometimes irrecoverably. Gradient clipping caps the gradient's magnitude before the update is applied — if the gradient's norm exceeds some threshold, it's rescaled down to that threshold while keeping its direction — turning a training-destroying spike into a merely large (but survivable) step.

Mixed precision training uses lower-precision numbers (16-bit floats, FP16 or BF16) for most of the computation instead of the default 32-bit (FP32), roughly halving memory usage and significantly speeding up computation on modern GPUs (which have specialized hardware for 16-bit math). The 'mixed' part matters: certain sensitive operations (like the running sums in gradient accumulation) are kept in 32-bit precision to avoid numeric instability, while the bulk of the matrix multiplications run in 16-bit. This is a training-side sibling of the quantization technique from the LLM Optimization unit — same core idea (fewer bits, faster/cheaper), applied during training instead of inference.

import math

def lr_with_warmup_and_cosine_decay(step: int, total_steps: int,
                                    warmup_steps: int, peak_lr: float) -> float:
    """Ramp LR up linearly during warmup, then decay it following a
    cosine curve down to ~0 by the end of training."""
    if step < warmup_steps:
        return peak_lr * (step / warmup_steps)         # linear warmup
    progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
    return peak_lr * 0.5 * (1 + math.cos(math.pi * progress))  # cosine decay

for step in [0, 50, 100, 500, 1000, 2000]:
    lr = lr_with_warmup_and_cosine_decay(step, total_steps=2000,
                                          warmup_steps=100, peak_lr=1e-3)
    print(f"step {step:5d}: lr = {lr:.6f}")

def clip_gradient(grad: list[float], max_norm: float = 1.0) -> list[float]:
    """Rescale the gradient vector down to max_norm if its norm exceeds
    it -- keeps direction, caps magnitude, preventing exploding updates."""
    norm = math.sqrt(sum(g ** 2 for g in grad))
    if norm <= max_norm:
        return grad
    scale = max_norm / norm
    return [g * scale for g in grad]

exploding_grad = [15.0, -22.0, 8.0]   # unusually huge gradient
print("clipped gradient:", [round(g, 3) for g in clip_gradient(exploding_grad)])
# Mixed precision (real usage): torch.cuda.amp.autocast() around the
# forward pass runs matmuls in FP16 while keeping sensitive ops in FP32