AI/ML Alignment Deep Learning PyTorch

LLM Post-Training: A Practitioner's First-Principles Guide

October 26, 2026 • ☕ 8 min read

Pre-training teaches language models grammar, vocabulary, and world facts via next-token prediction over trillions of tokens. But base models are fundamentally unsteerable: ask them a question, and they might generate an alternate question, ramble indefinitely, or parrot internet tropes.

Post-training is the staged discipline that turns an unaligned statistical next-token predictor into a steerable, reliable, and safe assistant. Here is an end-to-end breakdown of the modern alignment stack: SFT and LoRA, preference optimization (DPO, PPO), rule-based exploration (GRPO), and edge distillation, complete with exact loss formulations and minimal PyTorch implementations.

0. The Post-Training Pipeline

Modern post-training decomposes alignment into distinct, sequential stages where each phase tackles a specific failure mode:

Figure 1: Canonical post-training workflow from foundational weights to deployment.

1. Supervised Fine-Tuning (SFT)

SFT teaches the model conversation mechanics, turn-taking, and response formatting using high-quality demonstration pairs $(x, y)$.

Masked Cross-Entropy Objective

Crucially, we do not compute loss on the user prompt $x$. We compute negative log-likelihood solely over the assistant's response tokens $y$:

Masked Cross-Entropy Loss $$\mathcal{L}_{\text{SFT}}(\theta) = -\frac{1}{N_{\text{valid}}} \sum_{t=1}^{L} m_t \log \pi_\theta(y_t \mid x, y_{\lt t})$$ where $m_t = \mathbb{I}(\text{labels}[t] \neq -100)$ masks prompt and padding tokens so gradients only supervise target responses.

LoRA: Low-Rank Adaptation

Full parameter fine-tuning of 70B models requires massive optimizer states (16 bytes/param in AdamW). LoRA freezes the pre-trained weights $W_0 \in \mathbb{R}^{d \times k}$ and factorizes updates into rank $r \ll \min(d, k)$ adapters:

LoRA Weight Formulation $$W = W_0 + \Delta W = W_0 + \frac{\alpha}{r}\, B \cdot A$$ $A \sim \mathcal{N}(0, 1/r)$ and $B = 0$. At step zero, $\Delta W = 0$ exactly preserving base model weights while permitting non-zero initial gradients $\frac{\partial\mathcal{L}}{\partial A} \neq 0$.

QLoRA Note: In 4-bit QLoRA (NF4 quantiles), never merge LoRA adapters into quantized weights directly because discretization error degrades accuracy. Always dequantize or reload the 16-bit base before merging weights.

VRAM Footprint: 1 model instance. ~4.1 GB for a 7B model in NF4 with trainable adapters.
Key Hyperparameters: Learning rate 2e-4, rank $r=16$, scale $\alpha=32$, cosine decay with 5% warmup, 2–3 epochs.
Primary Failure Mode: Exposure Bias: model is conditioned only on ground-truth prefixes at train time, accumulating compounding errors at autoregressive inference time.

💻 PyTorch Implementation: SFT Loss Masking & LoRA LayerPure PyTorch

import torch
import torch.nn as nn
import torch.nn.functional as F

def sft_masked_loss(logits: torch.Tensor, labels: torch.Tensor, ignore_index: int = -100) -> torch.Tensor:
    """
    logits: [Batch, SeqLen, Vocab]  (raw model token predictions)
    labels: [Batch, SeqLen]         (target token IDs, user prompt masked to -100)
    """
    # Shift so that token at t-1 predicts token at t
    shift_logits = logits[..., :-1, :].contiguous()
    shift_labels = labels[..., 1:].contiguous()
    
    # Cross-entropy automatically normalizes by non-ignored tokens
    return F.cross_entropy(
        shift_logits.view(-1, shift_logits.size(-1)),
        shift_labels.view(-1),
        ignore_index=ignore_index,
        reduction="mean"
    )

class LoRALinear(nn.Module):
    def __init__(self, in_features: int, out_features: int, r: int = 16, alpha: float = 32.0):
        super().__init__()
        self.scaling = alpha / r
        self.weight = nn.Parameter(torch.randn(out_features, in_features), requires_grad=False)
        # B=0 guarantees Delta_W=0 at step 0
        self.lora_A = nn.Parameter(torch.randn(r, in_features) / (r ** 0.5))
        self.lora_B = nn.Parameter(torch.zeros(out_features, r))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        base_out = F.linear(x, self.weight)
        lora_out = F.linear(F.linear(x, self.lora_A), self.lora_B) * self.scaling
        return base_out + lora_out

2. Direct Preference Optimization (DPO)

SFT models frequently generate verbose, overly sycophantic, or inaccurate answers. Preference alignment tunes the model using paired pairwise rankings $(x, y_w, y_l)$, consisting of a chosen ($y_w$) and rejected ($y_l$) response.

Under the Bradley-Terry preference model with a KL penalty against reference policy $\pi_{\text{ref}}$, the analytically optimal policy satisfies $\pi^*(y \mid x) = \frac{1}{Z(x)} \pi_{\text{ref}}(y \mid x) \exp\!\bigl(\frac{1}{\beta} r(x,y)\bigr)$. Rearranging yields the implicit reward $r(x, y) = \beta \log \frac{\pi_\theta(y \mid x)}{\pi_{\text{ref}}(y \mid x)} + \beta \log Z(x)$. When computing the pairwise difference $r(x, y_w) - r(x, y_l)$, the partition function $\beta \log Z(x)$ cancels identically, enabling direct optimization without an auxiliary reward model.

Figure 2: Direct Preference Optimization (DPO) architecture with pairwise implicit reward margin and policy feedback loop.

Mathematical Formulation

DPO Objective (One-Line Formulation) $$\mathcal{L}_{\text{DPO}}(\theta) = -\mathbb{E}_{(x, y_w, y_l)}\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]$$

RPO Objective (Likelihood Displacement Fix) $$\mathcal{L}_{\text{RPO}}(\theta) = \mathcal{L}_{\text{DPO}}(\theta) - \alpha_{\text{nll}} \log \pi_\theta(y_w \mid x)$$ The NLL anchor prevents winning response probability $\pi(y_w)$ from collapsing when optimizing pairwise margins.

VRAM Footprint: 2 models ($\pi_\theta$ + $\pi_{\text{ref}}$). When using LoRA, disabling adapters yields $\pi_{\text{ref}}$ on the same weights, requiring zero extra base model memory.
Key Hyperparameters: Learning rate 5e-6 (40× lower than SFT), $\beta \in [0.05, 0.15]$, RPO $\alpha_{\text{nll}} \in [0.1, 0.2]$, 1 epoch.
Primary Failure Mode: Length Bias: the policy exploits unnormalized token-sum margins by generating arbitrarily verbose outputs.

💻 PyTorch Implementation: DPO Loss with RPO AnchorPure PyTorch

import torch
import torch.nn.functional as F

def dpo_rpo_loss(
    policy_chosen_logps: torch.Tensor,    # [Batch] \sum log pi_theta(y_w | x)
    policy_rejected_logps: torch.Tensor,  # [Batch] \sum log pi_theta(y_l | x)
    ref_chosen_logps: torch.Tensor,       # [Batch] \sum log pi_ref(y_w | x)
    ref_rejected_logps: torch.Tensor,     # [Batch] \sum log pi_ref(y_l | x)
    beta: float = 0.1,
    alpha_nll: float = 0.0                # Set > 0 for RPO anchor
) -> torch.Tensor:
    # Compute margin between chosen and rejected implicit rewards
    pi_logratios = policy_chosen_logps - policy_rejected_logps
    ref_logratios = ref_chosen_logps - ref_rejected_logps
    logits = beta * (pi_logratios - ref_logratios)
    
    # Binary cross entropy on Bradley-Terry formulation
    dpo_loss = -F.logsigmoid(logits).mean()
    
    # Optional RPO anchor to suppress likelihood displacement
    if alpha_nll > 0.0:
        nll_loss = -policy_chosen_logps.mean()
        return dpo_loss + alpha_nll * nll_loss
        
    return dpo_loss

3. Proximal Policy Optimization (PPO)

PPO represents the classical Reinforcement Learning from Human Feedback (RLHF) paradigm introduced in InstructGPT. Unlike DPO's offline static dataset, PPO generates online rollouts and updates policy weights against a learned neural reward model $R_\psi$.

Figure 3: Proximal Policy Optimization (PPO) complete architecture with 4 models in memory, Generalized Advantage Estimation, and policy feedback loop.

Mathematical Formulation

PPO Clipped Objective (One-Line Formulation) $$\mathcal{J}_{\text{PPO}}(\theta) = \mathbb{E}_{(x,y) \sim \mathcal{D},\, t} \left[ \min\!\left( \rho_t(\theta)\hat{A}_t,\; \text{clip}(\rho_t(\theta),\, 1-\varepsilon,\, 1+\varepsilon)\hat{A}_t \right) \right]$$ where $\rho_t(\theta) = \frac{\pi_\theta(y_t \mid x, y_{\lt t})}{\pi_{\theta_{\text{old}}}(y_t \mid x, y_{\lt t})}$ and advantage $\hat{A}_t = r_t - V_\phi(s_t)$ with composite reward $r_t = r_\varphi(x, y) - \beta D_{\text{KL},t}$.

Why PPO is Fading for Reasoning: PPO suffers from three severe weaknesses: (1) enormous multi-model VRAM overhead, (2) critic training instability where value estimates fail to converge across long chains of thought, and (3) reward hacking, where policies exploit flaws in the neural reward model rather than truly solving tasks.

💻 PyTorch Implementation: PPO Clipped Surrogate LossPure PyTorch

import torch

def ppo_clipped_loss(
    policy_logps: torch.Tensor,  # [Batch, Len] log pi_theta(a_t | s_t)
    old_logps: torch.Tensor,     # [Batch, Len] log pi_old(a_t | s_t)
    advantages: torch.Tensor,    # [Batch, Len] GAE normalized advantage A_t
    mask: torch.Tensor,          # [Batch, Len] 1.0 for valid generated tokens
    clip_eps: float = 0.2
) -> torch.Tensor:
    # Probability ratio: rho_t = exp(log pi_theta - log pi_old)
    ratios = torch.exp(policy_logps - old_logps)
    
    # Clipped surrogate objective: min(rho * A, clip(rho, 1-eps, 1+eps) * A)
    surr1 = ratios * advantages
    surr2 = torch.clamp(ratios, 1.0 - clip_eps, 1.0 + clip_eps) * advantages
    policy_loss = -torch.min(surr1, surr2)
    
    # Normalize exclusively over active response tokens
    return (policy_loss * mask).sum() / mask.sum().clamp(min=1.0)

4. Group Relative Policy Optimization (GRPO)

Introduced in DeepSeekMath (Shao et al., 2024) and popularized by DeepSeek-R1, GRPO eliminates the neural Value Critic ($V_\phi$) completely. Instead of learning a value function, GRPO samples $G$ responses for each prompt $q$, scores them with deterministic rubrics (e.g. unit test passes, compiler outputs, or exact math solutions), and normalizes advantages against group statistics.

Figure 4: Group Relative Policy Optimization (GRPO) complete architecture with group sampling, rule evaluation, and critic-free advantage normalization.

Mathematical Formulation

GRPO Clipped Objective (One-Line Formulation) $$\mathcal{J}_{\text{GRPO}}(\theta) = \mathbb{E}_{q \sim P(Q),\, \{o_i\}_{i=1}^G \sim \pi_{\theta_{\text{old}}}} \left[ \frac{1}{G}\sum_{i=1}^{G} \sum_{t=1}^{|o_i|} \min\!\left( \rho_{i,t}(\theta)\hat{A}_i,\; \text{clip}(\rho_{i,t}(\theta),\, 1-\varepsilon,\, 1+\varepsilon)\hat{A}_i \right) - \beta D_{\text{KL}}^{k_3}(\pi_\theta \| \pi_{\text{ref}}) \right]$$ where group advantage $\hat{A}_i = \frac{r_i - \mu(R)}{\sigma(R) + \epsilon}$ with group mean $\mu(R) = \frac{1}{G}\sum_{j=1}^G r_j$ and standard deviation $\sigma(R) = \sqrt{\frac{1}{G}\sum_{j=1}^G (r_j - \mu(R))^2}$.

Strictly Non-Negative k3 KL Estimator (One-Line Formulation) $$D_{\text{KL}}^{k_3}(\pi_\theta \| \pi_{\text{ref}}) = \exp(u) - u - 1 \ge 0, \qquad u = \log \frac{\pi_{\text{ref}}}{\pi_\theta}$$ The $k_3$ estimator is mathematically bounded below by zero, preventing the policy from receiving artificial positive rewards when diverging from reference weights.

The Dr. GRPO / DAPO Fix: Standard GRPO normalizes by actual response length. This accidentally assigns larger gradient updates to short, clipped responses and penalizes long, thorough reasoning. Modern implementations replace length normalization with a constant denominator $\frac{1}{L_{\max}}$ (e.g., 256 or 512), eliminating length gaming entirely.

VRAM Footprint: 2 models ($\pi_\theta$ + $\pi_{\text{ref}}$). ~50% VRAM reduction relative to PPO by cutting out the Critic model.
Key Hyperparameters: Learning rate 1e-6 to 3e-6, group size $G=8$, rollout temperature $T=0.8$, clipping $\varepsilon=0.2$, KL weight $\beta=0.03$.
Primary Failure Mode: Zero-Variance Stagnation: if all $G$ completions fail on a hard problem ($r_i=0$ for all $i$), $\sigma=0$ and gradients vanish. Solved via dense step-level shaping or hybrid curriculum prompting.

💻 PyTorch Implementation: GRPO Group Advantage & DAPO k3 LossPure PyTorch

import torch

def compute_group_advantages(rewards: torch.Tensor, group_size: int = 8, eps: float = 1e-4) -> torch.Tensor:
    """
    rewards: [Batch * G] scalar rewards evaluated from deterministic rubrics
    returns: [Batch * G] normalized group advantages without a value model
    """
    rewards_group = rewards.view(-1, group_size)                     # [Batch, G]
    mean = rewards_group.mean(dim=-1, keepdim=True)                  # [Batch, 1]
    std = rewards_group.std(dim=-1, keepdim=True, unbiased=False)    # [Batch, 1]
    adv = (rewards_group - mean) / (std + eps)
    return adv.view(-1)                                              # [Batch * G]

def grpo_loss(
    policy_logps: torch.Tensor,  # [Batch * G, Len]
    old_logps: torch.Tensor,     # [Batch * G, Len]
    ref_logps: torch.Tensor,     # [Batch * G, Len]
    advantages: torch.Tensor,    # [Batch * G]
    mask: torch.Tensor,          # [Batch * G, Len]
    clip_eps: float = 0.2,
    beta_kl: float = 0.04,
    l_max: int = 256             # DAPO constant denominator kills length gaming
) -> torch.Tensor:
    # 1. Clipped Policy Gradient objective normalized by constant L_max
    ratios = torch.exp(policy_logps - old_logps)
    adv = advantages.unsqueeze(-1)
    surr = torch.min(ratios * adv, torch.clamp(ratios, 1.0 - clip_eps, 1.0 + clip_eps) * adv)
    pg_loss = -(surr * mask).sum(dim=-1) / l_max
    
    # 2. Strictly Non-Negative k3 KL Divergence: exp(u) - u - 1 >= 0
    u = ref_logps - policy_logps
    kl_k3 = torch.exp(u) - u - 1.0
    kl_penalty = (kl_k3 * mask).sum(dim=-1) / l_max
    
    return (pg_loss + beta_kl * kl_penalty).mean()

5. Knowledge Distillation

Deploying 70B reasoning models in low-latency production environments is economically prohibitive. Distillation transfers reasoning capabilities from a large teacher $\pi_T$ to a small edge student $\pi_S$ (e.g. 1.5B–3B).

Hinton Distillation Objective

KD Loss with Temperature Scaling $$\mathcal{L}_{\text{KD}} = (1 - \alpha)\,\mathcal{L}_{\text{CE}}(y, \sigma(z_S)) + \alpha\, T^2 \cdot D_{\text{KL}}\!\left(\sigma\!\left(\frac{z_T}{T}\right) \;\Big\|\; \sigma\!\left(\frac{z_S}{T}\right)\right)$$ The $T^2$ multiplier offsets the $1/T^2$ shrinkage of softmax gradients at higher temperatures, ensuring soft probability supervision remains balanced with hard ground-truth labels.

Forward KL vs. Reverse KL

The direction of the KL divergence dictates how the student approximates the teacher's probability distribution:

Forward KL: Mode-Covering $$D_{\text{KL}}(P_T \| P_S) = \sum_{v \in V} P_T(v) \log \frac{P_T(v)}{P_S(v)}$$ Preserves reasoning paths: The student is heavily penalized if $P_S(v) \to 0$ when $P_T(v) > 0$. It forces the student to cover all exploratory tokens, including backtrack triggers (e.g. "Wait, let me recalculate...").

Reverse KL: Mode-Seeking $$D_{\text{KL}}(P_S \| P_T) = \sum_{v \in V} P_S(v) \log \frac{P_S(v)}{P_T(v)}$$ Mode-dropping risk: The student minimizes loss by collapsing onto a single sharp peak and zeroing out valid alternate branches, stripping away multi-turn recovery abilities.

💻 PyTorch Implementation: Forward vs. Reverse KL DistillationPure PyTorch

import torch
import torch.nn.functional as F

def knowledge_distillation_loss(
    student_logits: torch.Tensor,  # [Batch, Len, Vocab]
    teacher_logits: torch.Tensor,  # [Batch, Len, Vocab]
    labels: torch.Tensor,          # [Batch, Len]
    alpha: float = 0.5,            # Balance between hard CE and soft KD
    temp: float = 2.5,             # Temperature scaling factor T
    mode: str = "forward"          # 'forward' (mode-covering) or 'reverse' (mode-seeking)
) -> torch.Tensor:
    # 1. Hard cross-entropy loss on ground-truth target tokens
    shift_s = student_logits[..., :-1, :].contiguous()
    shift_y = labels[..., 1:].contiguous()
    ce_loss = F.cross_entropy(shift_s.view(-1, shift_s.size(-1)), shift_y.view(-1), ignore_index=-100)
    
    # 2. Temperature scaling
    s_log_p = F.log_softmax(student_logits / temp, dim=-1)
    t_log_p = F.log_softmax(teacher_logits / temp, dim=-1)
    s_p = F.softmax(student_logits / temp, dim=-1)
    t_p = F.softmax(teacher_logits / temp, dim=-1)
    
    # 3. Directional KL Divergence
    if mode == "forward":
        # Forward KL: KL(Teacher || Student)
        kl_loss = (t_p * (t_log_p - s_log_p)).sum(dim=-1).mean()
    else:
        # Reverse KL: KL(Student || Teacher)
        kl_loss = (s_p * (s_log_p - t_log_p)).sum(dim=-1).mean()
        
    # Scale soft loss by T^2 to match gradient magnitudes with hard CE
    return (1.0 - alpha) * ce_loss + alpha * (temp ** 2) * kl_loss

6. Head-to-Head Comparison

Dimension	PPO	DPO	GRPO
Paradigm	Online RL with rollouts	Offline Supervised log-ratio	Online RL with group sampling
Baseline Model	Neural Critic ($V_\phi$)	Implicit Reference ($\pi_{\text{ref}}$)	Statistical Group Mean $\mu(R)$
Reward Source	Neural Reward Model ($R_\psi$)	Implicit Policy Log-Ratio	Deterministic Tests / Compiler / Math
VRAM Footprint	Extreme (4 models in memory)	Low (1 base + LoRA adapter)	Moderate (Policy + Reference only)
Best Used For	Subjective conversational style	Rapid stylistic alignment & tone	Math, coding & verifiable reasoning

7. The Modern Alignment Recipe

The transition from classical RLHF to modern reasoning models (like DeepSeek-R1 and Qwen-2.5-Math) follows a clear, battle-tested recipe:

SFT Cold-Start: Fine-tune base models on a small set of clean, structured reasoning traces and multi-turn conversation formatting (using masked labels and LoRA).
GRPO for System 2 Reasoning: Apply GRPO with deterministic, verifiable outcome rewards (code test cases, math verification). Let the model develop self-correction and reflection organically.
DPO for Conversational Polish: Run a rapid DPO pass using chosen/rejected pairs to align helpfulness, eliminate sycophancy, and smooth tone.
Forward-KL Distillation: Distill the 70B teacher's soft probability distributions into 1.5B–3B student checkpoints for low-cost, real-time edge deployment.