Navigation
🏋️ Training Systems· 12 min read

RLHF/GRPO Pipeline

Generate rollouts → score with reward model → compute policy gradients → update — the infrastructure of reinforcement learning from human feedback.

rlhfgrpopporeinforcement-learningreward-modelalignment

RLHF aligns language models with human preferences. Unlike SFT (which imitates demonstrations), RLHF optimizes for a reward signal — requiring fundamentally different infrastructure with interleaved inference and training.

1. The RLHF Pipeline

RLHF Pipeline

┌─────────────┐    ┌──────────────┐    ┌────────────────┐    ┌─────────────┐
│ Policy Model │───→│ Generate     │───→│ Score with     │───→│ Compute     │
│ (the LLM)   │    │ Rollouts     │    │ Reward Model   │    │ Policy Grad │
└──────┬──────┘    └──────────────┘    └────────────────┘    └──────┬──────┘
       │                                                            │
       └────────────── Update weights ←─────────────────────────────┘

Each iteration:

  1. Generate: Policy model produces completions for a batch of prompts (inference)
  2. Score: Reward model rates each completion (inference)
  3. Train: Compute PPO/GRPO policy gradients and update the policy (training)

The Infrastructure Challenge

RLHF mixes inference (generation, scoring) with training (gradient computation, optimizer step) in the same pipeline. This requires:

  • Inference-optimized serving for generation (high throughput)
  • Training-optimized setup for gradient computation (high FLOPS utilization)
  • Coordination between the two phases

2. PPO vs GRPO

PPO (Proximal Policy Optimization)

Components needed:
  - Policy model (being trained)
  - Reference model (frozen copy — for KL penalty)
  - Reward model (frozen — scores completions)
  - Value model (trained — estimates expected reward)

Memory: 4× model copies!

For a 70B model: 4 × 140 GB = 560 GB just for BF16 weights. Infrastructure nightmare.

GRPO (Group Relative Policy Optimization)

Components needed:
  - Policy model (being trained)
  - Reference model (can be the policy from last step — no separate copy)
  - Reward model (or rule-based verifier)
  
Memory: 2-3× model copies (no value model)

DeepSeek-R1 uses GRPO — significantly simpler infrastructure:

  • No value model to train
  • Reward comes from group-relative ranking (compare completions within a batch)
  • Compatible with rule-based rewards (math verifiers, code execution)

3. Infrastructure Components

Generation Phase (Inference)

# Generate K completions per prompt
for prompt_batch in prompt_dataloader:
    completions = []
    for _ in range(K):
        # Use inference engine (vLLM) for fast generation
        completion = policy_model.generate(
            prompt_batch,
            max_new_tokens=512,
            temperature=0.8,
            do_sample=True,
        )
        completions.append(completion)

This is the bottleneck — autoregressive generation is slow. Optimizations:

  • Use inference engine (vLLM) for generation phase
  • Batch size matters: larger batches = better GPU utilization during generation
  • Can use speculative decoding during rollout generation

Scoring Phase (Inference)

# Score each completion with reward model
rewards = reward_model.score(prompts, completions)
# Or use rule-based verification
rewards = code_executor.verify(completions)  # correct/incorrect

Training Phase

# Compute policy gradient
advantages = compute_advantages(rewards, values)  # PPO
# or
advantages = compute_group_advantages(rewards)     # GRPO: rank within group

# Policy gradient loss
log_probs = policy_model.log_prob(completions)
log_probs_ref = reference_model.log_prob(completions)
kl = log_probs - log_probs_ref

loss = -advantages * log_probs + beta * kl  # KL-constrained objective
loss.backward()
optimizer.step()

4. GPU Allocation Strategies

StrategyGeneration GPUsTraining GPUsProsCons
Time-multiplexedAll GPUsAll GPUs (alternating)Simple, full utilizationGPUs idle during the other phase
Dedicated poolsPool APool BContinuous operationRequires 2× GPUs
Hybrid EngineAll GPUs (inference mode)All GPUs (training mode)Single cluster, mode switchingComplex mode transitions

DeepSpeed’s Hybrid Engine switches the same GPU cluster between inference and training modes:

  • Inference mode: Tensor parallelism, KV cache, continuous batching
  • Training mode: FSDP, activation checkpointing, gradient computation

For rollout generation infrastructure, see Rollout Orchestration.

Self-Check

  • Draw the RLHF pipeline: generation → scoring → training → update
  • Compare PPO vs GRPO memory requirements (4 vs 2-3 model copies)
  • Explain why RLHF requires both inference and training infrastructure
  • Calculate GPU memory for PPO with a 70B policy model
  • Choose between time-multiplexed and dedicated pool allocation

References