RLHF efficace : Réduction de l'utilisation de la mémoire dans PPO

Efficient RLHF: Reducing the Memory Usage of PPO

September 1, 2023

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

cs.AI

Résumé

L'apprentissage par renforcement avec feedback humain (RLHF) a révolutionné la modélisation du langage en alignant les modèles sur les préférences humaines. Cependant, l'étape de renforcement, l'Optimisation de Politique Proximale (PPO), nécessite plus de 3 fois la mémoire utilisée pour le Fine-Tuning Supervisé (SFT), la rendant impraticable pour la plupart des praticiens. Pour résoudre ce problème, nous présentons une analyse approfondie de l'utilisation de la mémoire, des performances et du temps d'entraînement des techniques d'économie de mémoire pour PPO. Nous introduisons Hydra-RLHF en intégrant d'abord les modèles SFT et de récompense, puis en désactivant dynamiquement LoRA pendant l'entraînement. Nos expériences montrent : 1. L'utilisation de LoRA pendant PPO réduit son utilisation de mémoire à un niveau inférieur à celui de SFT tout en améliorant l'alignement sur quatre benchmarks publics, et 2. Hydra-PPO réduit la latence par échantillon de LoRA-PPO jusqu'à 65 % tout en maintenant ses performances. Nos résultats démontrent que Hydra-PPO est une solution simple et prometteuse pour permettre une utilisation plus large de RLHF.

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.