Fine-tuning 20B LLMs with RLHF on a 24GB Consumer GPU

TL;DR

Hugging Face announced that the trl library now works with peft and 8‑bit quantization, allowing RLHF fine‑tuning of 20 B parameter LLMs on a single 24 GB consumer GPU. This makes large‑scale RL fine‑tuning affordable and accessible without multi‑GPU model‑parallel setups.

Why RLHF Needs Efficient Fine‑Tuning

Reinforcement Learning with Human Feedback (RLHF) typically follows three stages: (1) supervised fine‑tuning on instructions, (2) training a reward model from human annotations, and (3) PPO‑based RL fine‑tuning using the reward model. The RL step requires two copies of the model (active and reference) on each GPU, which quickly exceeds the memory of a single device for models larger than 10 B parameters.

TRL: A Library for PPO‑Based RL

trl provides a high‑level API for PPO training of language models. It leverages 🤗 Accelerate to run on single‑device or distributed setups. The PPO loop needs both an active model (being updated) and a reference model (kept frozen) to compute KL‑regularized rewards, effectively doubling the memory footprint.

Reducing Memory Footprint with PEFT and 8‑Bit Quantization

8‑Bit Matrix Multiplication

  • 8‑bit quantization (LLM.int8()) stores weights in 1 byte per parameter, cutting model size by a factor of four compared to float32.
  • The method splits each linear layer into an outlier‑handling float16 part and a bulk int8 part, preserving accuracy while gaining speed.

Low‑Rank Adaptation (LoRA) via PEFT

  • LoRA freezes the pretrained weights and injects low‑rank matrices (A and B) into the query and value projections of the attention blocks.
  • Only the adapter parameters are trainable, reducing optimizer memory dramatically.
  • The forward and backward passes are roughly twice as slow because of extra matrix multiplications, but the memory savings enable training of 20 B models on consumer hardware.

End‑to‑End Pipeline for a 20 B Model on a 24 GB GPU

Step 1 – Load the Model in 8‑Bit Precision

model = AutoModelForCausalLM.from_pretrained(
    "EleutherAI/gpt-neox-20b",
    load_in_8bit=True,
    device_map="auto",
)

Loading in 8‑bit reduces memory from ~80 GB (float32) to ~20 GB, fitting comfortably on a 24 GB card.

Step 2 – Attach Trainable LoRA Adapters with PEFT

from peft import get_peft_model, LoraConfig
config = LoraConfig(r=8, lora_alpha=32, target_modules=["q_proj", "v_proj"], bias="none")
model = get_peft_model(model, config)

Only the low‑rank matrices are stored in optimizer state, cutting optimizer memory from several gigabytes to a few hundred megabytes.

Step 3 – Use a Single Model for Reference and Active Logits

PEFT’s disable_adapters context manager temporarily deactivates LoRA layers, allowing the same underlying model to produce reference logits:

with model.disable_adapter():
    ref_logits = model(input_ids)
# active logits are computed with adapters enabled
active_logits = model(input_ids)

No second full model copy is needed, further reducing memory usage.

Training Scripts Overview

The blog post links three scripts that demonstrate the full workflow on a 20 B GPT‑NeoX model:

  1. clm_finetune_peft_imdb.py – Causal language‑model fine‑tuning of LoRA adapters on the IMDB sentiment dataset (one epoch).
  2. merge_peft_adapter.py – Merges LoRA weights into the base model for inference or further training.
  3. gpt-neo-20b_sentiment_peft.py – PPO fine‑tuning using an IMDB sentiment classifier as the reward model to generate positive movie reviews.

All scripts were executed on an NVIDIA RTX 4090 (24 GB). Full training runs were also tested on a single A100 in the 🤗 research cluster.

Results

  • The loss curve shows stable convergence after one epoch of supervised LoRA fine‑tuning on IMDB.
  • During PPO, the mean reward steadily increases, indicating the model learns to produce more positive reviews.
  • The entire pipeline runs on a single 24 GB GPU, demonstrating that RLHF is no longer limited to multi‑GPU clusters.

Implications for the Community

  • Lower barrier to entry – Researchers and developers can experiment with RLHF on consumer‑grade hardware.
  • Open‑source reproducibility – All code and adapters are hosted on the Hugging Face Hub, enabling easy sharing of fine‑tuned artifacts.
  • Scalable foundation – The same approach can be extended to larger models with data parallelism once multi‑GPU support is added.

Open Questions and Future Work

  • Multi‑GPU scaling – How well does the integration work with data parallelism across several GPUs?
  • Training speed – LoRA adds overhead; exploring faster kernels or mixed‑precision strategies could mitigate this.
  • Broader RL algorithms – While PPO is the default, integrating other RL methods (e.g., DPO) may broaden applicability.

References

Sources