Understanding and Alleviating Memory Consumption in RLHF for LLMs

Jin Zhou, Hanmei Yang, Steven, Tang, Mingcan Xiang, Hui Guan, Tongping Liu

Preprint: not peer reviewed. Results may change in later versions.

Abstract

Fine-tuning with Reinforcement Learning with Human Feedback (RLHF) is essential for aligning large language models (LLMs). However, RLHF often encounters significant memory challenges. This study is the first to examine memory usage in the RLHF context, exploring various memory management strategies and unveiling the reasons behind excessive memory consumption. Additionally, we introduce a simple yet effective approach that substantially reduces the memory required for RLHF fine-tuning.

Abstract: arXiv descriptive metadata (CC0). Thank you to arXiv for use of its open access interoperability. This service was not reviewed or approved by, nor does it necessarily express or reflect the policies or opinions of, arXiv.

Record

Archive ID
arxiv:2410.15651
Published
Oct 21, 2024
arXiv
2410.15651
Categories
cs.LG
Provenance
arxiv