How LLM Training Actually Works: From Tokenization to Loss Convergence

Training a modern large language model is a rigid, memory-bound distributed optimization process. It converts hundreds of terabytes of unstructured text into dense floating-point weight matrices across thousands of interconnected GPUs operating in tight synchronization.

Every step in the training pipeline is governed by physical and mathematical constraints: memory bandwidth limits on High Bandwidth Memory (HBM), network latency across InfiniBand fabrics, numerical range boundaries of floating-point representations, and numerical stability during gradient updates.

This guide walks through the exact mechanics of how large language models are trained from scratch. We cover data extraction, Byte-Pair Encoding (BPE) vocabulary construction and priority queue data structures, weight initialization, multi-head attention forward passes, FlashAttention-2 online softmax tiling, cross-entropy loss computation, reverse-mode automatic differentiation derivations for Softmax and RMSNorm, AdamW optimizer state accounting and weight decay decoupling, 3D cluster parallelism (Tensor, Pipeline, and Fully Sharded Data Parallelism), FP16 and BF16 numerical stability, and post-training alignment through Supervised Fine-Tuning (SFT) and Direct Preference Optimization (DPO).

Dataset Curation and Byte-Pair Encoding

The training process begins long before tensors touch a GPU. Modern pre-training datasets comprise several trillion tokens (for instance, 3 to 15 trillion tokens for state-of-the-art open models). The raw data is scraped from web crawls (such as Common Crawl), public code repositories, digital libraries, scientific archives, and encyclopedias.

Data Extraction and Quality Filtering

Raw web text contains HTML boilerplate, navigation headers, advertising copy, machine-generated spam, OCR corruption, and duplicated text. Ingesting uncurated text causes loss divergence, severe model hallucinations, or degradation in reasoning performance.

The data pipeline runs through several strict transformation stages:

1. Text Extraction and Normalisation: Raw HTML is parsed to strip markup tags while preserving paragraph structure. Unicode characters are normalized (typically using NFC or NFKC forms), and invalid byte sequences are discarded.
2. Language Identification: A fast text classification model (such as fastText) computes probability scores over language codes. Non-target languages or documents with low classification confidence are purged.
3. Quality and Heuristic Filtering: Rule-based heuristics discard low-quality documents. Thresholds are applied to word length distribution, mean line length, symbol-to-word ratios (discarding text where more than 10% of characters are non-alphanumeric punctuation), and stop-word counts.
4. Perplexity Filtering: A lightweight reference language model (such as a 5-gram language model trained on clean Wikipedia and textbook datasets) calculates the perplexity of incoming text blocks. Documents exhibiting excessively high perplexity (indicating ungrammatical gibberish or garbage encoding) or abnormally low perplexity (indicating repetitive boilerplate or domain lists) are removed.

+------------------+    +-------------------+    +----------------------+
| Raw Web Crawls   | -> | Text Extraction   | -> | Quality Heuristics   |
| (Common Crawl)   |    | (HTML stripping)  |    | (Symbol ratios, etc) |
+------------------+    +-------------------+    +----------------------+
                                                            |
                                                            v
+------------------+    +-------------------+    +----------------------+
| Tokenized Corpus | <- | MinHash LSH       | <- | Perplexity Filter    |
| (BPE Binary File)|    | Deduplication     |    | (5-gram KenLM Model) |
+------------------+    +-------------------+    +----------------------+

Deduplication via MinHash and Locality-Sensitive Hashing

Duplicate text distorts the empirical data distribution. If a news article or boilerplate disclaimers appear millions of times in the training set, the model overfits to those exact sequences, resulting in memorization and degraded generalization.

Deduplication occurs at both exact match and near-duplicate levels:

• Exact Deduplication: Exact string hashes (SHA-256 or MD5) are computed over line, paragraph, or document contents. Identical matches are removed.
• Fuzzy Deduplication (MinHash + LSH): To catch documents that differ by only a few words or formatting changes, datasets use Locality-Sensitive Hashing (LSH) over $k$-shingle sets.

For a document $D$, a set of character or word $k$-shingles is generated. A collection of $N$ distinct hash functions $h_1, h_2, \dots, h_N$ is applied to all shingles in the document. The minimum hash value for each hash function forms the document MinHash signature:

$$\text{Sig}(D) = \left[ \min_{s \in D} h_1(s), \min_{s \in D} h_2(s), \dots, \min_{s \in D} h_N(s) \right]$$

The probability that two documents $A$ and $B$ produce identical MinHash values for a given hash function equals their Jaccard similarity:

$$P(\min h(A) = \min h(B)) = J(A, B) = \frac{|A \cap B|}{|A \cup B|}$$

To scale signature comparison across billions of documents, signatures are partitioned into $b$ bands of $r$ rows each ($N = b \times r$). Documents that match identically in all $r$ entries of at least one band are hashed into the same LSH bucket as candidate duplicates. Candidate pairs exceeding a target Jaccard threshold (typically $J(A, B) \ge 0.8$) are purged from the corpus.

Pre-Tokenization Regex Splitting and Vocabulary Construction

Neural networks process vectors, not raw text strings. The tokenizer maps text into sequences of discrete integer token IDs drawn from a fixed vocabulary of size $V$ (typically between 32,768 and 128,000).

Modern LLMs utilize Byte-Pair Encoding (BPE). Before pair counting begins, production tokenizers (such as OpenAI tiktoken or Meta Llama tokenizers) run a pre-tokenization regex split. The pre-tokenization step prevents merges from crossing syntactic and whitespace boundaries (for example, stopping a space and a punctuation mark from merging into a single token across word boundaries).

A standard production BPE regex pattern enforces category isolation:

's|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+

This pattern breaks text into distinct groups:

1. Contraction suffixes ('s, 't, 're, etc.).
2. Letter sequences with an optional preceding space ( ?\p{L}+).
3. Numeric sequences with an optional preceding space ( ?\p{N}+).
4. Punctuation and non-alphanumeric character clusters ( ?[^\s\p{L}\p{N}]+).
5. Standalone whitespace runs (\s+).

The BPE algorithm starts with base byte characters (256 initial tokens representing raw bytes 0x00 through 0xFF). This byte fallback guarantees that the tokenizer can encode any arbitrary binary stream or UTF-8 sequence without producing out-of-vocabulary (<unk>) tokens.

BPE Merge Frequency Data Structures: $O(N \log V)$ Optimization

Naive implementation of BPE merge counting scans the tokenized dataset of size $N$ repeatedly for each merge iteration. For a vocabulary size $V$, this results in an unscalable $O(N \cdot V)$ time complexity.

Production BPE training pipelines achieve $O(N \log V)$ computational complexity by maintaining two co-dependent data structures:

1. Max-Heap (Priority Queue): Stores pairs of adjacent tokens indexed by their current frequency count. The heap allows $O(1)$ lookup of the most frequent pair and $O(\log K)$ frequency update operations (where $K$ is the number of unique active pairs).
2. Doubly-Linked Position Index: Stores the positions of tokens in the corpus as a doubly-linked list alongside an inverted hash index mapping every pair $(t_a, t_b)$ to its set of occurrence locations.

When the top pair $(t_i, t_j)$ is popped from the max-heap and merged into new token ID $t_{\text{new}}$, the tokenizer looks up all occurrences of $(t_i, t_j)$ using the inverted location index. At each location, neighboring token pair counts are updated in the max-heap in $O(\log K)$ time.

import heapq
from collections import defaultdict
 
class BPETrainer:
    def __init__(self, vocab_size: int):
        self.vocab_size = vocab_size
        self.vocab = {i: bytes([i]) for i in range(256)}
        self.merges = {}
 
    def train(self, corpus_bytes: bytes):
        tokens = list(corpus_bytes)
        pair_counts = defaultdict(int)
        pair_positions = defaultdict(set)
        
        for idx in range(len(tokens) - 1):
            pair = (tokens[idx], tokens[idx + 1])
            pair_counts[pair] += 1
            pair_positions[pair].add(idx)
            
        heap = [(-count, pair) for pair, count in pair_counts.items()]
        heapq.heapify(heap)
        
        next_token_id = 256
        while next_token_id < self.vocab_size and heap:
            neg_count, best_pair = heapq.heappop(heap)
            current_count = -neg_count
            
            if pair_counts[best_pair] != current_count or current_count == 0:
                continue
                
            self.merges[best_pair] = next_token_id
            self.vocab[next_token_id] = self.vocab[best_pair[0]] + self.vocab[best_pair[1]]
            pair_counts[best_pair] = 0
            next_token_id += 1
 
        return self.merges

During tokenization, special control tokens are added to the vocabulary, including sequence boundaries (<|endoftext|>), instruction identifiers (<|im_start|>, <|im_end|>), and structural padding markers.

Transformer Weight Initialization and Forward Pass Computation

Once the dataset is tokenized, text is structured into fixed-length sequences of length $T$ (the context window, e.g., $T = 4096$ or $8192$) and batched into matrices of shape $(B, T)$, where $B$ is the micro-batch size per GPU.

Input Token IDs: [B, T]
       |
       v
Embedding Matrix (E): Lookup -> [B, T, d_model]
       |
       +---> For each of L Layers:
       |        |
       |        v
       |     RMSNorm Layer
       |        |
       |        v
       |     Multi-Head / Grouped-Query Attention (MHA/GQA) + RoPE
       |        |
       |        v
       |     Residual Addition
       |        |
       |        v
       |     RMSNorm Layer
       |        |
       |        v
       |     SwiGLU MLP Block
       |        |
       |        v
       |     Residual Addition
       |
       v
Final RMSNorm -> LM Head Projection [W_lm_head] -> Logits: [B, T, V]

Parameter Initialization Distributions

Before training begins, model parameters must be initialized to ensure stable variance during initial forward and backward passes. Unscaled initialization causes activations to explode exponentially with layer depth or vanish to zero.

• Embedding Matrices: Initialized from a standard Gaussian distribution $\mathcal{N}(0, \sigma^2)$ where $\sigma = \frac{1}{\sqrt{d_{\text{model}}}}$.
• Linear Projection Matrices ($W_Q, W_K, W_V, W_1, W_3$): Initialized using Xavier (Glorot) or Kaiming (He) normal distributions:

$$W \sim \mathcal{N}\left(0, \frac{2}{d_{\text{in}} + d_{\text{out}}}\right)$$

• Residual Projection Output Matrices ($W_O, W_2$): To prevent residual stream variance from expanding linearly with the number of layers $L$, output projections within attention and MLP blocks are scaled down by $\frac{1}{\sqrt{2L}}$:

$$W_{\text{out}} \sim \mathcal{N}\left(0, \frac{1}{2L \cdot d_{\text{in}}}\right)$$

Rotary Position Embeddings (RoPE)

Because attention operations are permutation-invariant, positional information must be injected. Modern decoder architectures use Rotary Position Embeddings (RoPE).

Instead of adding static position vectors to embeddings, RoPE rotates the Query and Key projection vectors in 2D vector slices by a position-dependent angle matrix. For a vector $\mathbf{x} \in \mathbb{R}^{d_{\text{head}}}$ at token index $m$:

$$R_{\Theta, m}^d \mathbf{x}_m = \begin{pmatrix} \mathbf{x}_m^{(1)} \cos m\theta_1 - \mathbf{x}_m^{(2)} \sin m\theta_1 \ \mathbf{x}_m^{(1)} \sin m\theta_1 + \mathbf{x}_m^{(2)} \cos m\theta_1 \ \vdots \ \mathbf{x}m^{(d-1)} \cos m\theta{d/2} - \mathbf{x}m^{(d)} \sin m\theta{d/2} \ \mathbf{x}m^{(d-1)} \sin m\theta{d/2} + \mathbf{x}m^{(d)} \cos m\theta{d/2} \end{pmatrix}$$

where $\theta_i = 10000^{-2(i-1)/d}$. RoPE ensures that the inner product between rotated Query $\mathbf{q}_m$ and Key $\mathbf{k}_n$ depends purely on the relative offset $(m - n)$:

$$\langle R_{\Theta, m}^d \mathbf{q}m, R{\Theta, n}^d \mathbf{k}_n \rangle = \mathbf{q}m^T R{\Theta, n-m}^d \mathbf{k}_n$$

Normalization Layers: LayerNorm vs RMSNorm

Standard LayerNorm normalizes inputs across the feature dimension using both mean and variance:

$$\text{LayerNorm}(x) = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} \odot \gamma + \beta$$

where $\mu = \frac{1}{d}\sum_{i=1}^d x_i$ and $\sigma^2 = \frac{1}{d}\sum_{i=1}^d (x_i - \mu)^2$.

Modern LLMs substitute LayerNorm with Root Mean Square Normalization (RMSNorm). RMSNorm drops the mean calculation $\mu$ and shift parameter $\beta$, normalizing strictly by root mean square magnitude:

$$\text{RMSNorm}(x) = \frac{x}{\text{RMS}(x)} \odot \gamma, \quad \text{where } \text{RMS}(x) = \sqrt{\frac{1}{d} \sum_{i=1}^d x_i^2 + \epsilon}$$

Discarding mean centering reduces memory overhead and reduces GPU memory access instructions without degrading training convergence.

Attention Computation and Causal Masking

For a hidden state input $H \in \mathbb{R}^{B \times T \times d_{\text{model}}}$, linear projection matrices project $H$ into Query, Key, and Value matrices:

$$Q = H W_Q, \quad K = H W_K, \quad V = H W_V$$

In Grouped-Query Attention (GQA), $Q$ has $H_Q$ heads while $K$ and $V$ share $H_{KV}$ heads ($H_{KV} < H_Q$) to conserve memory during inference and training.

The scaled dot-product attention per head computes:

$$\text{Attention}(Q, K, V) = \text{softmax}\left( \frac{Q K^T}{\sqrt{d_k}} + M \right) V$$

where $d_k = d_{\text{model}} / H_Q$. $M$ is the causal attention mask that enforces autoregressive bounds. To prevent token at index $t$ from attending to future tokens $t' > t$:

$$M_{t, t'} = \begin{cases} 0 & \text{if } t' \le t \ -\infty & \text{if } t' > t \end{cases}$$

Adding $-\infty$ prior to softmax zeroes out the attention probabilities for all subsequent positions: $\exp(-\infty) = 0$.

FlashAttention-2 Memory-Efficient Online Softmax Tiling

Standard attention computation materializes the full $T \times T$ score matrix $S = \frac{Q K^T}{\sqrt{d_k}}$ and attention probability matrix $P = \text{softmax}(S)$ in High Bandwidth Memory (HBM). For context length $T = 8192$, materializing $S$ requires gigabytes of HBM read/write throughput per layer, causing severe memory bandwidth bottlenecks.

FlashAttention-2 eliminates $O(T^2)$ memory reads and writes by executing attention using online softmax tiling inside fast GPU SRAM (20 TB/s bandwidth on NVIDIA H100 vs 3.35 TB/s HBM3).

+-----------------------------------------------------------------+
|                       GPU HBM (3.35 TB/s)                       |
|           Q, K, V Tensors stored in global memory               |
+-----------------------------------------------------------------+
                 |                              ^
     Load Blocks |                              | Write Final Output O
   Br x d, Bc x d|                              |
                 v                              |
+-----------------------------------------------------------------+
|                       GPU SRAM (20 TB/s)                        |
|                                                                 |
|   1. Tile Q into blocks Q_i of size Br x d                      |
|   2. Tile K, V into blocks K_j, V_j of size Bc x d              |
|   3. Compute S_i^(j) = Q_i * (K_j)^T / sqrt(d_k)                |
|   4. Update running row-max m_i and running sum l_i             |
|   5. Rescale and accumulate partial output O_i                  |
+-----------------------------------------------------------------+

Online Softmax Algorithm Equations

The key mathematical identity behind FlashAttention is that Softmax can be computed incrementally without seeing all entries at once.

For a query block $Q_i \in \mathbb{R}^{B_r \times d_k}$ and key/value blocks $K_j, V_j \in \mathbb{R}^{B_c \times d_k}$:

1. Compute block dot-product score matrix:

$$S_i^{(j)} = \frac{Q_i K_j^T}{\sqrt{d_k}} \in \mathbb{R}^{B_r \times B_c}$$

2. Compute row-wise maximum of the current block:

$$\tilde{m}_i^{(j)} = \text{rowmax}\left(S_i^{(j)}\right) \in \mathbb{R}^{B_r}$$

3. Update the global running row-maximum:

$$m_i^{(j)} = \max\left(m_i^{(j-1)}, \tilde{m}_i^{(j)}\right)$$

4. Compute the unnormalized exponent matrix for the current block:

$$P_i^{(j)} = \exp\left(S_i^{(j)} - m_i^{(j)}\right) \in \mathbb{R}^{B_r \times B_c}$$

5. Update the running normalization row-sum factor $l_i^{(j)}$:

$$l_i^{(j)} = e^{m_i^{(j-1)} - m_i^{(j)}} \odot l_i^{(j-1)} + \text{rowsum}\left(P_i^{(j)}\right)$$

6. Rescale previous partial output accumulator $O_i^{(j-1)}$ and add the current block contribution:

$$O_i^{(j)} = \text{diag}\left(e^{m_i^{(j-1)} - m_i^{(j)}}\right)^{-1} O_i^{(j-1)} + P_i^{(j)} V_j$$

After iterating through all key/value blocks $j = 1 \dots T_c$ (where $T_c = T / B_c$), the final output block is normalized:

$$O_i = \text{diag}\left(l_i^{(T_c)}\right)^{-1} O_i^{(T_c)}$$

Backward Pass Activation Recomputation

FlashAttention-2 does not store the $T \times T$ attention score matrix $P$ for the backward pass. Instead, it stores only the scalar statistics $m_i$ and $l_i$ (size $O(T)$) in HBM. During backpropagation, FlashAttention recomputes the score blocks $S_i^{(j)}$ and probabilities $P_i^{(j)}$ on the fly in SRAM, reducing activation memory footprint by up to 95% while speeding up execution due to reduced HBM read/write traffic.

SwiGLU Feed-Forward Block

After the attention output is projected by $W_O$ and added back to the residual stream via RMSNorm, the activation vector passes through an MLP block. Modern architectures use the SwiGLU (Swish Gated Linear Unit) variant:

$$\text{SwiGLU}(H) = \left( \text{Swish}(H W_1) \otimes (H W_3) \right) W_2$$

where $\text{Swish}(z) = z \cdot \sigma(z)$, and $\otimes$ denotes element-wise Hadamard multiplication. The intermediate dimension $d_{\text{ff}}$ is typically set to $\frac{8}{3} d_{\text{model}}$, rounded up to a multiple of 256 for GPU memory alignment.

Logits and Cross-Entropy Loss Computation

After passing through $L$ transformer layers, the final hidden state $H_L \in \mathbb{R}^{B \times T \times d_{\text{model}}}$ is normalized via RMSNorm and projected to vocabulary space using the Language Modeling (LM) Head matrix $W_{\text{lm_head}} \in \mathbb{R}^{d_{\text{model}} \times V}$:

$$Z = H_L W_{\text{lm_head}} \in \mathbb{R}^{B \times T \times V}$$

$Z_{b, t, k}$ represents the unnormalized score (logit) for token $k$ at position $t$ in batch sample $b$.

The model converts logits to probability predictions via the Softmax function:

$$P(y_{b, t} = k \mid x_{b, <t}) = \frac{\exp(Z_{b, t, k})}{\sum_{j=1}^V \exp(Z_{b, t, j})}$$

To avoid numerical overflow during $\exp(Z_{b, t, k})$, the Log-Sum-Exp trick subtracts the maximum logit $m = \max_j Z_{b, t, j}$ prior to exponentiation:

$$\log \sum_{j=1}^V \exp(Z_{b, t, j}) = m + \log \sum_{j=1}^V \exp(Z_{b, t, j} - m)$$

The pre-training objective is next-token prediction over the target sequence $Y$ (which is sequence $X$ shifted left by one token). The cross-entropy loss $\mathcal{L}$ averaged over batch size $B$ and context length $T$ is formulated as:

$$\mathcal{L} = -\frac{1}{B \cdot T} \sum_{b=1}^B \sum_{t=1}^T \log P(y_{b, t} = x_{b, t+1} \mid x_{b, \le t})$$

Backward Pass, Gradients, and Optimizer Mechanics

Once the forward pass yields the scalar loss $\mathcal{L}$, the training runtime executes a backward pass using reverse-mode automatic differentiation.

Step-by-Step Mathematical Derivation of Softmax Cross-Entropy Gradients

To compute updates for the model weights, we derive the gradient of scalar loss $\mathcal{L}$ with respect to raw input logits $Z$.

Consider a single sequence position with logit vector $z = [z_1, z_2, \dots, z_V]^T \in \mathbb{R}^V$ and true target token index $y \in {1, \dots, V}$.

The softmax probability assigned to index $k$ is:

$$p_k = \frac{e^{z_k}}{\sum_{j=1}^V e^{z_j}}$$

The cross-entropy loss for this token position is:

$$\mathcal{L} = -\log p_y = -\left( z_y - \log \sum_{j=1}^V e^{z_j} \right) = \log \sum_{j=1}^V e^{z_j} - z_y$$

We evaluate partial derivative $\frac{\partial \mathcal{L}}{\partial z_i}$ for an arbitrary logit index $i \in {1, \dots, V}$.

Case 1: $i \neq y$ (Non-target logit index)

$$\frac{\partial \mathcal{L}}{\partial z_i} = \frac{\partial}{\partial z_i} \left( \log \sum_{j=1}^V e^{z_j} - z_y \right) = \frac{1}{\sum_{j=1}^V e^{z_j}} \cdot \frac{\partial}{\partial z_i} \left( \sum_{j=1}^V e^{z_j} \right) - 0 = \frac{e^{z_i}}{\sum_{j=1}^V e^{z_j}} = p_i$$

Case 2: $i = y$ (Target logit index)

$$\frac{\partial \mathcal{L}}{\partial z_y} = \frac{\partial}{\partial z_y} \left( \log \sum_{j=1}^V e^{z_j} - z_y \right) = \frac{e^{z_y}}{\sum_{j=1}^V e^{z_j}} - 1 = p_y - 1$$

Combining both cases using one-hot target vector $y_{\text{onehot}} \in \mathbb{R}^V$ yields the exact gradient vector expression:

$$\frac{\partial \mathcal{L}}{\partial z} = p - y_{\text{onehot}}$$

In matrix notation across batch size $B$ and sequence context length $T$:

$$\nabla_Z \mathcal{L} = \frac{1}{B \cdot T} \left( P - Y_{\text{onehot}} \right)$$

This expression shows that the error signal fed back into the network is the difference between predicted probabilities $P$ and ground-truth targets $Y_{\text{onehot}}$.

Step-by-Step Derivation of RMSNorm Backward Pass

During backpropagation, gradients must propagate backward through RMSNorm layers.

For input vector $x \in \mathbb{R}^d$ and gain parameter $\gamma \in \mathbb{R}^d$, the forward pass is:

$$\bar{x} = \text{RMS}(x) = \sqrt{\frac{1}{d} \sum_{k=1}^d x_k^2 + \epsilon}, \quad y_i = \frac{x_i}{\bar{x}} \gamma_i$$

Let $\hat{g}_i = \frac{\partial \mathcal{L}}{\partial y_i}$ be the incoming gradient from downstream operations. By the multivariable chain rule:

$$\frac{\partial \mathcal{L}}{\partial x_i} = \sum_{j=1}^d \frac{\partial \mathcal{L}}{\partial y_j} \frac{\partial y_j}{\partial x_i} = \sum_{j=1}^d \hat{g}_j \frac{\partial y_j}{\partial x_i}$$

Using the product rule on $y_j = x_j \bar{x}^{-1} \gamma_j$:

$$\frac{\partial y_j}{\partial x_i} = \frac{\partial x_j}{\partial x_i} \bar{x}^{-1} \gamma_j + x_j \gamma_j \frac{\partial \bar{x}^{-1}}{\partial x_i} = \delta_{ij} \frac{\gamma_j}{\bar{x}} + x_j \gamma_j \left( -\bar{x}^{-2} \frac{\partial \bar{x}}{\partial x_i} \right)$$

Evaluating $\frac{\partial \bar{x}}{\partial x_i}$:

$$\frac{\partial \bar{x}}{\partial x_i} = \frac{1}{2 \bar{x}} \cdot \frac{2 x_i}{d} = \frac{x_i}{d \bar{x}}$$

Substituting back into $\frac{\partial y_j}{\partial x_i}$:

$$\frac{\partial y_j}{\partial x_i} = \delta_{ij} \frac{\gamma_i}{\bar{x}} - \frac{x_j \gamma_j x_i}{d \bar{x}^3}$$

Multiplying by $\hat{g}_j$ and summing over all indices $j$:

$$\frac{\partial \mathcal{L}}{\partial x_i} = \hat{g}i \frac{\gamma_i}{\bar{x}} - \sum{j=1}^d \hat{g}_j \frac{x_j \gamma_j x_i}{d \bar{x}^3} = \frac{\gamma_i}{\bar{x}} \left( \hat{g}i - \frac{x_i}{d \bar{x}^2} \sum{j=1}^d \hat{g}_j \gamma_j x_j \right)$$

Vectorizing across all features yields the RMSNorm backward pass equation:

$$\nabla_x \mathcal{L} = \frac{\gamma}{\text{RMS}(x)} \odot \left( \nabla_y \mathcal{L} - \frac{x}{d \cdot \text{RMS}(x)^2} \left( (\nabla_y \mathcal{L} \odot \gamma) \cdot x \right) \right)$$

Activation Checkpointing (Gradient Checkpointing)

To compute gradients during backpropagation, intermediate activations (such as $Q, K, V$ projections, softmax score matrices, and SwiGLU inputs) must be saved in GPU memory during the forward pass. Activation checkpointing (gradient checkpointing) selectively discards intermediate activations during forward execution and recomputes them on demand during the backward pass, saving up to 70% of activation VRAM at the cost of approximately 20% extra FLOPs.

AdamW Optimizer Inner Mechanics and Decoupled Weight Decay

Standard Stochastic Gradient Descent (SGD) struggles with transformer optimization due to sparse gradients and ill-conditioned loss surfaces. Training pipelines utilize AdamW (Adam with decoupled weight decay).

Mathematical Derivation of Decoupled Weight Decay

In standard Adam with L2 regularization, weight decay penalty $\frac{\lambda'}{2} |\theta|2^2$ is added directly to loss objective $\mathcal{L}$, modifying gradient $g_t$ to $g_t' = g_t + \lambda' \theta{t-1}$.

When this combined gradient enters Adam's second moment accumulator:

$$v_t = \beta_2 v_{t-1} + (1 - \beta_2) (g_t + \lambda' \theta_{t-1})^2$$

The parameter update scales by $\frac{1}{\sqrt{\hat{v}t} + \epsilon}$. If a parameter receives small historical gradients $g_t \approx 0$, then $v_t \approx (1-\beta_2)(\lambda' \theta{t-1})^2$. Dividing the momentum update by $\sqrt{\hat{v}t}$ causes parameter magnitude $|\theta{t-1}|$ in numerator and denominator to cancel out. Consequently, parameters with tiny historical gradients experience disproportionately massive weight decay steps.

AdamW solves this by decoupling weight decay completely from gradient moment calculations.

For parameter vector $\theta$ and its computed gradient $g_t = \nabla_\theta \mathcal{L}_t$ at step $t$:

1. Compute exponentially decaying first moment vector (mean of gradients):

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

2. Compute exponentially decaying second moment vector (uncentered variance of gradients):

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

3. Correct for zero-initialization bias in early steps:

$$\hat{m}_t = \frac{m_t}{1 - \beta_1^t}, \quad \hat{v}_t = \frac{v_t}{1 - \beta_2^t}$$

4. Apply parameter update with decoupled weight decay parameter $\lambda$:

$$\theta_t = \theta_{t-1} - \eta_t \left( \frac{\hat{m}_t}{\sqrt{\hat{v}t} + \epsilon} + \lambda \theta{t-1} \right)$$

Hyperparameter defaults for LLM pre-training are typically set to: $\beta_1 = 0.9$, $\beta_2 = 0.95$, $\epsilon = 10^{-8}$, and $\lambda = 0.1$.

Optimizer Memory Accounting

The memory required to maintain model weights and optimizer states is the primary bottleneck in distributed LLM training.

For a model with $P$ parameters trained using AdamW in mixed precision:

A 70-billion parameter model requires over 1.12 TB of GPU VRAM purely for static state before accounting for activation tensors or KV caches. No single modern GPU (such as an NVIDIA H100 with 80 GB HBM3) can hold these states without distributed memory partitioning.

Gradient Accumulation and Global Batch Size

Optimal training stability requires large global batch sizes $B_{\text{global}}$ (typically between 2 million and 4 million tokens per step, e.g., 1,024 sequences of context length 4,096).

If hardware memory caps physical micro-batch size $b_{\text{micro}}$ to 2 sequences per GPU across $N_{\text{gpus}} = 128$ GPUs, the physical batch per step is $2 \times 128 = 256$ sequences.

To reach target $B_{\text{global}} = 1,024$, the runtime accumulates gradients over $A = 4$ micro-steps:

1. Zero optimizer gradients.
2. For micro-step $a = 1 \dots A$:
   • Run forward pass on micro-batch $b_{\text{micro}}$.
   • Compute scaled loss $\mathcal{L}_a = \frac{\mathcal{L}}{A}$.
   • Run backward pass and accumulate gradients into $g_{\text{acc}} \leftarrow g_{\text{acc}} + \nabla \mathcal{L}_a$.
3. Execute optimizer update step using accumulated gradients $g_{\text{acc}}$.

Gradient Norm Clipping

Spikes in training loss caused by corrupted data batches or architectural instabilities produce exploding gradient vectors. If unclipped, large gradient steps destroy learned representation spaces.

Before applying optimizer updates, global L2 gradient norm |\mathbf{g}|_2 is computed across all parameters:

$$|\mathbf{g}|2 = \sqrt{\sum{i} |g_i|_2^2}$$

If |\mathbf{g}|2 exceeds maximum threshold $M{\text{clip}}$ (typically $M_{\text{clip}} = 1.0$), gradients are scaled down in place:

$$\mathbf{g} \leftarrow \mathbf{g} \cdot \frac{M_{\text{clip}}}{\max(|\mathbf{g}|2, M{\text{clip}})}$$

import torch
 
def clip_gradients_(parameters, max_norm: float):
    parameters = [p for p in parameters if p.grad is not None]
    total_norm = torch.sqrt(sum(p.grad.detach().pow(2).sum() for p in parameters))
    clip_coef = max_norm / (total_norm + 1e-6)
    
    if clip_coef < 1.0:
        for p in parameters:
            p.grad.detach().mul_(clip_coef)
    return total_norm

Learning Rate Schedules

LLM pre-training does not use static learning rates. Training follows a two-stage schedule:

1. Linear Warmup: The learning rate ramps up linearly from $0$ to maximum rate $\eta_{\max}$ (e.g., $1.5 \times 10^{-4}$) over the first $N_{\text{warmup}}$ steps (typically initial 2,000 steps or 0.5% of total tokens). This prevents early gradient fluctuations from destabilizing initialized weights.
2. Cosine Decay: The rate decays following a cosine schedule down to minimum floor $\eta_{\min} = 0.1 \times \eta_{\max}$:

$$\eta_t = \eta_{\min} + \frac{1}{2}(\eta_{\max} - \eta_{\min})\left(1 + \cos\left(\frac{t - N_{\text{warmup}}}{T_{\text{max}} - N_{\text{warmup}}} \pi \right)\right)$$

Learning Rate (eta)
  ^
  |     /-------------\ 
eta_max|    /                 |   /                   |  /                     | /                       |/                       \--- eta_min
  +-----------------------------------> Training Steps (t)
   |<--->|                   |
   Warmup                   Cosine Decay

Distributed Cluster Parallelism (3D Parallelism and DeepSpeed/FSDP)

Scaling training to hundreds of billions of parameters requires splitting model parameters, gradients, optimizer states, and activations across clusters composed of thousands of GPUs. This orchestrates three distinct parallel dimensions: 3D Parallelism.

                  +-----------------------------------+
                  |        3D Parallel Grid           |
                  +-----------------------------------+
                  |                                   |
    +-------------+-------------+       +-------------+-------------+
    |   Pipeline Parallel (PP)  |       |    Tensor Parallel (TP)   |
    | (Layers split across nodes|       | (Matrices split in node   |
    |  via Inter-Node Network)  |       |   via high-speed NVLink)  |
    +-------------+-------------+       +-------------+-------------+
                  |                                   |
                  +-----------------+-----------------+
                                    |
                                    v
                     +------------------------------+
                     |  Data Parallel (DP / FSDP)   |
                     | (Data batches split, states  |
                     |  sharded across all GPUs)    |
                     +------------------------------+

Fully Sharded Data Parallelism (ZeRO / FSDP)

In standard Distributed Data Parallelism (DDP), every GPU retains a full copy of model parameters and optimizer states while processing a subset of batch samples. As shown in optimizer memory accounting, this crashes into VRAM limits.

DeepSpeed ZeRO (Zero Redundancy Optimizer) and PyTorch FSDP eliminate memory redundancy by sharding state components across data-parallel GPUs ($N_d$):

• ZeRO-Stage 1: Shards AdamW optimizer states ($m, v, \theta_{\text{master}}$) across $N_d$ GPUs. Memory per GPU drops from $16P$ bytes to $4P + \frac{12P}{N_d}$.
• ZeRO-Stage 2: Shards gradients alongside optimizer states. Memory per GPU drops to $2P + \frac{14P}{N_d}$.
• ZeRO-Stage 3 / FSDP: Shards model parameters, gradients, and optimizer states across all $N_d$ GPUs. Memory per GPU drops to $\frac{16P}{N_d}$.

Under FSDP / ZeRO-Stage 3:

1. Before forward execution of layer $l$, GPUs execute an All-Gather collective operation to dynamically reconstruct full parameters for layer $l$.
2. Forward computation runs for layer $l$.
3. Reconstructed parameters for layer $l$ are purged immediately from memory.
4. During backward execution, parameters for layer $l$ are gathered again via All-Gather.
5. Gradients are computed and sharded across GPUs using a Reduce-Scatter collective operation.

Forward Step (Layer l):
  1. All-Gather(Layer l parameters)  [GPU 0..N reconstruct full layer]
  2. Compute Forward Pass
  3. Free Layer l parameters         [Keep sharded fraction only]
 
Backward Step (Layer l):
  1. All-Gather(Layer l parameters)
  2. Compute Backward Gradients
  3. Reduce-Scatter(Gradients)       [Shard gradients to owner GPUs]
  4. Free Layer l parameters

Communication Overhead Derivation: DDP vs ZeRO-3

A common misconception is that sharding parameters increases network communication overhead drastically.

• Standard DDP: Transfers gradients once per backward pass using an All-Reduce operation. Total communication volume per step is $2 \cdot \frac{N_d - 1}{N_d} \cdot 2P \approx 4P$ bytes.
• ZeRO-Stage 3: Executes one All-Gather of parameters in the forward pass ($2P$ bytes), one All-Gather of parameters in the backward pass ($2P$ bytes), and one Reduce-Scatter of gradients in the backward pass ($2P$ bytes). Total communication volume is $6P$ bytes per step.

ZeRO-Stage 3 increases communication overhead by $1.5\times$ compared to standard DDP while reducing memory footprint by factor $N_d$. On modern InfiniBand networks (such as 400 Gbps NDR InfiniBand), this extra communication is overlapped with matrix multiplication compute kernels.

Tensor Parallelism (Megatron-LM TP)

When an individual layer's execution exceeds single-GPU VRAM or requires reduced latency, Tensor Parallelism splits individual weight matrices intra-node across GPUs linked by high-bandwidth NVLink interconnects (900 GB/s on H100).

Megatron-LM partitions linear transformations across Column-Parallel and Row-Parallel layers:

• Column-Parallel Layer (QKV Projections & MLP Gate/Up Projections $W_1, W_3$): Matrix $W \in \mathbb{R}^{h \times H}$ is split along columns into $W = [W^{(1)} \mid W^{(2)}]$. Input matrix $X$ is duplicated to each GPU via copy:

$$Y^{(1)} = X W^{(1)}, \quad Y^{(2)} = X W^{(2)} \implies Y = [Y^{(1)} \mid Y^{(2)}]$$

• Row-Parallel Layer (Attention Output Projection $W_O$ & MLP Down Projection $W_2$): Matrix $W \in \mathbb{R}^{H \times h}$ is split along rows into $W = \begin{bmatrix} W^{(1)} \ W^{(2)} \end{bmatrix}$. Input $X$ is split along columns $[X^{(1)} \mid X^{(2)}]$:

$$Y = X^{(1)} W^{(1)} + X^{(2)} W^{(2)}$$

Outputs from each GPU are summed across the Tensor Parallel group using an All-Reduce collective primitive.

Column-Parallel Linear (Split W by columns):
  X ---> [ GPU 0: W_1 ] ---> Y_1 --+
  X ---> [ GPU 1: W_2 ] ---> Y_2 --+--> Concatenate [Y_1 | Y_2]
 
Row-Parallel Linear (Split W by rows):
  X_1 -> [ GPU 0: W_1 ] ---> Y_1 --+
  X_2 -> [ GPU 1: W_2 ] ---> Y_2 --+--> All-Reduce (Sum) ---> Y_final

Pipeline Parallelism (Megatron PP and 1F1B Scheduling)

Pipeline Parallelism partitions the transformer network sequentially by layer groups across $P_{pp}$ distinct GPU nodes (for example, assigning layers 1-16 to Node 0, layers 17-32 to Node 1, and so on).

To prevent downstream GPUs from idling while waiting for upstream nodes (the pipeline bubble), batch sequences are divided into $m$ micro-batches.

The execution engine uses the 1F1B (One Forward, One Backward) schedule:

1. Warmup Phase: Upstream nodes execute forward passes for initial micro-batches until the pipeline fills.
2. Steady State: Each node executes exactly one forward pass on an incoming micro-batch, followed immediately by one backward pass on a completed micro-batch (1F1B).
3. Cooldown Phase: Remaining backward passes are drained.

Pipeline Bubble Overhead Derivation

Let $P_{pp}$ be the number of pipeline stages, and $m$ be the number of micro-batches per batch step.

In 1F1B execution, the total time required to execute one global batch step is proportional to $(m + P_{pp} - 1)$ time units.

The fractional pipeline bubble overhead $F_{\text{bubble}}$ is derived by dividing total idle time units by total compute units:

$$F_{\text{bubble}} = \frac{P_{pp} - 1}{m + P_{pp} - 1}$$

To maintain pipeline bubble overhead below 10% ($F_{\text{bubble}} < 0.10$), cluster configurations set micro-batch count $m \ge 8 \cdot P_{pp}$.

Numerical Stability and Mixed-Precision (FP32, FP16, BF16, and FP8)

Modern cluster hardware executes matrix multiplications using specialized hardware units (such as NVIDIA Tensor Cores). Execution speed depends directly on the numerical precision format used.

FP32:   [ S ] [ E E E E E E E E ] [ M M M M M M M M M M M M M M M M M M M M M M M ]
        1 bit      8 exponent bits                      23 mantissa bits
 
FP16:   [ S ] [ E E E E E ] [ M M M M M M M M M M ]
        1 bit   5 exp bits        10 mantissa bits
 
BF16:   [ S ] [ E E E E E E E E ] [ M M M M M M M ]
        1 bit      8 exp bits       7 mantissa bits

Numerical Formats Overview

Loss Scaling Mechanics for FP16

FP16 reserves only 5 bits for its exponent, capping maximum representable value at $65,504$ and underflow threshold at $6.10 \times 10^{-5}$. During backpropagation, transformer activation gradients frequently drop below $10^{-5}$, causing underflow to absolute zero (0.0). Conversely, unnormalized attention logits can exceed $65,504$, triggering overflow to infinity (Inf) and returning NaN loss vectors.

To train in FP16 without underflow, runtimes apply Dynamic Loss Scaling:

1. Multiply forward scalar loss by scale factor $S$ (initial value $S = 2^{16} = 65,536$):

$$\mathcal{L}_{\text{scaled}} = \mathcal{L} \cdot S$$

2. Execute backpropagation. Gradients are scaled up by $S$, shifting small values into FP16 representable range:

$$g_{\text{scaled}} = g \cdot S$$

3. Prior to optimizer update, inspect gradients for Inf or NaN values.
   • If non-finite values are detected: skip the optimizer update step, discard gradients, and halve scale factor: $S \leftarrow S / 2$.
   • If all gradients are finite: unscale gradients $g = g_{\text{scaled}} / S$ and execute FP32 master weight update. If no overflows occur for $N_{steps}$ (e.g., 2,000 steps), double scale factor: $S \leftarrow S \times 2$.

Why BF16 Replaced FP16 for Pre-Training

bfloat16 (BF16) preserves the full 8-bit exponent field of FP32 while reducing mantissa precision to 7 bits.

Because BF16 shares the exact dynamic range of FP32 ($\approx 10^{-38} \dots 10^{38}$):

• Gradients cannot underflow or overflow under normal training dynamics.
• Dynamic loss scaling logic ($S$) is entirely eliminated.
• Tensor Cores run at peak hardware FLOP throughput without numerical instability or NaN loss spikes.

BF16 has replaced FP16 as the universal standard format for pre-training models across modern GPU architectures (NVIDIA Ampere, Hopper, Blackwell, and AMD Instinct).

Post-Training Alignment (SFT and DPO)

Pre-training yields a raw base language model capable of completing matching text patterns. It does not yield an assistant that follows user instructions or aligns with safety criteria. Transforming a base model into a usable chat model requires post-training alignment.

+-------------------+
| Base Model        |  Pre-trained on trillions of raw text tokens
+-------------------+
          |
          v
+-------------------+
| SFT Stage         |  Trained on instruction-response pairs (Causal Masking)
+-------------------+
          |
          v
+-------------------+
| DPO Alignment     |  Direct Preference Optimization over (y_w, y_l) pairs
+-------------------+
          |
          v
+-------------------+
| Aligned Model     |  Final chat/assistant weights
+-------------------+

Supervised Fine-Tuning (SFT)

Supervised Fine-Tuning trains the base model on curated instruction-response pairs:

User: Configure an Nginx reverse proxy block listening on port 8080.
Assistant: server { listen 8080; location / { proxy_pass http://127.0.0.1:3000; } }

Loss Masking on Prompt Tokens

During SFT, running standard cross-entropy loss over the entire sequence causes the model to optimize next-token prediction for the user's prompt tokens. This degrades performance because prompt syntax is fixed.

SFT applies Causal Loss Masking. The loss mask $M_t$ is set to $0$ for prompt tokens $x_{1 \dots P}$ and $1$ for assistant response tokens $y_{1 \dots R}$:

$$\mathcal{L}{\text{SFT}} = -\frac{1}{\sum{t=1}^T M_t} \sum_{t=1}^T M_t \log P_\theta(x_t \mid x_{<t})$$

import torch
 
def compute_sft_loss(logits, targets, prompt_lengths):
    B, T, V = logits.shape
    loss_fn = torch.nn.CrossEntropyLoss(reduction='none')
    shift_logits = logits[:, :-1, :].contiguous().view(-1, V)
    shift_targets = targets[:, 1:].contiguous().view(-1)
    
    unmasked_loss = loss_fn(shift_logits, shift_targets).view(B, T - 1)
    mask = torch.zeros_like(unmasked_loss)
    for b in range(B):
        mask[b, prompt_lengths[b]-1:] = 1.0
        
    masked_loss = (unmasked_loss * mask).sum() / mask.sum()
    return masked_loss

Direct Preference Optimization (DPO)

Historically, preference alignment required Reinforcement Learning from Human Feedback (RLHF) using PPO (Proximal Policy Optimization). PPO requires training a separate Reward Model, maintaining four simultaneous neural networks in memory (Actor, Critic, Reference, and Reward models), and balancing unstable policy gradient updates.

Direct Preference Optimization (DPO) mathematically reparameterizes the reward function, enabling direct policy optimization over preference pairs without explicit reward modeling or RL policy loops.

Mathematical Derivation of DPO Loss and Implicit Reward

Given prompt $x$, the dataset contains a preferred response $y_w$ (winning sample) and a dispreferred response $y_l$ (losing sample).

Under Bradley-Terry preference modeling, human preference probability is expressed via a ground-truth reward function $r(x, y)$:

$$P(y_w \succ y_l \mid x) = \sigma(r(x, y_w) - r(x, y_l))$$

RLHF solves for optimal policy $\pi_\theta$ maximizing reward while penalizing KL-divergence from the frozen reference SFT policy $\pi_{\text{ref}}$:

$$\max_{\pi} \mathbb{E}{x \sim \mathcal{D}, y \sim \pi(y \mid x)} [r(x, y)] - \beta D{\text{KL}}(\pi(y \mid x) \parallel \pi_{\text{ref}}(y \mid x))$$

The exact analytical solution to this constrained optimization objective yields the implicit relationship between reward $r(x, y)$ and optimal policy $\pi_\theta$:

$$r(x, y) = \beta \log \frac{\pi_\theta(y \mid x)}{\pi_{\text{ref}}(y \mid x)} + \beta \log Z(x)$$

where $Z(x)$ is the partition function. Substituting this analytical reward identity directly back into the Bradley-Terry preference model causes the partition function $\log Z(x)$ to cancel out completely:

$$r(x, y_w) - r(x, y_l) = \beta \log \frac{\pi_\theta(y_w \mid x)}{\pi_{\text{ref}}(y_w \mid x)} - \beta \log \frac{\pi_\theta(y_l \mid x)}{\pi_{\text{ref}}(y_l \mid x)}$$

This yields the explicit DPO objective function optimized directly over parameter weights $\theta$:

$$\mathcal{L}{\text{DPO}}(\pi\theta; \pi_{\text{ref}}) = -\mathbb{E}{(x, y_w, y_l) \sim \mathcal{D}} \left[ \log \sigma \left( \beta \log \frac{\pi\theta(y_w \mid x)}{\pi_{\text{ref}}(y_w \mid x)} - \beta \log \frac{\pi_\theta(y_l \mid x)}{\pi_{\text{ref}}(y_l \mid x)} \right) \right]$$

Derivation of DPO Loss Gradient w.r.t Policy Parameters $\theta$

To understand how DPO updates policy weights during backpropagation, we derive the gradient $\nabla_\theta \mathcal{L}_{\text{DPO}}$.

Define implicit reward estimate $\hat{r}\theta(x, y) = \beta \log \frac{\pi\theta(y \mid x)}{\pi_{\text{ref}}(y \mid x)}$.

Let $u = \hat{r}\theta(x, y_w) - \hat{r}\theta(x, y_l)$.

The DPO loss for a single preference tuple $(x, y_w, y_l)$ is:

$$\mathcal{L}_{\text{DPO}}(\theta) = -\log \sigma(u)$$

Taking the derivative with respect to scalar $u$:

$$\frac{d}{du} (-\log \sigma(u)) = -\frac{\sigma'(u)}{\sigma(u)} = -\frac{\sigma(u)(1 - \sigma(u))}{\sigma(u)} = -(1 - \sigma(u)) = -\sigma(-u)$$

Now take the gradient of $u$ with respect to policy parameters $\theta$:

$$\nabla_\theta u = \beta \nabla_\theta \log \pi_\theta(y_w \mid x) - \beta \nabla_\theta \log \pi_\theta(y_l \mid x)$$

Applying the chain rule:

$$\nabla_\theta \mathcal{L}{\text{DPO}}(\theta) = -\beta \cdot \sigma\left(\hat{r}\theta(x, y_l) - \hat{r}\theta(x, y_w)\right) \left[ \nabla\theta \log \pi_\theta(y_w \mid x) - \nabla_\theta \log \pi_\theta(y_l \mid x) \right]$$

This gradient expression reveals an essential adaptive weighting mechanism:

1. Adaptive Weighting Factor: The gradient is scaled by $\sigma\left(\hat{r}\theta(x, y_l) - \hat{r}\theta(x, y_w)\right)$, which evaluates how incorrectly the policy model currently rates the pair. If the policy model already assigns much higher reward to $y_w$ than $y_l$ (i.e., $\hat{r}\theta(x, y_w) \gg \hat{r}\theta(x, y_l)$), this weighting factor approaches zero, preventing over-fitting on already aligned pairs. Conversely, if the policy model incorrectly rates $y_l$ higher than $y_w$, the factor approaches one, applying a strong gradient push.
2. Dual Likelihood Push/Pull: The update pushes $\log \pi_\theta(y_w \mid x)$ upward (increasing likelihood of preferred completion) while pulling $\log \pi_\theta(y_l \mid x)$ downward (suppressing likelihood of dispreferred completion).

import torch
import torch.nn.functional as F
 
def compute_dpo_loss(
    policy_chosen_logps: torch.FloatTensor,
    policy_rejected_logps: torch.FloatTensor,
    reference_chosen_logps: torch.FloatTensor,
    reference_rejected_logps: torch.FloatTensor,
    beta: float = 0.1
) -> torch.FloatTensor:
    pi_logratios = policy_chosen_logps - policy_rejected_logps
    ref_logratios = reference_chosen_logps - reference_rejected_logps
    logits = pi_logratios - ref_logratios
    losses = -F.logsigmoid(beta * logits)
    return losses.mean()

The hyperparameter $\beta$ (typically between $0.01$ and $0.1$) controls strength of the KL penalty against the reference model.

DPO increases the probability of preferred completion $y_w$ while suppressing dispreferred completion $y_l$. The implicit reward formulation ensures that if the policy model increases response probability for $y_w$ beyond reference policy baseline $\pi_{\text{ref}}$, the sample contributes positively to loss convergence without requiring complex actor-critic reward modeling infrastructure.