RLHF Eficiente: Reducción del Uso de Memoria en PPO

Efficient RLHF: Reducing the Memory Usage of PPO

September 1, 2023

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

cs.AI

Resumen

El Aprendizaje por Refuerzo con Retroalimentación Humana (RLHF) ha revolucionado el modelado del lenguaje al alinear los modelos con las preferencias humanas. Sin embargo, la etapa de RL, Optimización de Políticas Proximales (PPO), requiere más de 3 veces la memoria del Ajuste Supervisado (SFT), lo que hace inviable su uso para la mayoría de los profesionales. Para abordar este problema, presentamos un análisis exhaustivo del uso de memoria, el rendimiento y el tiempo de entrenamiento de técnicas de ahorro de memoria para PPO. Introducimos Hydra-RLHF integrando primero los modelos SFT y de Recompensa y luego desactivando dinámicamente LoRA durante el entrenamiento. Nuestros experimentos muestran: 1. El uso de LoRA durante PPO reduce su uso de memoria a un nivel menor que SFT mientras mejora la alineación en cuatro puntos de referencia públicos, y 2. Hydra-PPO reduce la latencia por muestra de LoRA-PPO hasta en un 65% manteniendo su rendimiento. Nuestros resultados demuestran que Hydra-PPO es una solución simple y prometedora para permitir un uso más generalizado 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.