RLHF Eficiente: Reduzindo o Uso de Memória do 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

Resumo

O Aprendizado por Reforço com Feedback Humano (RLHF) revolucionou a modelagem de linguagem ao alinhar modelos com as preferências humanas. No entanto, a etapa de RL, o Proximal Policy Optimization (PPO), requer mais de 3 vezes a memória do Ajuste Fino Supervisionado (SFT), tornando-o inviável para a maioria dos profissionais. Para resolver esse problema, apresentamos uma análise abrangente do uso de memória, desempenho e tempo de treinamento de técnicas de economia de memória para o PPO. Introduzimos o Hydra-RLHF, primeiro integrando os modelos SFT e de Recompensa e, em seguida, desativando dinamicamente o LoRA durante o treinamento. Nossos experimentos mostram: 1. O uso do LoRA durante o PPO reduz seu uso de memória para menos que o SFT, enquanto melhora o alinhamento em quatro benchmarks públicos, e 2. O Hydra-PPO reduz a latência por amostra do LoRA-PPO em até 65%, mantendo seu desempenho. Nossos resultados demonstram que o Hydra-PPO é uma solução simples e promissora para permitir um uso mais amplo do 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.