PaperlayerView exact revision on arXiv

Source-linked AI summary

Efficient RLHF: Reducing the Memory Usage of PPO

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

arXiv:2309.00754v1cs.LGcs.AIcs.CL

TL;DR

PPO-based RLHF is complex and memory-intensive, limiting its use. Hydra-RLHF shares models and dynamically switches LoRA during PPO; the paper reports maintained performance with up to 65% faster per-sample latency.

  • Problem

    RLHF-PPO is complex and memory-intensive because multiple models must be loaded and trained, limiting its usage.

    Source: §1 · p. 1

  • Method

    Hydra-RLHF shares models and dynamically switches the active LoRA module during PPO.

    Source: §1 · p. 1, §3 · p. 3, §6 · p. 8

  • Results

    Hydra-RLHF saves memory while maintaining performance, with up to 65% faster per-sample latency when the saved memory supports larger batches.

    Source: §6 · p. 8

  • Takeaways & Limitations

    Hydra-RLHF may open RLHF use to a wider variety of models and applications.

    Source: §6 · p. 8

    PPO performance can be highly inconsistent and may decline through reward-model exploitation, which the authors observe on StackExchange.

    Source: §4 · p. 4

Abstract

from arXiv · show

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.

1 Introduction

PPO makes RLHF complex and memory-intensive, while the need to evaluate its speed and performance motivates this study. The authors identify model sharing and propose Hydra-PPO to reduce memory use and per-sample latency.

1 Introduction

  • RLHF-PPO is complex and memory-intensive because it loads and trains multiple models, limiting its usage.The authors also call for evaluating RLHF variations in speed and performance.Source: §1 · p. 1
  • The study identifies model sharing between Reference/Reward Models and Actor/Critic Models as an opportunity to reduce memory and computation costs.Source: §1 · p. 1
  • Up to 65% lower per-sample PPO latency is reported using memory savings to increase training batch size.Source: §1 · p. 1
  • Table 1 compares memory usage and runtime across methods for Llama 7b on StackExchange per A100 80GB GPU.Full Fine-Tuning PPO memory is a scaled-up estimate.Source: Table 1 · p. 2

2 RLHF

The section describes standard RLHF as a three-stage process: supervised fine-tuning, reward-model training, and PPO. In the common setup, PPO uses reference, actor, reward, and critic models, with LoRA applied to the actor and value function.

2 RLHF

  • SFT trains an input language model on causal language modeling data to produce πSFT.Training all parameters is called FFT-SFT, while using LoRA is called LoRA-SFT.Source: §2 · p. 2
  • Reward-model training replaces the language-model head with a scalar output trained to predict human preferences from prompt-completion pairs.The reward is often normalized to improve PPO stability.Source: §2 · p. 2
  • The common PPO setup uses a frozen reference model, a trained actor, a frozen reward model, and a trained critic initialized from SFT and reward models.The paper uses the most common and simplest four-model setup.Source: §2 · p. 2
  • LoRA-PPO applies LoRA to all linear layers of the actor and value function, and the authors omit Full Fine-Tuning PPO experiments because of its extreme cost.The PPO objective clips the output probability ratio, while generalized advantage estimation uses the value function to form advantage estimates.Source: §2 · p. 3

3 Hydra-RLHF

Hydra-RLHF combines causal and reward heads in πhydra, then uses LoRA sharing to reduce the number of models held in memory during PPO. The section also describes Hydra-SFT’s data requirements and a joined-weight PPO ablation.

Stage 1: Hydra-SFT

  • Hydra-SFT trains πhydra with next-token cross-entropy plus a weighted pairwise ranking loss; γ = 0.1 generally works well.Source: §3 · p. 3
  • Hydra-SFT requires pairwise preference data, and ranked lists should be paired only with the best-ranked sample to avoid training the SFT head on other samples.Source: §3 · p. 3

Dynamic LoRA

  • About 20% memory is saved by recovering reference and reward models with LoRA disabled, while maintaining performance equivalent to LoRA-PPO.The reference is recovered from the actor, and the reward model from the critic.Source: §3 · p. 3

Stage 2: Hydra-PPO

  • Hydra-PPO uses separate actor and critic LoRA weights on one πhydra base, requiring only one full base model and giving similar memory usage to LoRA fine-tuning at the same batch size.Source: §3 · p. 3
  • Joined Hydra-PPO shares one LoRA-weight set between actor and critic, saving a small amount of memory and runtime but performing worse than Hydra-PPO.By contrast, joining the models in Hydra-SFT does not affect performance.Source: §3 · p. 3

4 Experiments

Across datasets, PPO improves alignment over its base model on average, and Hydra-PPO performs better than LoRA-PPO for Llama 7b. The experiments also compare methods across tasks and examine runtime as sequence length increases.

4 Experiments

  • Experiments compare five approaches across four datasets using GPT-4 evaluation, with ROUGE scores also used for summarization.LoRA-PPO starts from SFT, while J-Hydra-PPO and Hydra-PPO start from Hydra-SFT.Source: §4 · pp. 3–4
  • PPO outperforms SFT on average, Hydra-PPO similarly improves Hydra-SFT, and Hydra-PPO achieves better alignment than LoRA-PPO for Llama 7b.Table 3 reports Llama 7b expected aggregate win-rates; PPO performance can be inconsistent, and reward-model flaws can reduce it on StackExchange.Source: §4 · p. 4, Table 3 · p. 4

4.3 Learning to Summarize

  • For Learning to Summarize, SFT-based methods typically achieve better precision, while PPO-based methods substantially improve recall.The reported ROUGE results align with the win-rate findings; the authors suggest longer PPO generations may contribute to the recall pattern.Source: §4 · p. 4, Table 4 · p. 4

4 Experiments

  • Hydra-PPO saves exponentially more time as sequence length increases, while its latency advantage over J-Hydra-PPO narrows at length 1024.Latency sums inference and parameter-update time per sample; LoRA-PPO does not fit in memory at context length 2048 in this setup.Source: §4 · p. 5, Figure 2 · p. 5

4.1 GPT-4-LLM

  • GPT-4-LLM produced the most consistent and well-behaved training runs among the datasets tested.The dataset pairs its highest-scoring response with each other response, and the study attempts full RLHF on it.Source: §4.1 · p. 5

4.2 Open-Source Assistant Datasets

  • The Open-Source Assistant Datasets combine several assistant-style datasets, including one with Helpful & Harmless data, used without modification.The included datasets are hosted on HuggingFace.Source: §4.2 · p. 5

4.4 StackExchange

  • On StackExchange, PPO and J-Hydra-PPO often learn to repeat answers, whereas Hydra-PPO does not.The dataset uses 150k samples and pairs the best answer with up to three others; its questions are diverse and specific.Source: §4.4 · p. 6, Table 8 · p. 6

4.5 Changing Model Size and Family

  • For OPT-1.3b, Hydra-SFT performs worse than SFT, and LoRA-PPO has better overall alignment than Hydra-PPO.The authors speculate that combining language and reward models may be more difficult at this model size; PPO methods still improve over their respective base models.Source: §4.5 · p. 6, Table 9 · p. 7

5 Related Works

The section reviews alignment via human-preference reward models and alternatives to PPO, then distinguishes Hydra-SFT from approaches that replace PPO. It also states dataset and reward-model-size conditions relevant to Hydra-RLHF.

5 Related Works

  • Alignment modifies foundation models’ behavior for human interaction, while reward models trained on preference labels can guide models toward those preferences.The cited work reports that such reward-model guidance can outperform SFT across many domains.Source: §5 · p. 7
  • Other feedback work uses foundation models to train or evaluate models, while RAFT, RRHF, PRO, and DPO combine preference data with supervised fine-tuning to replace or reframe PPO.RAFT and RRHF draw on best-of-n sampling; PRO and DPO reframe the supervised objective.Source: §5 · p. 7
  • Hydra-SFT integrates ranked feedback into supervised fine-tuning, but Hydra-RLHF is presented as orthogonal to PPO-replacement methods and aims to make RLHF more usable.Source: §5 · pp. 7–8
  • Hydra-RLHF requires identical SFT and reward-model training datasets; the experiments use pairwise comparisons and report no overfitting, while noting possible extensions for limited data.Source: §5 · p. 8
  • The experiments keep reward and language models the same size for fair comparisons, though memory savings are smaller when standard RLHF uses a smaller reward model.In applied use, Hydra-RLHF can instead use a larger reward model at lower training cost.Source: §5 · p. 8

6 Conclusion

The conclusion presents Hydra-RLHF as a memory-saving PPO approach that combines reference and reward models and dynamically switches LoRA, with up to 65% faster per-sample latency. The appendix provides pseudocode for Hydra-RLHF and J-Hydra-RLHF.

6 Conclusion

  • Hydra-RLHF combines reference and reward models and dynamically switches the active LoRA module; it can use a higher batch size and reduce per-sample latency by up to 65% while maintaining performance.Source: §6 · p. 8
  • The conclusion identifies balancing SFT and reward-model datasets, improving J-Hydra-PPO performance, and improving PEFT methods for RLHF as future work.Source: §6 · p. 8

A Algorithm Pseudocode

  • The appendix supplies J-Hydra-RLHF and Hydra-RLHF pseudocode, including steps that compute reward and reference log probabilities, calculate advantages, and optimize the surrogate objective.Source: A Algorithm Pseudocode · p. 12
  • The pseudocode snippets include retaining critic-head values for the final sequence step and turning on the critic LoRA to calculate critic values.Source: A Algorithm Pseudocode · p. 12

B Experiment Details

Experiments ran on an 8×A100 80GB node with shared SFT, reward-model, and PPO settings. Runtime measurements use Llama 7b on StackExchange, with a scaled OPT 1.3b run used when full PPO exceeded the setup.

B Experiment Details

  • Experiments ran on an 8×A100 80GB node with DeepSpeed ZeRO-1 and a cosine-decaying learning-rate schedule; the hyperparameter tables cover Llama 7b and OPT 1.3b.Batch size is per GPU, so total batch size equals 8 × batch size × gradient accumulation steps.Source: B Experiment Details · p. 12, Table 12 · p. 13, Table 13 · p. 14
  • SFT, reward-model, and Hydra-SFT runs use weight decay 0.1, normalized rewards, and validation-selected models after learning-rate sweeps; PPO variants share KL β = 0.02, GAE γ = 1.0, and λ = 0.95.LoRA rank is 128 for SFT and for PPO actors and critics; PPO variants run for one epoch.Source: B Experiment Details · p. 12
  • Runtime measurements use Llama 7b on StackExchange with sequence length 800; because full PPO overflows the setup, OPT 1.3b results are scaled using the Llama 7b LoRA-PPO ratio.Memory was measured with PyTorch, and total per-sample latency sums the two PPO stages.Source: B Experiment Details · p. 12
  • J-Hydra-PPO is described as highly unstable and requiring multiple attempts to find solid hyperparameters, whereas Hydra-PPO is generally as stable as LoRA-PPO.Source: B Experiment Details · p. 13

C GPT-4 Grading

The grading appendix reports that reversing two answers in a direct comparison could substantially change the preferred answer. It therefore uses repeated 0–9 ratings with step-by-step reasoning, then compares averaged scores side by side.

C GPT-4 Grading

  • Reversing answer order in a two-answer preference prompt substantially changed which answer was preferred.Source: C GPT-4 Grading · p. 14
  • The chosen evaluation grades each answer separately on a 0–9 scale with step-by-step reasoning, then compares grades side by side; each output is rated three times at temperature 1 and averaged.Source: C GPT-4 Grading · p. 14
  • Grading cleanup removes EOT tokens and stops generation when a model begins speaking for the other side or repeats a 5-gram phrase; evaluation uses 500 validation samples per dataset.Source: C GPT-4 Grading · pp. 14–15
  • The grading prompt asks for an integer score out of 9 and a one-sentence evaluation, judging helpfulness, precision, concision, grammar, friendliness, and non-confrontational language.Source: C GPT-4 Grading · p. 15

Insert conversation and question

Insert answer to judge

The section supplies a prompt for GPT-4 to evaluate a summary, with a one-sentence assessment and a score out of nine. It then provides placeholders for the original text and its summary.

Insert answer to judge

Insert original text

D Reward Model Grading

The reward-model grading section reports that reward models rate each PPO model as strongly superior to its respective input model. The next section notes that LoRA-SFT underperforms FFT-SFT and cautions that LoRA does not always match its performance.

D Reward Model Grading

  • Reward models rate each PPO model as strongly superior to its respective input model.Table 14 compares FFT-SFT with LoRA-PPO using a separated reward model, and Hydra-SFT with J-Hydra-PPO and Hydra-PPO using Hydra-RM.Source: D Reward Model Grading · p. 16, Table 14 · p. 16

E LoRA-SFT vs. LoRA-PPO

  • LoRA-SFT underperforms FFT-SFT, and LoRA does not always match FFT-SFT performance.Table 15 reports GPT-4-judged win rates for LoRA-SFT versus LoRA-PPO with FFT-SFT as the base; no cell values are supplied here.Source: E LoRA-SFT vs. LoRA-PPO · p. 16, Table 15 · p. 16
  • If a method improves LoRA-SFT for alignment, the full RLHF process could use roughly the same footprint as LoRA-SFT.The passage says further investigation into PEFT for alignment is required.Source: E LoRA-SFT vs. LoRA-PPO · p. 16

F Pre-PPO Performance

The pre-PPO performance results cover SFT, reward-model, and Hydra-SFT models on their respective validation datasets. The Hydra-SFT reward-model head usually outperforms the standard reward model.

F Pre-PPO Performance

Loading 2309.00754v1…