Rollout Orchestration — infra.bea.ai

Rollout generation is the inference phase of RLHF — producing completions from the policy model. It’s often the bottleneck and requires different optimizations than training.

1. Why Rollouts Are the Bottleneck

Phase	Operation	Arithmetic intensity	Typical time (70B, batch=256)
Generation	Autoregressive decode	Low (memory-bound)	~60-120 seconds
Scoring	Reward model forward	Moderate	~5-10 seconds
Training	Forward + backward	High (compute-bound)	~10-30 seconds

Generation dominates — 60-80% of RLHF wall-clock time. Every token requires a full model forward pass, sequentially.

2. Optimizing Generation Throughput

Use an Inference Engine

Don’t generate with the training framework — use vLLM or SGLang:

from vllm import LLM, SamplingParams

# Load policy model for inference (separate from training instance)
inference_engine = LLM(
    model=policy_model_path,
    tensor_parallel_size=4,
    max_num_seqs=256,  # large batch for throughput
)

sampling_params = SamplingParams(
    temperature=0.8,
    top_p=0.95,
    max_tokens=512,
)

# Generate completions for all prompts
completions = inference_engine.generate(prompts, sampling_params)

Variable-Length Batching

RLHF completions have highly variable lengths — some finish in 10 tokens, others use all 512. Continuous batching reclaims resources from finished sequences:

Without continuous batching:
  All 256 sequences run until the longest finishes → many GPUs idle

With continuous batching (vLLM):
  Sequence finishes → slot freed → new prompt inserted
  → GPU stays busy throughout generation

Multiple Completions per Prompt

GRPO generates K completions per prompt (e.g., K=8) for group-relative ranking:

Prompt "Solve 2+3":
  Completion 1: "5" → reward: 1.0
  Completion 2: "The answer is 5" → reward: 1.0
  Completion 3: "7" → reward: 0.0
  ...
  
Advantage = (reward - group_mean) / group_std

This multiplies generation cost by K — making efficient inference even more critical.

3. GPU Allocation Patterns

Pattern 1: OpenRLHF (Separate Clusters)

Generation cluster: 8 GPUs running vLLM (inference optimized)
Training cluster: 8 GPUs running FSDP (training optimized)
Communication: Transfer completions + rewards via shared storage or network

Pattern 2: DeepSpeed Hybrid Engine (Shared GPUs)

Phase 1: All 8 GPUs in inference mode → generate rollouts
Phase 2: All 8 GPUs switch to training mode → compute gradients
Phase 3: All 8 GPUs in inference mode → next batch of rollouts

Pattern 3: Async Generation (Overlap)

GPUs 0-3: Generate batch N+1 rollouts
GPUs 4-7: Train on batch N rollouts
→ When both finish, swap: 0-3 train on N+1, 4-7 generate N+2

4. Weight Synchronization

After a training step, the updated policy weights must be loaded into the inference engine:

Approach	Latency	Complexity
Save checkpoint, reload in vLLM	~10-30 seconds	Simple
Shared memory (same node)	~1-2 seconds	Moderate
NCCL broadcast	~0.5-1 second	Complex

For reward model hosting, see Reward Model Serving. For the overall RLHF pipeline, see RLHF/GRPO Pipeline.

Self-Check

Explain why rollout generation is the RLHF bottleneck
Configure vLLM for rollout generation with continuous batching
Compare separate-cluster vs hybrid-engine GPU allocation
Describe weight synchronization between training and inference engines