RLHF Performance Optimization: PPO Algorithm Acceleration Techniques

This article discusses performance optimization techniques for Reinforcement Learning from Human Feedback (RLHF), focusing on the PPO algorithm. The author addresses the challenge that RLHF training throughput is significantly lower than pre-training or SFT due to its complex multi-stage pipeline involving four models: Actor, Critic, Reward, and Reference.

The article presents three main optimization strategies:

1. Text Generation Speed Optimization : Using NVIDIA's TRT-LLM framework to accelerate inference. The authors identified three key bottlenecks: KV cache memory limitations affecting batch size, inefficient generation due to uneven prompt lengths, and model parallelism overhead from training. TRT-LLM addresses these through paged attention, in-flight batching, and flexible model deployment. They also proposed a refit solution for online model updates, reducing parameter synchronization time from 15 minutes to 20 seconds.

2. Dynamic Memory Optimization : Implementing selective activation recomputation combined with sequence parallelism (from Megatron-LLM), which reduced activation memory by 50% with only 20% speed degradation. Additional optimizations include micro-batch padding and temporary memory management following two principles: early release and avoiding over-allocation.

3. System Parallel Optimization : Identifying parallelization opportunities in the PPO pipeline, including parallel execution of ref_logp and logp computations, and overlapping text generation with reward calculation.

Experimental results on a 10B parameter model with 16x A100 (80GB) GPUs showed baseline throughput of 0.012 samples/gpu/s. After optimizations, throughput improved to 0.054 samples/gpu/s, representing a 350% improvement. The baseline iteration time of 13,376 seconds was reduced to 295 seconds.