[2.4.1] - RLHF

Colab: exercises | solutions

Please send any problems / bugs on the #errata channel in the Slack group, and ask any questions on the dedicated channels for this chapter of material.

If you want to change to dark mode, you can do this by clicking the three horizontal lines in the top-right, then navigating to Settings → Theme.

Links to all other chapters: (0) Fundamentals, (1) Transformer Interpretability, (2) RL.

Introduction

🚧 Work in progress: this material is new and still being polished, so don't treat it as final.

This section takes you through a full implementation of RL fine-tuning of a language model, the "RL" in RLHF (Reinforcement Learning from Human Feedback). We'll use GRPO (Group Relative Policy Optimization), the algorithm introduced in the DeepSeekMath paper and used to train DeepSeek-R1, which has become the standard algorithm for RL on language models. Unlike PPO, GRPO needs no critic (value network): we instead estimate the advantage by running many rollouts, scoring each of them, and using the average score obtained by the group as our baseline from which we can estimate advantages. This is more efficient than using PPO, which would require modifying the LLM and grafting on a value head to the end of the residual stream, and then training it to try and estimate how much the reward model would score a particular generation. An older version of this material did exactly this, but GRPO (and other such methods) are more standard to use now.

We'll be working directly with Hugging Face transformers models, rather than transformer lens which is not optimized for performance. We start by discussing how the RL framework of states, actions and rewards fits onto autoregressive language models, then build GRPO up piece by piece (sampling, rewards, advantages, and the clipped objective), and put everything together into a GRPO_Trainer class. In the second half, you'll implement LoRA (Low-Rank Adaptation), a Parameter Efficient Fine-Tuning (PEFT) method, from scratch, and use it to RL fine-tune a model that would otherwise be too big to train on a single GPU.

You don't need to have done the PPO material to do this section, although if you have, you'll recognise a few pieces (the clipped surrogate objective in particular).

Note - these exercises assume you're running on a GPU with at least 24GB of memory (e.g. an RTX 4090, A10, L4 or A100). If you're running on a machine with less than 22 GiB of VRAM (or with no GPU at all), the variable LOW_GPU_MEM will be set to True automatically, and we'll switch to smaller models. You won't miss out on learning anything about RLHF, but you'll be less able to play around with training larger models on more interesting reward functions.

For a lecture on the material today, which provides some high-level understanding before you dive into the material, watch the video below:

Content & Learning Objectives

1️⃣ RLHF via GRPO

Most of the exercises in this section build towards the implementation of the GRPO_Trainer class.

Learning Objectives
  • Understand how the RL agent / action / environment paradigm works in the context of autoregressive language models
  • Understand GRPO: group-relative advantages, the clipped surrogate objective, and the KL penalty to a reference model
  • Understand why we estimate the KL penalty from sampled tokens
  • Write a full RLHF training loop, and train against a toy reward function.
  • Observe and understand reward hacking and mode collapse
  • Experiment with a more complex reward function (a sentiment classifier)

2️⃣ LoRA

Learning Objectives
  • Understand the mechanism behind Low-Rank Adaptation, and why it allows fine-tuning with far less memory (both in terms of memory saved during the optimization step, and memory saved on the reference model).
  • Implement LoRA from scratch for a Hugging Face model, by wrapping its nn.Linear layers
  • Use GRPO + LoRA together to RL fine-tune a model otherwise too large to fully fine-tune on a single GPU

☆ Bonus

This section offers some suggested ways to extend the core exercises.

Learning Objectives
  • Compare your LoRA to the peft library's
  • Learn about variants of GRPO, and libraries like TRL which implement it at scale
  • Train your own reward model from human preferences

Optional Reading

Setup code

import gc
import math
import os
import sys
import time
from contextlib import contextmanager
from dataclasses import dataclass, field
from functools import partial
from pathlib import Path
from typing import Callable, Literal

# Allow expandable memory segments on CUDA to avoid OOMs from fragmentation
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import einops
import matplotlib.pyplot as plt
import numpy as np
import torch as t
import torch.nn as nn
import wandb
from jaxtyping import Float, Int
from tabulate import tabulate
from torch import Tensor
from tqdm.auto import tqdm
from eindex import eindex
from transformers import (
    AutoModelForCausalLM,
    AutoModelForSequenceClassification,
    AutoTokenizer,
    PreTrainedModel,
    PreTrainedTokenizerBase,
)

# Make sure exercises are in the path
chapter = "chapter2_rl"
section = "part4_rlhf"
root_dir = next(p for p in Path.cwd().parents if (p / chapter).exists())
exercises_dir = root_dir / chapter / "exercises"
section_dir = exercises_dir / section
if str(exercises_dir) not in sys.path:
    sys.path.append(str(exercises_dir))

from part4_rlhf import tests, tests_lora

# PyTorch defaults to one thread per core. On machines that report many more cores than they let you use (common in
# cloud containers), CPU-side work then spends longer waiting on threads than computing. Adjust if needed.
t.set_num_threads(min(4, t.get_num_threads()))

device = t.device("mps" if t.backends.mps.is_available() else "cuda" if t.cuda.is_available() else "cpu")
# TF32 matmuls on NVIDIA GPUs: noticeably faster training (no effect on CPU, MPS or pre-Ampere GPUs)
t.set_float32_matmul_precision("high")


MAIN = __name__ == "__main__"

We'll use two families of models today, both loaded from the Hugging Face Hub:

  • GPT-2 medium (355M parameters, or GPT-2 small with 124M parameters if LOW_GPU_MEM), which we'll fully fine-tune in the first section. It's an old model, but small and fast to sample from, and you know its architecture well from chapter 1.
  • Qwen3 (1.7B parameters, or 0.6B with LOW_GPU_MEM), which we'll fine-tune with LoRA in the second section. Fully fine-tuning a 1.7B model in float32 needs well over 24GB, as we'll see.

Both are base models (pretrained only, no instruction tuning or RLHF).

# LOW_GPU_MEM fires automatically on GPUs with less than 22 GiB (e.g. Colab's 16GB T4), or with no CUDA GPU at all.
LOW_GPU_MEM = not t.cuda.is_available() or t.cuda.get_device_properties(device).total_memory / (1024**3) < 22
BASE_MODEL = "gpt2" if LOW_GPU_MEM else "gpt2-medium"
LORA_MODEL = "Qwen/Qwen3-0.6B-Base" if LOW_GPU_MEM else "Qwen/Qwen3-1.7B-Base"
# bfloat16 halves the memory of the frozen base weights in the LoRA section; GPUs before Ampere (e.g. T4) don't support it
LORA_DTYPE = t.bfloat16 if t.cuda.is_available() and t.cuda.is_bf16_supported() else t.float32
RUN_BASE_RLHF = True
print(f"{LOW_GPU_MEM=}, {BASE_MODEL=}, {LORA_MODEL=}, {LORA_DTYPE=}, {RUN_BASE_RLHF=}")