Effizientes RLHF: Reduzierung des Speicherverbrauchs von PPO

Efficient RLHF: Reducing the Memory Usage of PPO

September 1, 2023

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

cs.AI

Zusammenfassung

Reinforcement Learning with Human Feedback (RLHF) hat die Sprachmodellierung revolutioniert, indem es Modelle mit menschlichen Präferenzen in Einklang bringt. Allerdings erfordert die RL-Phase, Proximal Policy Optimization (PPO), mehr als das Dreifache des Speichers im Vergleich zum Supervised Fine-Tuning (SFT), was ihre Anwendung für die meisten Praktiker unpraktikabel macht. Um dieses Problem zu lösen, präsentieren wir eine umfassende Analyse des Speicherverbrauchs, der Leistung und der Trainingszeit von speichereffizienten Techniken für PPO. Wir stellen Hydra-RLHF vor, indem wir zunächst das SFT- und das Belohnungsmodell integrieren und dann LoRA während des Trainings dynamisch „abschalten“. Unsere Experimente zeigen: 1. Die Verwendung von LoRA während PPO reduziert den Speicherverbrauch auf ein Niveau, das unter dem von SFT liegt, und verbessert gleichzeitig die Ausrichtung über vier öffentliche Benchmarks hinweg, und 2. Hydra-PPO reduziert die Latenz pro Sample von LoRA-PPO um bis zu 65 %, während die Leistung erhalten bleibt. Unsere Ergebnisse zeigen, dass Hydra-PPO eine einfache und vielversprechende Lösung ist, um eine breitere Nutzung von RLHF zu ermöglichen.

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.