3.3 The Pretraining Objective & Loss¶

Before a language model can answer questions, write code, or follow instructions, it must first acquire a dense statistical model of language. That foundation is laid during pretraining, and the entire multi-trillion-token, multi-thousand-GPU campaign is guided by a single scalar number: the cross-entropy loss on the next-token prediction task. Understanding that loss — what it measures, why it is chosen, how it is computed in practice, and what its values mean — is arguably the most important prerequisite for every topic in the rest of this book.

This chapter dissects the pretraining objective from first principles. We cover the probabilistic framing, the exact mechanics of teacher forcing, the engineering of causal masking and label shifting, loss masking strategies for packing and multi-document batches, the bits-per-byte normalization, and the UL2 family of span-corruption alternatives. Throughout, we ground everything in concrete code you can run today.

Why Next-Token Prediction?¶

A language model assigns probabilities to sequences. Given a vocabulary \(\mathcal{V}\) of size \(V\), we want to model the probability of a sequence \(x = (x_1, x_2, \ldots, x_T)\). By the chain rule of probability this factorizes exactly:

There is no approximation here — this is a theorem. It says that modeling a joint distribution over sequences is equivalent to modeling a sequence of conditional distributions, each predicting the next token given all previous tokens. A decoder-only transformer with a causal mask does exactly this.

Why this objective rather than, say, a reconstruction autoencoder or a contrastive loss?

1. Every token in a long sequence contributes signal. A 2048-token sequence yields 2047 gradient-carrying predictions per example, making pretraining extraordinarily data-efficient compared to objectives that emit one signal per sequence.
2. It is a proper scoring rule. Maximizing log-likelihood under the true data distribution is provably equivalent to minimizing the KL divergence from the model distribution to the data distribution. The model is incentivized to be calibrated, not just accurate on a held-out label.
3. Scalability. The computation is embarrassingly parallelizable across the time dimension during training (unlike autoregressive inference). The forward and backward passes can be batched over all positions simultaneously.
4. Generality. Nothing about this objective is task-specific. A model that can predict the next token of code, math, natural language, and structured data simultaneously acquires cross-domain representations that transfer well.

The connection to tokenization and vocabulary is covered in Tokenization: BPE, WordPiece, Unigram & Byte-Level. Scaling behavior of this loss is analyzed in Scaling Laws: Kaplan, Chinchilla & Beyond.

The Cross-Entropy Loss and Negative Log-Likelihood¶

Formal Definition¶

Let \(f_\theta : \mathcal{V}^{<t} \to \mathbb{R}^V\) be the transformer that maps a prefix to logits over the vocabulary. The model’s predicted probability at position \(t\) is:

The negative log-likelihood (NLL) loss for a single sequence of length \(T\) is:

This is numerically identical to the categorical cross-entropy between the one-hot target distribution and the predicted softmax distribution:

The inner sum collapses because all probability mass in the one-hot target is on the true token \(x_t\).

The Information-Theoretic Viewpoint¶

The per-token loss equals the cross-entropy \(H(p_\text{data}, p_\theta)\). By decomposition:

The entropy \(H(p_\text{data})\) of the true data distribution is a constant. So minimizing cross-entropy is the same as minimizing the KL divergence from the model to the data. The irreducible entropy of natural language — Shannon’s classic estimates put English at roughly 0.6–1.3 bits per character, which for a BPE token covering ~4 characters is on the order of 2–5 bits per token — sets a floor the model can never beat.

Perplexity (PPL) is the exponentiated average loss:

Perplexity has the intuitive interpretation of the effective branching factor: a perplexity of 10 means the model is on average as uncertain as a uniform distribution over 10 equiprobable tokens. Well-pretrained 70B-class models reach perplexities on the order of 3–8 on standard benchmarks such as WikiText-103, depending on tokenizer and evaluation setup. (Always compare perplexity numbers only within the same tokenizer, as vocabulary size strongly affects the value.)

Teacher Forcing and the Causal Mask¶

What is Teacher Forcing?¶

During training we use teacher forcing: at every position \(t\) we feed the ground-truth token \(x_{t-1}\) as input, not the model’s own prediction \(\hat{x}_{t-1}\). This is what makes the training forward pass efficient — all positions can be computed simultaneously rather than sequentially.

Without teacher forcing (scheduled sampling or curriculum approaches), each prediction at time \(t\) depends on potentially erroneous earlier predictions, causing gradient signal to be noisy and training to be slow. Teacher forcing decouples positions and allows a single GPU kernel to produce all \(T\) predictions in one shot.

The Causal Mask¶

The quid pro quo for teacher forcing is that position \(t\) must not be allowed to see tokens \(x_{t+1}, \ldots, x_T\) — otherwise the model would trivially predict the next token by looking at it. This is enforced by the causal (lower-triangular) attention mask:

Applied before the softmax in each attention layer, entries with \(-\infty\) vanish after exponentiation, making the attention weight exactly zero. Position \(i\) can only attend to positions \(\le i\) (including itself).

The causal mask is discussed in detail with implementation in The Attention Mechanism From Scratch and Building a GPT From Scratch (nanoGPT-style).

Input sequence:  [BOS] T h e   c a t   s a t
Positions:         0   1 2 3 4 5 6 7 8 9 10 11

Teacher-forced input  (i-th logit predicts token i+1):
 input:  [BOS] T h e   c a t   s a
 target:  T   h e     c a t   s a t

Causal mask (T=6 shown, ✓=attend, ✗=masked):
         pos0 pos1 pos2 pos3 pos4 pos5
  pos0:    ✓    ✗    ✗    ✗    ✗    ✗
  pos1:    ✓    ✓    ✗    ✗    ✗    ✗
  pos2:    ✓    ✓    ✓    ✗    ✗    ✗
  pos3:    ✓    ✓    ✓    ✓    ✗    ✗
  pos4:    ✓    ✓    ✓    ✓    ✓    ✗
  pos5:    ✓    ✓    ✓    ✓    ✓    ✓

Shifting Labels¶

In PyTorch, the cross-entropy loss requires aligned (logits, targets) tensors. For a sequence of length \(T\), the model produces \(T\) logit vectors (one per input position), but the first useful prediction is at position 0 (predicting position 1). The standard idiom is:

• Input to the model: tokens \([x_0, x_1, \ldots, x_{T-1}]\)
• Targets: tokens \([x_1, x_2, \ldots, x_T]\) — shifted by 1

import torch
import torch.nn.functional as F

# Suppose tokens is shape (B, T+1): batch of sequences with one extra token
# e.g., tokens = tokenizer.encode(text) + [EOS]

def compute_lm_loss(logits: torch.Tensor, tokens: torch.Tensor) -> torch.Tensor:
    """
    logits : (B, T, V)   — model output for each of T input positions
    tokens : (B, T+1)    — the full sequence including the token AFTER the last input

    We align by:
      inputs  = tokens[:, :-1]   shape (B, T)  — fed to the model
      targets = tokens[:, 1:]    shape (B, T)  — what each position should predict
    """
    # targets are the tokens one step ahead of each input position
    targets = tokens[:, 1:]          # (B, T)

    # logits come from running the model on tokens[:, :-1]
    # Reshape for F.cross_entropy: expects (N, C) or (N, C, ...) 
    B, T, V = logits.shape
    loss = F.cross_entropy(
        logits.reshape(B * T, V),    # (B*T, V)
        targets.reshape(B * T),      # (B*T,)
        reduction='mean'             # average over all non-masked positions
    )
    return loss

This is the entirety of the loss function for standard causal language modeling. Everything else — masking, packing, weighting — is an elaboration of this core.

What F.cross_entropy Actually Computes¶

Every training-context code path in this chapter calls F.cross_entropy on a flattened (B*T, V) logit tensor, treating it as a black box. It isn’t one: over class indices, it is exactly a numerically stable log-softmax followed by a gather of the true-class log-probability and a mean over the non-ignored rows. The stable log-softmax itself is the logsumexp shift trick derived in Numerical Computing, Floating Point & Precision — \(\text{log\_softmax}(z)_i = z_i - (m + \log \sum_j \exp(z_j - m))\) with \(m = \max_j z_j\) — which is also where that chapter’s fused NumPy backward for this exact operation is worked out.

import torch
import torch.nn.functional as F

def cross_entropy_from_scratch(logits: torch.Tensor, targets: torch.Tensor, ignore_index: int = -100) -> torch.Tensor:
    """
    Reimplements F.cross_entropy(logits, targets, reduction='mean', ignore_index=ignore_index)
    from scratch: stable log-softmax + gather + masked mean.

    logits  : (N, V) float — raw scores, N = B*T when called on flattened LM logits
    targets : (N,)   long  — true class index per row, or ignore_index to skip that row
    returns : scalar — mean negative log-likelihood over rows where targets != ignore_index
    """
    # 1. Stable log-softmax via the max-shift trick (see foundations 1.4)
    m = logits.max(dim=-1, keepdim=True).values                    # (N,1)
    shifted = logits - m                                           # (N,V)
    logsumexp = shifted.exp().sum(dim=-1, keepdim=True).log()      # (N,1)
    log_probs = shifted - logsumexp                                # (N,V) == log_softmax(logits)

    # 2. Mask ignored rows: clamp targets so gather never indexes out of bounds
    #    on ignore_index=-100 rows; those rows are dropped in step 3 anyway.
    valid = targets != ignore_index                                # (N,) bool
    safe = targets.clamp_min(0)                                    # (N,) long, safe for gather

    # 3. Negative log-likelihood of the true class, averaged over active positions only
    nll = -log_probs.gather(1, safe.unsqueeze(1)).squeeze(1)       # (N,)
    return nll[valid].mean()                                       # scalar


# Verification against F.cross_entropy
torch.manual_seed(0)
N, V = 100, 32000
logits = torch.randn(N, V)
targets = torch.randint(0, V, (N,))
targets[::7] = -100  # ignore roughly 1 in 7 rows (padding / masked positions)

ref  = F.cross_entropy(logits, targets, ignore_index=-100, reduction='mean')
mine = cross_entropy_from_scratch(logits, targets, ignore_index=-100)

print(f"F.cross_entropy: {ref.item():.6f}")   # 11.051692
print(f"from scratch:    {mine.item():.6f}")  # 11.051691
assert torch.allclose(ref, mine, atol=1e-5)

Both lines print ~11.05 — a randomly-initialized model scoring a 32000-token vocabulary lands near \(\log V\) plus a logit-variance term, here about 11.05 nats — and the assert passes to machine precision. That confirms the flattened (B*T, V) call used throughout this chapter is nothing more than stable log-softmax, a gather at the true-class index, and a masked mean.

The Logit Tensor Is the Memory Bottleneck¶

That innocuous logits.reshape(B * T, V) is, at real batch sizes, the single largest allocation in a training step — and it is the first thing that will OOM when you try to train your own model. Count the bytes for a micro-batch of \(B \cdot T = 32 \times 2048 = 65{,}536\) tokens against a \(V = 32{,}768\) vocabulary:

For a 100M-parameter model the entire transformer trunk — weights, optimizer state, and all block activations — is smaller than this. Two properties make it especially nasty: the cost grows as \(B \cdot T \cdot V\), so it fights every attempt to raise throughput by increasing the micro-batch; and activation checkpointing does nothing about it, because the head sits outside the checkpointed blocks.

The 2026 standard fix is chunked (fused) linear cross-entropy: fold the lm_head matmul into the loss and process the token dimension in chunks, so at most chunk × V logits ever exist, and recompute those logits in the backward pass instead of storing them.

import torch
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint


def _chunk_ce(h, w, t):
    """One chunk -> (sum of CE over valid rows, valid count). The (C, V) logits are
    local to this function, so `checkpoint` frees them right after the forward and
    rebuilds one chunk at a time during backward."""
    logits = F.linear(h, w).float()                    # (C, V) fp32 — transient
    lse = torch.logsumexp(logits, dim=-1)              # (C,)
    valid = t != -100
    tgt = t.clamp_min(0).unsqueeze(-1)                 # keep gather in range
    ce = lse - logits.gather(-1, tgt).squeeze(-1)      # CE = logsumexp - logit[target]
    return (ce * valid).sum(), valid.sum()


def chunked_linear_cross_entropy(hidden, weight, targets, chunk: int = 8192):
    """Mathematically identical to F.cross_entropy(hidden @ weight.T, targets),
    but peak logit memory is chunk*V instead of (B*T)*V — independent of batch size.
    Gradients w.r.t. both `hidden` and `weight` accumulate correctly because autograd
    sums across the checkpointed calls."""
    h = hidden.reshape(-1, hidden.shape[-1])           # (B*T, d)
    t = targets.reshape(-1)                            # (B*T,)
    tot = h.new_zeros((), dtype=torch.float32)
    n = torch.zeros((), dtype=torch.long, device=h.device)
    for i in range(0, h.shape[0], chunk):
        s, c = checkpoint(_chunk_ce, h[i:i + chunk], weight, t[i:i + chunk],
                          use_reentrant=False)
        tot, n = tot + s, n + c
    return tot / n.clamp_min(1)


# Verify against the naive path on a small problem
torch.manual_seed(0)
BT, d, V = 512, 64, 1000
hidden = torch.randn(BT, d, requires_grad=True)
weight = (torch.randn(V, d) * 0.02).requires_grad_()
targets = torch.randint(0, V, (BT,))
targets[::9] = -100

ref = F.cross_entropy(F.linear(hidden, weight).float(), targets, ignore_index=-100)
mine = chunked_linear_cross_entropy(hidden, weight, targets, chunk=128)
assert torch.allclose(ref, mine, atol=1e-5), (ref.item(), mine.item())
print(f"naive: {ref.item():.6f}   chunked: {mine.item():.6f}")

The price is recomputing the lm_head matmul in backward: \(2 B T d V\) extra FLOPs, which for a 100M-class model is only a few percent on top of the \(6ND\) step cost — a very cheap trade for a 10× cut in peak memory.

In production you do not write this yourself. The libraries that implement it:

• Liger Kernel (liger-kernel) — Triton kernels for LLM training; from liger_kernel.transformers import LigerFusedLinearCrossEntropyLoss drops in for the head + loss and is wired into TRL and Axolotl by a single config flag.
• Cut Cross-Entropy (cut-cross-entropy, Wijmans et al.) — from cut_cross_entropy import linear_cross_entropy; computes the loss without ever materializing the logit matrix by fusing the matmul, the logsumexp, and the true-class gather in one kernel.
• torch.compile over a hand-written chunked loop like the one above, which is what TorchTitan-style stacks use when they want no extra dependency.
• Megatron-LM shards the vocabulary across tensor-parallel ranks and uses vocab_parallel_cross_entropy, which computes the log-partition with a single all-reduce of per-rank maxima and sums — so no rank ever holds the full \((B \cdot T, V)\) tensor (see Distributed Training II: Tensor, Pipeline, Sequence & Expert Parallelism).

Stack-100M uses the chunked path with loss_chunk = 8192, fused with its z-loss so the logsumexp is computed once and reused; the full kernel and its memory accounting are in The Pretraining Run: A Complete Single-GPU Training Loop.

Loss Masking¶

Why Mask?¶

Not all token positions carry equal pedagogical value. There are three major situations where we want to zero out (mask) certain positions in the loss:

1. Padding tokens. Sequences in a batch are padded to the same length. Padding positions have no linguistic content and including them would dilute the gradient.
2. Prompt tokens in supervised fine-tuning. During SFT or instruction tuning, we typically want the model to learn the response, not regurgitate the system prompt. (Covered in Supervised Fine-Tuning & Instruction Tuning.)
3. Document boundary tokens in packed sequences. When multiple documents are concatenated into a single long sequence (see §Packing below), we must prevent cross-document loss bleed.

Implementation¶

import torch
import torch.nn.functional as F

def compute_lm_loss_masked(
    logits: torch.Tensor,     # (B, T, V)
    tokens: torch.Tensor,     # (B, T+1) — full sequences
    mask: torch.Tensor,       # (B, T)   — 1 for positions to train on, 0 to skip
) -> torch.Tensor:
    """
    Masked causal language modeling loss.

    mask=0 at padding tokens, prompt tokens, or cross-document positions.
    We use ignore_index=-100 (PyTorch convention) to exclude masked positions.
    """
    targets = tokens[:, 1:].clone()  # (B, T) — shift targets

    # Replace masked positions with ignore_index so they contribute 0 to the loss
    targets[mask == 0] = -100        # -100 is the default ignore_index in F.cross_entropy

    B, T, V = logits.shape
    loss = F.cross_entropy(
        logits.reshape(B * T, V),
        targets.reshape(B * T),
        ignore_index=-100,
        reduction='mean',    # averages only over non-ignored positions
    )
    return loss


# Quick sanity check -------------------------------------------------------
torch.manual_seed(42)
B, T, V = 2, 8, 32000
logits  = torch.randn(B, T, V)
tokens  = torch.randint(0, V, (B, T + 1))
mask    = torch.ones(B, T, dtype=torch.long)
mask[0, 6:] = 0   # mask last 2 positions of first example (padding)
mask[1, :3] = 0   # mask first 3 positions of second example (prompt)

loss = compute_lm_loss_masked(logits, tokens, mask)
print(f"Masked loss: {loss.item():.4f}")  # a finite float; masked positions don't contribute

A subtle point: when using reduction='mean', PyTorch averages over non-ignored positions only. If you use reduction='sum' and divide manually, be careful to count only active tokens — dividing by the full sequence length (including masked positions) will produce a number that is systematically too small and will silently hurt training.

Document Packing¶

The Efficiency Problem¶

Real pretraining datasets contain documents of wildly varying lengths — a sentence here, a book chapter there. Padding every example to the longest sequence in a batch wastes enormous amounts of compute on meaningless tokens. For a mixed-length corpus, naive padding can waste 30–60% of tokens.

Document packing (also called sequence packing or bin packing) solves this by concatenating multiple short documents end-to-end until the combined sequence reaches the target context length \(L\). The concatenation uses a special separator token (e.g., <|endoftext|> in GPT-style models).

Context window = 1024 tokens

Unpacked (padded):
  [Doc A: 300 tokens | PAD x 724]       — 71% waste
  [Doc B: 512 tokens | PAD x 512]       — 50% waste
  [Doc C: 128 tokens | PAD x 896]       — 88% waste

Packed:
  [Doc A: 300 | SEP | Doc B: 512 | SEP | Doc C: 128 | PAD x 83]
  — only 8% waste

The Cross-Document Contamination Problem¶

Naive packing causes a subtle loss contamination issue: the model’s prediction of the first token of document B is conditioned on the final tokens of document A, which is semantically meaningless. This inflates the loss on document-boundary tokens and, more insidiously, teaches the model to expect arbitrary tokens as context — potentially hurting coherence of long-form generation.

There are three responses, in increasing order of rigor:

Option 0: Do nothing. Classic GPT-2/GPT-3-style pretraining simply concatenates documents with <|endoftext|> and trains on every position, letting the model learn that the separator means “the previous document ended, condition on nothing before it.” At a 1–2k context the contaminated positions are a small fraction of all tokens, and this is what nanoGPT and most small reproductions do.

Option 1: Loss masking at boundaries. Zero out the loss for the first token of each document in a packed sequence (its prediction is “poisoned” by the previous unrelated document). One line of code, no change to the attention kernel.

Option 2: Intra-document (block-diagonal) attention masking. Forbid attention from crossing a document boundary at all, so each document is trained exactly as if it had been the only thing in the context window. The Llama 3 report describes using such a mask and notes it mattered most during long-context training, where a single sequence packs many documents. Two production paths give you this without giving up a fast kernel:

• varlen FlashAttention. flash_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, causal=True) from the flash-attn package takes the packed batch flattened to (total_tokens, n_heads, head_dim) plus an int32 cu_seqlens array of cumulative document offsets, and runs each document as an independent attention problem at full FlashAttention speed.
• FlexAttention. In PyTorch ≥ 2.5, torch.nn.attention.flex_attention with create_block_mask(lambda b, h, q, kv: (q >= kv) & (doc_id[b, q] == doc_id[b, kv]), B, None, T, T) compiles a fused causal-and-same-document mask and skips blocks that are entirely masked, so the cost stays close to dense causal attention.

Megatron-LM exposes the same behaviour with --reset-attention-mask, and its companion --reset-position-ids restarts position indices at each <|endoftext|> so that RoPE does not give the second document in a pack an artificial positional offset — forgetting the position reset is a classic packing bug. The same machinery is reused for post-training packs (see Chat Templates, Data Formatting & Sequence Packing), and Stack-100M’s packed uint16 shards carry per-token seq_ids for exactly this purpose (see The Pretraining Run: A Complete Single-GPU Training Loop).

import torch

def build_packed_loss_mask(
    doc_ids: torch.Tensor,  # (B, T) — integer doc ID for each token position
) -> torch.Tensor:
    """
    Returns a loss mask (B, T) where position t is 1 (active) unless it is
    the first token of a new document (in which case its loss is contaminated
    by the previous document's context and should be excluded).

    doc_ids example for one sequence:
       [0, 0, 0, 1, 1, 1, 1, 2, 2]
    First positions of docs 1 and 2 (indices 3 and 7) get mask=0.
    """
    B, T = doc_ids.shape
    # A position starts a new document when its doc_id differs from the previous one
    # Position 0 is also the start of a document, but it has no "poisoned" context
    # so we keep it active (its input is just the BOS or the context start).
    mask = torch.ones(B, T, dtype=torch.long, device=doc_ids.device)

    # Detect document boundaries: where doc_id[t] != doc_id[t-1]
    # doc_ids[:, 1:] != doc_ids[:, :-1] gives True at boundary positions (t >= 1)
    boundary = (doc_ids[:, 1:] != doc_ids[:, :-1])  # (B, T-1)

    # The first token AFTER a boundary (i.e., position t where boundary[t-1] is True)
    # has its loss masked out. In the targets tensor (which is shifted by 1),
    # we mask the target at position t-1 when boundary[t-1] is True.
    # Equivalently: in the loss over targets[:, t], mask when doc changes at t.
    # Targets are tokens[:, 1:], so target[t] corresponds to predicting token t+1
    # from prefix up to token t. If token t+1 starts a new doc, mask it.
    new_doc_at_next = doc_ids[:, 1:] != doc_ids[:, :-1]  # (B, T-1): True when t+1 starts new doc
    mask[:, :-1][new_doc_at_next] = 0  # mask positions t where next token is a new doc

    return mask   # (B, T): 1 = train on this position, 0 = ignore

A Complete, Annotated Loss Computation¶

The following is a self-contained, runnable example that ties together all the pieces: packing, masking, and loss computation.

"""
Minimal pretraining loss pipeline.

Demonstrates:
  - document packing with SEP tokens
  - loss mask construction (exclude cross-doc boundaries and padding)
  - causal LM loss computation

Runnable with: python -c "exec(open('this_file.py').read())"
Requires: torch >= 2.0
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List, Tuple


SEP_TOKEN_ID = 2   # <|endoftext|> or equivalent separator
PAD_TOKEN_ID = 0
VOCAB_SIZE    = 256  # tiny vocab for illustration


def pack_documents(
    documents: List[List[int]],
    context_len: int,
    sep_id: int = SEP_TOKEN_ID,
    pad_id: int = PAD_TOKEN_ID,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Pack a list of token-ID lists into a single sequence of length context_len.
    Returns:
      tokens  : (context_len,)  padded token sequence
      doc_ids : (context_len,)  document ID per position (-1 for padding)
    """
    tokens  = []
    doc_ids = []
    doc_idx = 0

    for doc in documents:
        # Add SEP before each document (except the very first)
        if tokens:
            tokens.append(sep_id)
            doc_ids.append(doc_idx - 1)  # SEP belongs to the preceding doc
        for tok in doc:
            if len(tokens) >= context_len:
                break
            tokens.append(tok)
            doc_ids.append(doc_idx)
        doc_idx += 1
        if len(tokens) >= context_len:
            break

    # Pad to context_len
    pad_len = context_len - len(tokens)
    tokens  = tokens  + [pad_id] * pad_len
    doc_ids = doc_ids + [-1]     * pad_len   # -1 marks padding

    return (
        torch.tensor(tokens,  dtype=torch.long),
        torch.tensor(doc_ids, dtype=torch.long),
    )


def loss_mask_from_doc_ids(
    doc_ids: torch.Tensor,   # (T,)  — -1 for padding
) -> torch.Tensor:
    """
    Build loss mask of shape (T,).
    Active (1) unless:
      - padding position (doc_id == -1)
      - first token of a new document that follows a different document
        (cross-doc context contamination)
    """
    T = doc_ids.shape[0]
    mask = (doc_ids >= 0).long()   # 0 at padding, 1 elsewhere

    # Also zero out the target positions where the *next* token starts a new doc.
    # Target at position t is tokens[t+1]; if tokens[t+1] belongs to a new doc,
    # the model's context (tokens[:t+1]) is from the wrong doc, so mask it.
    for t in range(T - 1):
        if doc_ids[t] >= 0 and doc_ids[t + 1] >= 0 and doc_ids[t] != doc_ids[t + 1]:
            mask[t] = 0   # predicting the first token of doc[t+1] from doc[t] context
    return mask


class TinyTransformerLM(nn.Module):
    """A minimal decoder-only LM for illustration (not optimized for performance)."""

    def __init__(self, vocab_size: int, d_model: int = 64, n_heads: int = 4, n_layers: int = 2, max_seq_len: int = 128):
        super().__init__()
        self.embed   = nn.Embedding(vocab_size, d_model)
        self.pos_emb = nn.Embedding(max_seq_len, d_model)
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=n_heads, dim_feedforward=d_model * 4,
            batch_first=True, dropout=0.0,
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=n_layers)
        self.head = nn.Linear(d_model, vocab_size, bias=False)
        self.d_model = d_model

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """x: (B, T) → logits: (B, T, V)"""
        B, T = x.shape
        positions = torch.arange(T, device=x.device).unsqueeze(0)   # (1, T)
        h = self.embed(x) + self.pos_emb(positions)                 # (B, T, d_model)

        # Causal mask: upper-triangular with -inf (additive mask for PyTorch)
        causal_mask = nn.Transformer.generate_square_subsequent_mask(T, device=x.device)
        h = self.transformer(h, mask=causal_mask, is_causal=True)   # (B, T, d_model)
        return self.head(h)                                           # (B, T, V)


def compute_causal_lm_loss(
    model:  nn.Module,
    tokens: torch.Tensor,   # (B, T+1)
    masks:  torch.Tensor,   # (B, T)
) -> torch.Tensor:
    inputs  = tokens[:, :-1]    # (B, T) — model input
    targets = tokens[:, 1:].clone()   # (B, T) — what to predict

    # Mask out undesired positions
    targets[masks == 0] = -100  # PyTorch's ignore_index convention

    logits = model(inputs)      # (B, T, V)
    B, T, V = logits.shape

    loss = F.cross_entropy(
        logits.reshape(B * T, V),
        targets.reshape(B * T),
        ignore_index=-100,
    )
    return loss


# ---- Demo run ---------------------------------------------------------------
if __name__ == "__main__":
    torch.manual_seed(0)
    CONTEXT = 64

    # Simulate two documents with different lengths
    doc_a = list(range(10, 30))     # 20 tokens
    doc_b = list(range(50, 90))     # 40 tokens

    tokens_1d, doc_ids_1d = pack_documents([doc_a, doc_b], context_len=CONTEXT)
    mask_1d = loss_mask_from_doc_ids(doc_ids_1d)

    print(f"Tokens shape: {tokens_1d.shape}")
    print(f"Active positions: {mask_1d.sum().item()} / {CONTEXT}")
    print(f"Doc boundaries masked: {(mask_1d == 0).sum().item()} positions")

    # Batch of 2 sequences (in real training, batch of hundreds)
    tokens_batch   = tokens_1d.unsqueeze(0).expand(2, -1).clone()      # (2, 64)
    # mask_1d[t] flags whether the prediction of token t+1 (from doc[t]'s
    # context) is valid; it has one entry per input position (63 of them),
    # so the last entry (which has no "next token" in this array) is dropped
    # to align with targets = tokens_batch[:, 1:] of length T=63.
    mask_batch     = mask_1d[:-1].unsqueeze(0).expand(2, -1).clone()   # (2, 63)

    # tokens_batch has T+1=64 tokens; model sees first T=63, predicts last T=63
    model = TinyTransformerLM(vocab_size=VOCAB_SIZE, max_seq_len=CONTEXT)
    loss  = compute_causal_lm_loss(model, tokens_batch, mask_batch)

    print(f"Loss:       {loss.item():.4f} nats/token")
    print(f"Perplexity: {loss.exp().item():.2f}")
    # Expected: loss ≈ log(256) ≈ 5.55 for a random initialized model over 256-token vocab

Bits-per-Byte: A Tokenizer-Agnostic Metric¶

Perplexity as defined above is tokenizer-dependent. A model trained on a byte-pair encoding (BPE) tokenizer with vocabulary size 50 000 will report different perplexity from one with a 100 000-token vocabulary, even if they have identical predictive power over raw text. This makes cross-model comparisons treacherous.

Bits-per-byte (BPB) normalizes by the number of bytes (or characters) each token represents, yielding a tokenizer-independent measure.

where \(\mathcal{L}_{\text{NLL}}\) is the mean per-token NLL loss in nats, \(\log_2 e \approx 1.4427\) converts nats to bits, and \(\bar{r}\) is the average number of UTF-8 bytes per token for the tokenizer. Equivalently:

where the sum runs over all tokens and \(N_\text{bytes}\) is the total byte count of the text.

import math

def nats_per_token_to_bpb(
    loss_nats: float,        # average NLL in nats per token
    avg_bytes_per_token: float,   # tokenizer-specific compression ratio
) -> float:
    """Convert per-token NLL (nats) to bits-per-byte."""
    bits_per_token = loss_nats * math.log2(math.e)
    return bits_per_token / avg_bytes_per_token


def compute_avg_bytes_per_token(tokenizer, sample_texts: list[str]) -> float:
    """Estimate the average bytes-per-token ratio for a given tokenizer."""
    total_bytes  = 0
    total_tokens = 0
    for text in sample_texts:
        total_bytes  += len(text.encode("utf-8"))
        total_tokens += len(tokenizer.encode(text))
    return total_bytes / total_tokens


# Illustrative usage (requires `transformers` installed):
# from transformers import AutoTokenizer
# tok = AutoTokenizer.from_pretrained("gpt2")
# r = compute_avg_bytes_per_token(tok, ["Hello world.", "The cat sat on the mat."])
# bpb = nats_per_token_to_bpb(loss_nats=2.85, avg_bytes_per_token=r)

# Manual example matching the worked example above:
bpb = nats_per_token_to_bpb(2.85, 4.0)
print(f"BPB: {bpb:.3f}")   # → 1.028

You rarely compute this by hand for a released model: lm-evaluation-harness (EleutherAI) implements exactly this normalization for its rolling-loglikelihood tasks, which score a whole document by sliding the model over it and summing per-token log-probabilities.

pip install lm-eval
# wikitext is a `loglikelihood_rolling` task: it reports word_perplexity,
# byte_perplexity and bits_per_byte, all derived from the same summed NLL.
lm_eval --model hf \
        --model_args pretrained=./my-100m-checkpoint,dtype=bfloat16 \
        --tasks wikitext --batch_size 8

Because the harness divides the summed NLL by the raw byte count of the original text rather than by the token count, bits_per_byte is directly comparable across tokenizers — which is why it, and not perplexity, is the number to track when you swap your tokenizer or compare against a model with a different vocabulary. Harness internals and how to add your own task are covered in Building Eval Harnesses.

UL2 and Span Corruption Alternatives¶

Beyond Causal LM¶

The causal language modeling objective is the dominant pretraining strategy for decoder-only models, but it is not the only option. Encoder-decoder models (T5, mT5) were pretrained on masked span corruption (also called masked language modeling with contiguous spans), popularized by Raffel et al. in the T5 paper (2020).

In span corruption, a fraction of token spans in the input is replaced with sentinel tokens (e.g., <extra_id_0>), and the model must reconstruct the original spans:

Input:  "The quick <extra_id_0> over the <extra_id_1> dog."
Target: "<extra_id_0> brown fox jumps <extra_id_1> lazy <eos>"

The trade-off against causal LM runs the other way on token efficiency: only the masked-out tokens (typically 15% of the sequence) carry loss, so a span-corruption model extracts roughly an order of magnitude fewer supervised predictions per token of corpus. What it buys is that each prediction is conditioned bidirectionally — the encoder sees text on both sides of the span — which historically produced stronger representations for classification and extraction at the same parameter count.

The UL2 Family¶

UL2 (Tay et al., 2022, “Unifying Language Learning Paradigms”) showed that a single model pretrained on a mixture of different denoising objectives can match or outperform models trained on any single objective, while gaining versatility.

UL2 defines three classes of denoising modes, each selected at inference by a mode token prepended to the input:

During pretraining the model sees a mixture of all three modes with different sampling probabilities — this “mixture-of-denoisers” is the whole idea, and it is orthogonal to the architecture. UL2 20B (2022) reported outperforming T5-XXL (span corruption only) across most of its evaluation suite while matching or exceeding GPT-3 175B on several zero-shot benchmarks despite being far smaller — evidence that the objective mixture, not just scale, was doing work.

Flan-UL2 (2023) further instruction-tuned the UL2 checkpoint and was released publicly.

Prefix Language Modeling¶

A middle ground between masked LM and causal LM is the prefix LM (or non-causal prefix model): the input portion attends bidirectionally (full attention), while the output portion attends causally. This is the architecture of GLM (General Language Model, Du et al.) and PaLM’s initial pretraining variant. The loss is computed only on the output (continuation) portion.

Architecture variants and their relationship to these objectives are covered in depth in Architecture Variants: Encoder-Decoder, Decoder-Only & Prefix-LM.

Two Variants That Survived Into 2026: FIM and Multi-Token Prediction¶

Span corruption and UL2 largely lost to plain causal LM for decoder-only frontier models, but two modifications of next-token prediction are standard in 2026 base models, and both are data or head changes rather than new losses.

Fill-in-the-middle (FIM). A left-to-right model cannot natively condition on text that comes after the cursor — precisely what code completion in an editor requires. Bavarian et al. (2022) showed you can buy that ability essentially for free by rearranging a fraction of pretraining documents: split a document into prefix/middle/suffix, then train on prefix + suffix + middle with sentinel tokens marking the pieces. The loss is still ordinary next-token cross-entropy — only the token order changed. Trained at a FIM rate around 0.5–0.9, the model retains its left-to-right quality (the “FIM-for-free” property) while gaining infilling. Every serious code-capable base model — StarCoder, DeepSeek-Coder, the Qwen-Coder line — is trained this way.

import random

def fim_transform(doc_ids, fim_rate=0.5, sentinels=(50281, 50282, 50283), rng=random):
    """Rearrange a document into PSM (prefix-suffix-middle) FIM order with probability
    fim_rate; otherwise return it unchanged. sentinels = (<fim_prefix>, <fim_suffix>,
    <fim_middle>) — three IDs you must reserve in the tokenizer BEFORE pretraining.
    The training target is still just `tokens[1:]`: FIM is a data transform, not a loss."""
    if len(doc_ids) < 3 or rng.random() > fim_rate:
        return list(doc_ids)
    i, j = sorted(rng.sample(range(1, len(doc_ids)), 2))   # two cut points
    prefix, middle, suffix = doc_ids[:i], doc_ids[i:j], doc_ids[j:]
    pre, suf, mid = sentinels
    # At inference you feed everything up to and including <fim_middle> and decode.
    return [pre] + prefix + [suf] + suffix + [mid] + middle


demo = list(range(20))
out = fim_transform(demo, fim_rate=1.0, rng=random.Random(0))
print(out[:3], "...", len(out))   # 3 sentinels + 20 original tokens = 23
assert len(out) == len(demo) + 3

Multi-token prediction (MTP). Instead of one head predicting \(x_{t+1}\), attach \(n\) heads (or \(n\) lightweight modules) on the shared trunk and predict \(x_{t+1}, \ldots, x_{t+n}\), summing the cross-entropies. Gloeckle et al. (2024) showed this improves downstream generation and coding at scale; DeepSeek-V3 adopted an MTP objective as an auxiliary loss with a small weight, and — the practical kicker — reused the trained extra head as a draft head for speculative decoding at inference, so an objective change bought a decoding speedup for free (see Speculative Decoding: Draft Models, Medusa, EAGLE & Lookahead). Note that MTP multiplies the loss-head memory discussed above by \(n\), which is another reason chunked cross-entropy is not optional at scale.

Practical Loss Engineering: What Practitioners Actually Do¶

Loss Normalization Across Heterogeneous Batches¶

In large-scale distributed training (see Distributed Training I: Data Parallelism, DDP, ZeRO & FSDP), each GPU holds a micro-batch. Two normalization choices are common:

Token-normalized loss (recommended): $$ \mathcal{L} = \frac{\sum_{b,t} \mathbb{1}[\text{active}{b,t}] \cdot \ell $$}}{\sum_{b,t} \mathbb{1}[\text{active}_{b,t}]

This divides by the count of active (unmasked) tokens and is independent of batch construction.

Sequence-normalized loss (legacy): $$ \mathcal{L} = \frac{1}{B} \sum_{b} \frac{1}{T_b} \sum_{t} \ell_{b,t} $$

This averages per-sequence first, then averages sequences. It implicitly weights short sequences more heavily — usually undesired.

Label Smoothing¶

Label smoothing replaces the one-hot target with a softer distribution: instead of probability 1 on the true token, the true token gets \(1 - \epsilon\) and each other token gets \(\epsilon / (V-1)\):

With \(\epsilon = 0.1\), this has been shown to improve calibration and slightly regularize training. PyTorch’s F.cross_entropy supports it directly via the label_smoothing argument. Many large LM runs skip label smoothing or use very small \(\epsilon \le 0.05\), since at scale the training signal is already rich enough.

Z-Loss for Softmax Stability¶

At scale, the logits fed to the softmax can grow very large, causing numerical instability (the softmax exponentials overflow in float16). One solution is z-loss (Chowdhery et al., PaLM, 2022), which adds a regularizer to penalize large logit norms:

where \(z_v\) are the pre-softmax logits. With a small coefficient \(\alpha\) (e.g., \(10^{-4}\)), this penalizes large log-partition values without significantly affecting the primary loss. It dramatically reduces loss spikes during training, as documented in the PaLM technical report.

def z_loss(logits: torch.Tensor, alpha: float = 1e-4) -> torch.Tensor:
    """
    Z-loss regularizer for softmax stability (PaLM / Chowdhery et al. 2022).

    logits : (*, V) — pre-softmax logit tensor
    alpha  : coefficient; typical value 1e-4 to 1e-5

    Returns a scalar to be added to the primary cross-entropy loss.
    """
    # log(sum_v exp(z_v)) = log-sum-exp, numerically stable via torch.logsumexp
    log_z = torch.logsumexp(logits, dim=-1)   # (*)  — one value per token
    return alpha * (log_z ** 2).mean()


# Integrate with main loss:
# total_loss = ce_loss + z_loss(logits)

Training stability issues and how loss diagnostics help debug them are covered in Training Stability, Loss Spikes & Debugging Large Runs.

Interview Corner¶

Key Takeaways¶

Further Reading¶

• Raffel et al., “Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer” (T5), JMLR 2020. Introduces span-corruption pretraining and systematically compares objectives, architectures, and data scales.
• Tay et al., “Unifying Language Learning Paradigms” (UL2), ICLR 2023. Defines the R/X/S-denoising taxonomy and shows that a single model benefits from a mixture of objectives.
• Brown et al., “Language Models are Few-Shot Learners” (GPT-3), NeurIPS 2020. Demonstrates that scale of causal LM pretraining alone produces powerful few-shot models, establishing the primacy of next-token prediction for general-purpose LLMs.
• Chowdhery et al., “PaLM: Scaling Language Modeling with Pathways,” JMLR 2023. Documents z-loss and other engineering choices that stabilize cross-entropy at large scale.
• Hoffmann et al., “Training Compute-Optimal Large Language Models” (Chinchilla), NeurIPS 2022. Uses the cross-entropy loss as the primary dependent variable to derive scaling laws; directly connected to Scaling Laws: Kaplan, Chinchilla & Beyond.
• Karpathy, “nanoGPT” (GitHub, 2022–present). A ~300-line clean implementation of the full pretraining loop including the loss computation; the best codebase to read alongside this chapter.
• Shannon, “A Mathematical Theory of Communication,” Bell System Technical Journal, 1948. The original derivation of entropy as a lower bound on compression, foundational to understanding what cross-entropy loss measures.

Exercises¶

1. Perplexity intuition. A randomly-initialized model over a vocabulary of size \(V = 32000\) has just been created; no training has happened. Assuming the softmax outputs are approximately uniform over the vocabulary at initialization, what per-token NLL loss (in nats) and what perplexity do you expect on the first batch? Then explain in one sentence why the chapter’s from-scratch verification snippet actually prints \(\approx 11.05\) rather than exactly \(\log V\).

2. Label shifting and tensor shapes. In the chapter’s compute_lm_loss, the input is tokens of shape (B, T+1). Explain precisely why the model is fed tokens[:, :-1] and supervised with tokens[:, 1:], and state how many gradient-carrying predictions a single unpadded sequence of length \(T+1 = 2048\) produces. What goes wrong if you instead (incorrectly) align the logits with tokens[:, :-1] as targets?

3. Bits-per-byte comparison. Two models are evaluated on the same English text.

- Model A: BPE tokenizer, NLL $= 3.10$ nats/token, average $4.2$ UTF-8 bytes/token.
- Model B: larger-vocab tokenizer, NLL $= 3.55$ nats/token, average $5.4$ UTF-8 bytes/token.

Compute the bits-per-byte for each (use \(\log_2 e \approx 1.4427\)) and state which model is better in a tokenizer-agnostic sense. Then explain in one sentence why comparing their raw per-token NLL values directly would be misleading.

4. Z-loss arithmetic and purpose. Consider a single token whose pre-softmax logits, before any normalization, all share a common additive offset \(c\) — i.e. the logits are \(z_v = a_v + c\) for some fixed base pattern \(a_v\). (a) Show that the cross-entropy loss for the true token is invariant to \(c\). (b) Show that the z-loss term \(\alpha \log^2\!\big(\sum_v e^{z_v}\big)\) is not invariant to \(c\), and compute the z-loss for \(\alpha = 10^{-4}\) when \(\log\sum_v e^{z_v} = 30\). © In one sentence, explain why penalizing this quantity improves numerical stability.

5. Implement per-document (sequence-normalized) loss and compare to token-normalized loss. The chapter contrasts token-normalized and sequence-normalized loss and warns that sequence-normalization “implicitly weights short sequences more heavily.” Implement a function sequence_normalized_loss(logits, tokens, mask) that computes the loss by first averaging over the active positions within each sequence, then averaging those per-sequence means across the batch. Then, using the chapter’s masking conventions, construct a 2-sequence batch where one sequence has far fewer active tokens than the other and demonstrate numerically that the sequence-normalized value differs from the standard token-normalized compute_lm_loss_masked.

import torch
import torch.nn.functional as F

def sequence_normalized_loss(
    logits: torch.Tensor,   # (B, T, V)
    tokens: torch.Tensor,   # (B, T+1)
    mask:   torch.Tensor,   # (B, T) — 1 = active, 0 = ignore
) -> torch.Tensor:
    targets = tokens[:, 1:].clone()          # (B, T)
    B, T, V = logits.shape

    # Per-position NLL with no reduction, so we can control the averaging.
    per_pos = F.cross_entropy(
        logits.reshape(B * T, V),
        targets.reshape(B * T),
        reduction='none',                    # (B*T,)
    ).reshape(B, T)                          # (B, T)

    m = mask.to(per_pos.dtype)               # (B, T)
    # Sum active loss per sequence, divide by active count per sequence.
    seq_active = m.sum(dim=1).clamp_min(1.0)         # (B,)
    seq_mean   = (per_pos * m).sum(dim=1) / seq_active   # (B,)
    # Then average the per-sequence means with equal weight.
    return seq_mean.mean()

def compute_lm_loss_masked(logits, tokens, mask):
    targets = tokens[:, 1:].clone()
    targets[mask == 0] = -100
    B, T, V = logits.shape
    return F.cross_entropy(
        logits.reshape(B * T, V),
        targets.reshape(B * T),
        ignore_index=-100,
        reduction='mean',                    # token-normalized
    )

# ---- Demonstration -----------------------------------------------------
torch.manual_seed(0)
B, T, V = 2, 8, 32000
logits = torch.randn(B, T, V)
tokens = torch.randint(0, V, (B, T + 1))

mask = torch.ones(B, T, dtype=torch.long)
mask[0, 2:] = 0   # sequence 0: only 2 active tokens (a short doc)
mask[1, :]  = 1   # sequence 1: all 8 active tokens (a long doc)

tok_norm = compute_lm_loss_masked(logits, tokens, mask)
seq_norm = sequence_normalized_loss(logits, tokens, mask)

print(f"token-normalized:    {tok_norm.item():.4f}")
print(f"sequence-normalized: {seq_norm.item():.4f}")
# The two values differ: the short 2-token sequence gets weight 1/2 under
# sequence-normalization but only 2/10 of the tokens under token-normalization.

6. Budgeting the loss head. You are training a 100M-parameter model with \(d_{\text{model}} = 512\) and \(V = 32768\), using a micro-batch of \(B \times T = 32 \times 2048\) tokens on a 40 GB GPU. (a) Compute the peak bytes of the naive path: bf16 logits, their fp32 upcast, and the fp32 log_softmax saved for backward. (b) Compute the peak logit memory of the chunked path with chunk = 8192 (fp32 logits plus their gradient). © Compute the extra FLOPs the chunked path spends recomputing the lm_head matmul in backward, as a fraction of the \(6ND\) step cost. (d) Why does activation checkpointing on the transformer blocks not help with (a)?