Efficiënte RLHF: Het geheugengebruik van PPO verminderen

Efficient RLHF: Reducing the Memory Usage of PPO

September 1, 2023

Auteurs: Michael Santacroce, Yadong Lu, Han Yu, Yuanzhi Li, Yelong Shen

cs.AI

Samenvatting

Reinforcement Learning met Menselijke Feedback (RLHF) heeft een revolutie teweeggebracht in taalmodellering door modellen af te stemmen op menselijke voorkeuren. Echter vereist de RL-fase, Proximal Policy Optimization (PPO), meer dan 3x het geheugen van Supervised Fine-Tuning (SFT), wat het voor de meeste praktijkbeoefenaars onhaalbaar maakt. Om dit probleem aan te pakken, presenteren we een uitgebreide analyse van het geheugengebruik, de prestaties en de trainingsduur van geheugenbesparende technieken voor PPO. We introduceren Hydra-RLHF door eerst de SFT- en beloningsmodellen te integreren en vervolgens LoRA dynamisch "uit" te schakelen tijdens de training. Onze experimenten tonen aan: 1. Het gebruik van LoRA tijdens PPO vermindert het geheugengebruik tot minder dan dat van SFT, terwijl de afstemming over vier openbare benchmarks wordt verbeterd, en 2. Hydra-PPO vermindert de latentie per sample van LoRA-PPO met tot wel 65% terwijl de prestaties behouden blijven. Onze resultaten laten zien dat Hydra-PPO een eenvoudige en veelbelovende oplossing is om een breder gebruik van RLHF mogelijk te maken.

English

Reinforcement Learning with Human Feedback (RLHF) has revolutionized language modeling by aligning models with human preferences. However, the RL stage, Proximal Policy Optimization (PPO), requires over 3x the memory of Supervised Fine-Tuning (SFT), making it infeasible to use for most practitioners. To address this issue, we present a comprehensive analysis the memory usage, performance, and training time of memory-savings techniques for PPO. We introduce Hydra-RLHF by first integrating the SFT and Reward models and then dynamically turning LoRA "off" during training. Our experiments show: 1. Using LoRA during PPO reduces its memory usage to be smaller than SFT while improving alignment across four public benchmarks, and 2. Hydra-PPO reduces the latency per sample of LoRA-PPO by up to 65% while maintaining its performance. Our results demonstrate that Hydra-PPO is a simple and promising solution for enabling more widespread usage of RLHF.