PyTorch Foundations For LLMs

import torch

def stable_softmax(x: torch.Tensor, dim: int = -1) -> torch.Tensor:
    shifted = x - x.max(dim=dim, keepdim=True).values
    exp = shifted.exp()
    return exp / exp.sum(dim=dim, keepdim=True)

def cross_entropy_from_logits(logits: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    # logits: [batch, classes], target: [batch]
    log_probs = logits - torch.logsumexp(logits, dim=-1, keepdim=True)
    nll = -log_probs.gather(dim=-1, index=target[:, None]).squeeze(-1)
    return nll.mean()

LayerNorm centers and scales. RMSNorm scales by root mean square without subtracting the mean. Many decoder LLMs use RMSNorm because it is simple, stable, and efficient.

def layer_norm(x, weight, bias, eps=1e-5):
    mean = x.mean(dim=-1, keepdim=True)
    var = (x - mean).pow(2).mean(dim=-1, keepdim=True)
    return (x - mean) * torch.rsqrt(var + eps) * weight + bias

def rms_norm(x, weight, eps=1e-6):
    rms = x.pow(2).mean(dim=-1, keepdim=True)
    return x * torch.rsqrt(rms + eps) * weight

Implementation Drill: Test RMSNorm Without Accidentally Recentering

RMSNorm keeps the rescaling part of LayerNorm but removes mean subtraction. The original RMSNorm paper frames this as preserving re-scaling invariance while dropping re-centering, and PyTorch exposes the layer as torch.nn.RMSNorm with an optional per-element affine weight. A good unit test should therefore compare your function with the library module and also prove that shifted inputs are not treated like LayerNorm.

def test_rms_norm_matches_pytorch():
    torch.manual_seed(0)
    x = torch.randn(3, 5, dtype=torch.float32)
    weight = torch.randn(5)

    expected = torch.nn.RMSNorm(5, eps=1e-6)
    expected.weight.data.copy_(weight)

    actual = rms_norm(x, weight, eps=1e-6)
    torch.testing.assert_close(actual, expected(x), rtol=1e-6, atol=1e-6)

    shifted = x + 10.0
    ln = torch.nn.LayerNorm(5, elementwise_affine=False)
    assert torch.allclose(ln(x), ln(shifted), atol=1e-5)
    assert not torch.allclose(rms_norm(x, weight), rms_norm(shifted, weight))

1. Match the module: compare shape, dtype, epsilon, and affine-weight behavior against torch.nn.RMSNorm before using a hand-written layer.
2. Do not subtract the mean: RMSNorm should change when a large constant is added to the hidden state, while affine-free LayerNorm should not.
3. Set epsilon deliberately: PyTorch accepts eps=None, but course code should pass an explicit value so tests are stable across dtypes.

Sources: Root Mean Square Layer Normalization, PyTorch RMSNorm reference, PyTorch LayerNorm reference.

def train_step(model, batch, optimizer, scheduler=None, grad_clip=1.0):
    model.train()
    optimizer.zero_grad(set_to_none=True)
    loss = model(**batch).loss
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
    optimizer.step()
    if scheduler is not None:
        scheduler.step()
    return float(loss.detach())

Implementation Drill: Accumulate Microbatches Without Changing The Optimizer Clock

Gradient accumulation works because PyTorch adds gradients into each parameter's .grad buffer when backward() runs. That means the training loop must deliberately choose when to clear gradients and when to advance optimizer state. Hugging Face Accelerate's accumulation guide frames the same idea as stepping only after several batches have contributed gradients. For a hand-written loop, divide the microbatch loss, call backward() on every microbatch, then clip, step the optimizer, step the scheduler, and zero gradients only at the update boundary. If the last accumulation window is shorter, scale by that window's actual size.

def train_epoch_accum(model, dataloader, optimizer, scheduler=None,
                      accumulate_steps=4, grad_clip=1.0):
    model.train()
    optimizer.zero_grad(set_to_none=True)
    total_loss = 0.0
    total_batches = len(dataloader)

    for micro_step, batch in enumerate(dataloader, start=1):
        window_start = ((micro_step - 1) // accumulate_steps) * accumulate_steps + 1
        window_end = min(window_start + accumulate_steps - 1, total_batches)
        window_size = window_end - window_start + 1

        loss = model(**batch).loss
        total_loss += float(loss.detach())
        (loss / window_size).backward()

        should_update = micro_step % accumulate_steps == 0
        is_last = micro_step == total_batches
        if should_update or is_last:
            torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
            optimizer.step()
            if scheduler is not None:
                scheduler.step()
            optimizer.zero_grad(set_to_none=True)

    return total_loss / total_batches

1. Scale the loss: if the loss is averaged per microbatch, divide by the accumulation-window size before backward so the accumulated gradient approximates the larger batch gradient.
2. Zero at boundaries: zero gradients before the accumulation window and after the optimizer step, not between microbatches.
3. Clip once: clip the accumulated gradient just before optimizer.step(); clipping each microbatch changes the effective direction.
4. Step schedulers by update: for per-update schedulers, call scheduler.step() with the optimizer step, not on every microbatch.
5. Log both clocks: record microbatch count, optimizer update count, effective batch size, and learning rate so resumed runs are debuggable.

Sources: PyTorch zeroing gradients recipe, PyTorch clip_grad_norm_ reference, Hugging Face Accelerate gradient accumulation guide.

Activation checkpointing reduces training memory by not saving every intermediate activation during the forward pass. During backward, PyTorch recomputes the checkpointed region so gradients can still be calculated. That tradeoff is useful for deep transformer blocks when activation memory is the bottleneck, but it spends extra compute and can hide bugs if the recomputed forward is not deterministic.

from torch.utils.checkpoint import checkpoint

class CheckpointedBlock(torch.nn.Module):
    def __init__(self, block):
        super().__init__()
        self.block = block

    def forward(self, hidden_states, attention_mask):
        def run_block(x, mask):
            return self.block(x, attention_mask=mask)

        return checkpoint(
            run_block,
            hidden_states,
            attention_mask,
            use_reentrant=False,
        )

1. Checkpoint large regions: wrap transformer blocks or groups of blocks, not tiny cheap ops whose recompute overhead dominates the saved memory.
2. Prefer the explicit API: pass use_reentrant=False deliberately and keep the choice in the experiment config because PyTorch documents behavioral differences between variants.
3. Preserve determinism: dropout, random sampling, and global state changes inside the checkpointed region can make recomputation disagree with the original forward.
4. Measure both sides: report peak memory, tokens per second, and loss parity against the non-checkpointed run before treating the change as a win.

Policy Drill: Save Expensive Ops, Recompute Cheap Ops

Plain checkpointing recomputes the whole wrapped region. Selective activation checkpointing lets the policy keep expensive operations, such as matrix multiplies, while still recomputing cheaper pointwise work. In transformer training, this gives a more useful knob than "checkpoint this block or not": keep attention and MLP matmuls from being replayed first, then measure whether the remaining recompute saves enough activation memory to justify the complexity.

import functools
from torch.utils.checkpoint import (
    CheckpointPolicy,
    checkpoint,
    create_selective_checkpoint_contexts,
)

ops_to_save = [
    torch.ops.aten.mm.default,
    torch.ops.aten.bmm.default,
    torch.ops.aten.addmm.default,
]

def checkpoint_policy(ctx, op, *args, **kwargs):
    if op in ops_to_save:
        return CheckpointPolicy.MUST_SAVE
    return CheckpointPolicy.PREFER_RECOMPUTE

context_fn = functools.partial(
    create_selective_checkpoint_contexts,
    checkpoint_policy,
)

hidden_states = checkpoint(
    run_transformer_block,
    hidden_states,
    attention_mask,
    use_reentrant=False,
    context_fn=context_fn,
)

1. Start from the profiler: save the ops that dominate recompute time before adding a broad policy.
2. Keep the policy small: document which ops are saved and why, because every saved tensor gives back some memory.
3. Check mutation errors: PyTorch raises by default if tensors cached by selective checkpointing are mutated, which is usually a correctness signal.
4. Compare three runs: benchmark eager, plain checkpointing, and selective checkpointing on the same batch shape and sequence length.

Sources: PyTorch activation checkpointing reference, PyTorch activation checkpointing techniques.

Treat AMP and torch.compile as training-loop changes that need correctness checks, not as one-line speed toggles. PyTorch AMP uses autocast to choose lower precision for suitable ops while keeping numerically sensitive work in float32. With float16 training, GradScaler protects small gradients from underflow; if you clip gradients, unscale them before clipping so the threshold applies to real gradients. torch.compile can optimize a model or function, but shape changes, Python control flow, and graph breaks can cause recompilation or fallback, so benchmark it against eager mode.

def amp_train_step(model, batch, optimizer, scaler, grad_clip=1.0):
    model.train()
    optimizer.zero_grad(set_to_none=True)
    with torch.amp.autocast("cuda", dtype=torch.float16):
        loss = model(**batch).loss
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
    scaler.step(optimizer)
    scaler.update()
    return float(loss.detach())

compiled_model = torch.compile(model, dynamic=True)

# Debug checklist:
# 1. Compare eager vs compiled losses on the same batch.
# 2. Run fixed-shape and variable-length batches.
# 3. Try fullgraph=True once to expose graph breaks early.
# 4. Use TORCH_LOGS=guards,dynamic when recompilation is surprising.

1. Establish parity: compare loss, grad norms, and a short eval sample before accepting a speedup.
2. Clip after unscale: scaled gradients make a clipping threshold meaningless until the scaler has unscaled them.
3. Bucket shapes: LLM batches often vary by sequence length, so use padding buckets or dynamic tracing before blaming the compiler.
4. Measure warm and cold: separate first-compile overhead from steady-state tokens or batches per second.

Implementation Drill: Keep AMP Clipping On Real Gradients

Mixed precision changes the training loop's state machine. PyTorch's AMP examples show that gradient clipping should happen after scaler.unscale_(optimizer), because the norm threshold is defined for the real gradients, not the temporarily scaled gradients. Treat the scaler as the owner of the optimizer step: call scaler.step(optimizer), then scaler.update(), and log skipped updates when overflow makes the step a no-op.

def amp_step_with_clip(model, batch, optimizer, scaler, max_norm=1.0):
    optimizer.zero_grad(set_to_none=True)

    with torch.amp.autocast("cuda", dtype=torch.float16):
        loss = model(**batch).loss

    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)  # Gradients are now in their real scale.
    grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)

    scale_before = scaler.get_scale()
    scaler.step(optimizer)
    scaler.update()
    skipped = scaler.get_scale() < scale_before

    return {
        "loss": float(loss.detach()),
        "grad_norm": float(grad_norm),
        "amp_step_skipped": skipped,
    }

1. Never double step: do not call optimizer.step() after scaler.step(optimizer); the scaler already delegates to the optimizer when gradients are finite.
2. Clip once: unscale once per optimizer per step, then clip the accumulated gradients immediately before the scaler step.
3. Log skips: a falling scale or skipped update can explain flat loss curves even when the forward pass looks healthy.
4. Compare fp32: run a tiny deterministic batch in fp32 and AMP to verify loss direction and gradient norm scale before a long training job.

Sources: PyTorch AMP reference, PyTorch AMP examples, PyTorch torch.compile reference.

PyTorch evaluation has two independent switches. model.eval() changes module behavior such as dropout and batch normalization, while grad modes decide whether operations are recorded by autograd. Use torch.inference_mode() for pure validation or serving paths whose tensors will not later re-enter a gradient-recorded computation. Keep torch.no_grad() for mixed workflows where you only want to skip recording a block but still need normal tensor behavior afterward.

def validate(model, dataloader, metric_fn):
    was_training = model.training
    model.eval()
    totals = []
    with torch.inference_mode():
        for batch in dataloader:
            logits = model(**batch).logits
            totals.append(metric_fn(logits, batch["labels"]))
    if was_training:
        model.train()
    return torch.stack(totals).mean().item()

# Use no_grad, not inference_mode, when outputs must flow into later autograd work.
with torch.no_grad():
    frozen_features = encoder(inputs)
loss = trainable_head(frozen_features).loss(labels)
loss.backward()

1. Set both switches: call eval() for module semantics and a grad mode for autograd memory and bookkeeping.
2. Restore training state: validation inside a training loop should return the model to the previous mode.
3. Use inference mode narrowly: prefer it for read-only eval and serving, but avoid it when tensors created inside the block must participate in later backward passes.
4. Test stochastic layers: compare two eval passes on the same batch to catch dropout, batch norm, or generation-mode mistakes.

Bug Drill: Grad Mode Does Not Rewrite Parameter Intent

Grad-disabled contexts affect operation recording, not every tensor's long-term role. PyTorch documents that factory functions which accept requires_grad are not affected by torch.no_grad(), so creating a Parameter inside validation can still produce trainable state. Conversely, tensors created inside torch.inference_mode() are intentionally restricted from later autograd-recorded use. Keep parameter construction, frozen feature extraction, and read-only serving paths separate in tests.

with torch.no_grad():
    probe = torch.nn.Parameter(torch.randn(8))
assert probe.requires_grad

with torch.no_grad():
    frozen_features = encoder(input_ids)
loss = classifier(frozen_features).loss(labels)
loss.backward()  # OK: no_grad skipped encoder recording only.

with torch.inference_mode():
    serving_features = encoder(input_ids)
# Do not feed serving_features into a later training loss.

1. Construct modules outside eval: validation should not create new trainable parameters as a side effect.
2. Choose the weaker mode when needed: use no_grad for frozen-feature workflows that still feed trainable heads.
3. Reserve inference mode: use inference_mode for pure read-only evaluation, export checks, and serving code paths.

Sources: PyTorch locally disabling gradient computation, PyTorch no_grad reference, PyTorch inference_mode reference.