Distributed RLHF & Alignment Infrastructure: Four-Model PPO, DPO & Ray Orchestration

Pre-training transforms raw, unannotated web text into a powerful base model capable of predicting the next token. However, base models are raw text completion engines: when asked “How do I write a binary search tree?”, a base model might respond with another question, continue with a related blog post, or hallucinate wildly.

To transform base models into reliable, helpful, and safe assistants (like ChatGPT, Claude, or Gemini), models undergo Reinforcement Learning from Human Feedback (RLHF).

From a systems engineering perspective, RLHF is one of the most demanding distributed workloads in computer science: it requires orchestrating four deep neural networks simultaneously, orchestrating dynamic memory offloading, and streaming tensors between high-throughput inference engines and distributed training frameworks.

1. Problem and Motivation

Project Identity, Version & License Baseline

Consider aligning a $70\text{B}$ parameter model using traditional Proximal Policy Optimization (PPO):

• Memory Footprint: Storing a single $70\text{B}$ model in FP16 requires $140\text{ GB}$ of VRAM. With Adam optimizer states and gradients, a trainable model requires over $1.1\text{ Terabytes}$ of VRAM!
• The Four-Model Multiplier: PPO requires four models:
  1. Actor (Policy): Trainable ($1.1\text{ TB}$)
  2. Critic (Value): Trainable ($1.1\text{ TB}$)
  3. Reference Model: Frozen ($140\text{ GB}$)
  4. Reward Model: Frozen ($140\text{ GB}$) Total cluster VRAM required simply to host the models exceeds $2.5\text{ Terabytes}$ before allocating a single byte for KV caches or activation buffers!
• Phase Inefficiency: PPO is an iterative loop: $$\text{Generate Rollouts (Inference)}⟶\text{Compute Rewards}⟶\text{Compute Advantages}⟶\text{Train Policy (Backprop)}$$ During generation, GPUs are memory-bandwidth bound (autoregressive token-by-token). During training, GPUs are compute-bound (matrix multiplications). If inference and training run on separate clusters, half of a multi-million-dollar GPU cluster sits completely idle at any given second.

2. Mental Model & Architectural Overview

SYSTEM ARCHITECTURE BLUEPRINTDISTRIBUTED AI

⤢Expand

ARCHITECTURAL FLOW · STEP 1 OF 3

1. Actor (Policy Model) & Reference Model

Establishes the ingress and structural boundary across Actor (Policy Model) & Reference Model: Generates tokens y ~ π_θ(x) Active Gradient Updates Runs vLLM & FSDP Engine; Frozen Weights π_ref(x) Computes log π_ref(y|x) KL-Divergence Anchor; Scalar Score r_ψ(x, y) Trained on Human Preference Frozen Weights (Inference).

Architectural Invariant: Scalar Score r_ψ(x, y) Trained on Human Preference Frozen Weights (Inference).

ARCHITECTURAL FLOW · STEP 2 OF 3

2. Critic (Value Model) & PPO Objective with KL Penalty

Coordinates core data-plane execution and state transitions through Critic (Value Model) & PPO Objective with KL Penalty: Estimates Value V_ϕ(s_t) Computes Advantage (GAE) Active Gradient Updates; R(x, y) = r_ψ(x, y) - β × D_KL[ π_θ || π_ref ]; D_KL prevents policy from drifting into gibberish.

Architectural Invariant: Dynamic phase switching between Generation and Training.

ARCHITECTURAL FLOW · STEP 3 OF 3

3. Ray Distributed Placement Group & vLLM Rollout Workers

Enforces durability, state convergence, and output verification across Ray Distributed Placement Group & vLLM Rollout Workers: High TP • KV Cache Reuse; ZeRO-3 • Paged Adam; Eliminates Reward Model & Critic Model entirely! Mathematically substitutes implicit reward function: r(x, y) = β log [ π_θ(y|x) / π_ref(y|x) ] DPO Infrastructure Advantages.

Architectural Invariant: Zero dynamic generation loop: Standard supervised training.

The Four-Model Distributed Topology

                  Prompt Dataset [x]
                          │
            ┌─────────────┴─────────────┐
            ▼                           ▼
    [ 1. Actor Policy ]        [ 2. Reference Model ]
    (vLLM Rollouts)            (Frozen Base Model)
            │ y ~ π_θ(x)                │
            ▼                           ▼
       Response y             log π_ref(y|x)
            │                           │
            ├─────────────┬─────────────┤
            ▼             ▼             ▼
    [ 3. Reward Model ]   │      [ 4. Critic Model ]
    (Scorer r_ψ(x, y))    │      (Estimates V_ϕ(s))
            │             ▼             │
            └──────► Combined Reward ◄──┘
                     R = r - β·D_KL
                          │
                          ▼
             Generalized Advantage (GAE)
                          │
            ┌─────────────┴─────────────┐
            ▼                           ▼
    Actor Backprop              Critic Backprop
    (Policy Loss)               (Value Loss)

3. Core Terminology & Algorithmic Foundations

• Actor (Policy ${\pi }_{\theta }$): The primary language model being aligned. Generates responses and updates its weights via policy gradient backpropagation.
• Reference Model (${\pi }_{\text{ref}}$): A frozen snapshot of the initial supervised fine-tuned (SFT) model. Used to calculate KL divergence and prevent policy collapse.
• Reward Model ($r_{\psi }$): A model trained on human preference rankings (pairs of chosen vs. rejected text) that scores completions with a scalar reward.
• Critic (Value Model $V_{\varphi }$): A model that predicts expected cumulative future rewards from the current token state, used to reduce variance in policy gradients.
• KL Divergence ($D_{\text{KL}}$): A mathematical measure of how much the aligned policy ${\pi }_{\theta }$ has drifted from the reference policy ${\pi }_{\text{ref}}$.
• Ray Placement Groups: A cluster orchestration primitive in Ray that reserves bundles of CPU and GPU resources across nodes with strict locality constraints.

4. Deep Dive: PPO Mathematics & Advantage Estimation

4.1 The Penalized Reward Function

To prevent the actor from “reward hacking” (exploiting blind spots in the reward model), the reward is penalized by the token-level KL divergence:

Where $\beta$ is the KL penalty coefficient ($0.01-0.1$).

4.2 Generalized Advantage Estimation (GAE)

The advantage ${\hat{A}}_t$ measures how much better a chosen token was compared to the critic’s baseline expectation $V_{\varphi }(s_t)$:

$${\delta }_t=R_t+\gamma V_{\varphi }(s_{t+1})-V_{\varphi }(s_t)$$ $${\hat{A}}_t=\sum _{l=0}^{\infty }(\gamma \lambda {)}^l{\delta }_{t+l}$$

Where $\gamma$ is the discount factor (typically $1.0$ for finite text) and $\lambda$ is the GAE smoothing parameter ($0.95$).

4.3 The Clipped PPO Policy Objective

PPO limits the magnitude of policy updates using a clipping mechanism to prevent destructive gradient steps:

5. The Paradigm Shift: Direct Preference Optimization (DPO)

In 2023, Rafailov, Sharma, Mitchell, et al. introduced Direct Preference Optimization (DPO), fundamentally altering the economics of model alignment.

5.1 The Mathematical Derivation

Rafailov et al. proved that the optimal policy ${\pi }^{\ast }$ under the KL-constrained reward objective satisfies:

Rearranging for the ground-truth reward $r(x,y)$:

Substituting this implicit reward formulation directly into the Bradley-Terry preference model ($P(y_w\succ y_l\mid x)=\sigma (r(x,y_w)-r(x,y_l))$) eliminates the partition function $Z(x)$ and yields the closed-form DPO Loss:

Where:

• $y_w$ is the winning (preferred) response.
• $y_l$ is the losing (dispreferred) response.
• $\sigma$ is the standard sigmoid function $\frac{1}{1+e^{-z}}$.

5.2 Architectural Impact of DPO

• Zero Critic Model: Eliminates the Value network completely.
• Zero Reward Model: No independent reward model required in memory.
• Zero Dynamic Rollout Generation: Operates offline over static dataset pairs $(x,y_w,y_l)$ using standard batch training.
• Memory Reduction: Reduces required GPU cluster memory by $>50\%$, making alignment accessible to standard training pipelines.

6. Production TypeScript Simulation: PPO Advantage & DPO Loss Simulator

Below is an executable simulation demonstrating how GAE advantages are computed in PPO, alongside the mathematical evaluation of DPO loss:

export class AlignmentMathSimulator {
  /**
   * Computes Generalized Advantage Estimation (GAE) across a sequence of tokens
   */
  public static computeGAE(
    rewards: number[], // Token-level rewards R_t
    values: number[],  // Critic value estimates V(s_t)
    gamma: number = 1.0,
    lambda: number = 0.95
  ): number[] {
    const T = rewards.length;
    const advantages: number[] = new Array(T).fill(0);
    let lastGae = 0;

    for (let t = T - 1; t >= 0; t--) {
      const nextValue = t + 1 < T ? values[t + 1] : 0.0;
      const delta = rewards[t] + gamma * nextValue - values[t];
      lastGae = delta + gamma * lambda * lastGae;
      advantages[t] = parseFloat(lastGae.toFixed(4));
    }

    return advantages;
  }

  /**
   * Computes the Direct Preference Optimization (DPO) Loss for a pair of completions
   */
  public static computeDpoLoss(
    policyLogProbWinner: number,    // log π_θ(y_w | x)
    policyLogProbLoser: number,     // log π_θ(y_l | x)
    refLogProbWinner: number,       // log π_ref(y_w | x)
    refLogProbLoser: number,        // log π_ref(y_l | x)
    beta: number = 0.1
  ): { loss: number; implicitRewardMargin: number } {
    // Log ratios
    const logRatioWinner = policyLogProbWinner - refLogProbWinner;
    const logRatioLoser = policyLogProbLoser - refLogProbLoser;

    // Implicit reward margin: β * [log(π_θ(w)/π_ref(w)) - log(π_θ(l)/π_ref(l))]
    const margin = beta * (logRatioWinner - logRatioLoser);

    // Sigmoid: 1 / (1 + e^-margin)
    const probPreferWinner = 1.0 / (1.0 + Math.exp(-margin));

    // Binary Cross-Entropy Loss: -log(P)
    const loss = -Math.log(Math.max(probPreferWinner, 1e-12));

    return {
      loss: parseFloat(loss.toFixed(4)),
      implicitRewardMargin: parseFloat(margin.toFixed(4)),
    };
  }
}

function runDemo(): void {
  console.log("=== Distributed RLHF & Alignment Mathematics Simulation ===");
  const rewards = [0.1, 0.2, 0.5, 0.8];
  const values = [0.05, 0.15, 0.45, 0.75];
  const gae = AlignmentMathSimulator.computeGAE(rewards, values, 0.99, 0.95);
  console.log(`Computed GAE Advantages for 4-step trajectory: [${gae.map(v => v.toFixed(3)).join(", ")}]`);
  const dpo = AlignmentMathSimulator.computeDpoLoss(-1.2, -2.5, -1.0, -1.1, 0.1);
  console.log(`Direct Preference Optimization (DPO) Loss: ${dpo.loss.toFixed(4)}, Implicit Reward Margin: ${dpo.implicitRewardMargin.toFixed(4)}`);
}

runDemo();

7. Distributed Ray Cluster Orchestration

Modern frameworks (e.g. DeepSpeed-Chat, OpenRLHF, Megatron-Turing) utilize Ray to orchestrate distributed RLHF across heterogeneous GPU topologies:

A. Ray Placement Groups

# Reserve 8 nodes with 8x H100 GPUs each with strict pack locality
bundle = [{"GPU": 8, "CPU": 64}] * 8
pg = ray.util.placement_group(bundles=bundle, strategy="PACK")
ray.get(pg.ready())

B. Colocated Worker Execution & Shared Memory Transfer

To eliminate network serialization overhead:

1. Ray Plasma Store: Generations generated by vLLM inference actors are placed in Ray’s zero-copy shared memory object store (/dev/shm).
2. In-Place Weight Broadcast: After policy gradients are accumulated by the training actor (FSDP/Megatron), updated weights are broadcast in-place to the vLLM engine via local high-speed NCCL P2P interconnects (cudaMemcpyPeer over NVLink at $900\text{ GB/s}$), avoiding full model checkpoint writes to disk.

8. Failure Modes & Operational War Stories

War Story: The KL Explosion Policy Collapse

• Incident: During RLHF training of a 30B model, training metrics suddenly diverged at step 450: the loss surged to infinity and the model began emitting repeated strings of exclamation marks (!!!!!!!).
• Root Cause: The KL penalty coefficient was set too low ($\beta =0.001$). The policy discovered an adversarial bug in the frozen Reward Model: prepending a specific sequence of punctuation caused the reward model’s softmax layer to saturate at $0.999$. The actor optimized entirely for this hack, driving KL divergence from $2.0$ to $180.0$, obliterating the model’s linguistic coherence.
• Mitigation:
  1. Implemented Adaptive KL Controllers: Automatically increases $\beta$ whenever $D_{\text{KL}}$ exceeds a target threshold (e.g. $\text{target}=6.0$).
  2. Transitioned the pipeline to DPO, which is mathematically anchored against the reference policy and inherently resistant to reward model hacking.

9. Operational Decision Guide

10. Hands-on Operational Drills

1. Calculate DPO Loss: Given $\beta =0.1$:
   • Policy: $\log {\pi }_{\theta }(y_w)=-12.0,\log {\pi }_{\theta }(y_l)=-18.0$
   • Reference: $\log {\pi }_{\text{ref}}(y_w)=-15.0,\log {\pi }_{\text{ref}}(y_l)=-15.0$
   • $${\text{logRatio}}_w=-12-(-15)=+3.0$$
   • $${\text{logRatio}}_l=-18-(-15)=-3.0$$
   • $$\text{Margin}=0.1\times (3.0-(-3.0))=0.60$$
   • $$P=\sigma (0.60)=\frac{1}{1+e^{-0.60}}\approx 0.6456$$
   • $\text{Loss}=-\log (0.6456)=0.4375$.
2. Execute GAE Advantage Simulation: Run AlignmentMathSimulator.computeGAE([1, 0, 0, 10], [2, 2, 2, 0]) and trace how the terminal $+10$ reward propagates upstream to earlier tokens via the discount factor.

11. Trade-offs & Alignment Architecture Alternatives

12. Authoritative References & Foundational Papers

• John Schulman et al. (2017), “Proximal Policy Optimization Algorithms”, arXiv — arxiv.org/abs/1707.06347
• Long Ouyang et al. (2022), “Training language models to follow instructions with human feedback”, NeurIPS ‘22 — arxiv.org/abs/2203.02155
• Rafael Rafailov et al. (2023), “Direct Preference Optimization: Your Language Model is Secretly a Reward Model”, NeurIPS ‘23 — arxiv.org/abs/2305.18290
• Philipp Moritz et al. (2018), “Ray: A Distributed Framework for Emerging AI Applications”, OSDI ‘18 — usenix.org/conference/osdi18/presentation/moritz
• Version Context: Implementations target Ray 2.35+ with DeepSpeed 0.14+ and vLLM 0.5+.