RLHF with TRLX: GPT‑2 Sentiment and Summarization Demos

This repository contains small, reproducible examples of training and evaluating GPT‑2 with TRLX (CarperAI) using Reinforcement Learning from Human Feedback (RLHF) techniques:

• PPO training pipeline for summarization (scripts)
• ILQL training to generate positive‑sentiment movie reviews (notebook)
• Inference and interactive chat comparison (base GPT‑2 vs PPO‑trained model)

The code is designed to run on a single machine (CPU or GPU), with conservative defaults to reduce memory usage.

---

Contents

├── summarize_rlhf/
│   ├── chat_inference.py              # Compare GPT‑2 base vs trained PPO model in an interactive chat
│   ├── demo_chat.py                   # Demo: run a fixed set of prompts through both models
│   ├── minimal_rlhf_demo.py           # Minimal, CPU‑friendly PPO training with tiny config
│   ├── ppo_with_reward_scores.csv     # Example output file written by inference
│   ├── trlx_gpt2_text_summarization.py# PPO training for TL;DR summarization (length‑based reward)
│   └── trlx_inference_gpt2.py         # Batch inference + ROUGE/evaluation utilities
├── trlx_sentiments.ipynb              # ILQL: positive IMDB sentiment generation with GPT‑2
└── trlx_simulacra.ipynb               # ILQL example (Simulacra) – for reference

---

Requirements

• Python 3.9+ (3.10 recommended)
• pip and virtual environments
• Optional: CUDA‑capable GPU and a compatible PyTorch build

Key Python packages used:

• trlx (installed from GitHub)
• torch, transformers, datasets, accelerate
• evaluate, pandas, tqdm

---

Quickstart (Windows PowerShell)

1. Create and activate a virtual environment

python -m venv .venv
.\.venv\Scripts\Activate.ps1
python -m pip install --upgrade pip

2. Install dependencies

pip install "git+https://github.com/CarperAI/trlx.git"
pip install torch transformers datasets accelerate evaluate pandas tqdm

Note: For GPU support, install the correct PyTorch build for your CUDA version from pytorch.org’s selector, then install the remaining packages.

3. Run the minimal PPO demo (CPU‑friendly)

python summarize_rlhf\minimal_rlhf_demo.py

This performs a tiny PPO training run on a few prompts to verify your setup.

---

PPO Summarization Training

The script below trains GPT‑2 with PPO on the CarperAI TL;DR dataset and uses a simple length‑based reward (kept small to avoid OOM on single GPUs/CPUs):

python summarize_rlhf\trlx_gpt2_text_summarization.py

Outputs/checkpoints are expected under ckpts/ (as configured by TRLX). The inference scripts look for ckpts/best_checkpoint/hf_model by default.

Important notes:

• The script intentionally skips loading a heavy reward model and uses a lightweight, length‑based reward for stability.
• If you see ModuleNotFoundError: reward_model, you can safely comment out the import line from reward_model.reward_model import GPTRewardModel near the top of the file (the reward model block is already disabled; the import is not required for the length‑based path).
• GPU usage can be toggled by CUDA_VISIBLE_DEVICES environment variable inside the script. Set it to an empty string to force CPU if needed.

---

Inference and Evaluation

After training, you can run batch inference and compute ROUGE on a small split:

python summarize_rlhf\trlx_inference_gpt2.py

This loads the trained PPO model from ckpts/best_checkpoint/hf_model, generates summaries, prints ROUGE metrics, and writes ppo_with_reward_scores.csv with simple reward scores.

Interactive chat comparison (base GPT‑2 vs trained PPO model):

python summarize_rlhf\demo_chat.py      # runs a fixed set of example prompts
python summarize_rlhf\chat_inference.py # interactive prompt loop

Both scripts will load two models: gpt2 and your trained model at ckpts/best_checkpoint/hf_model.

---

ILQL Sentiment Notebook

• trlx_sentiments.ipynb demonstrates using ILQL with TRLX to produce positive‑sentiment movie reviews on the IMDB dataset.
• It sets up a sentiment reward model (lvwerra/distilbert-imdb) and trains GPT‑2 to maximize positive sentiment.
• Open the notebook in Jupyter or run it on Colab. The installation cells will install TRLX and required libraries.

trlx_simulacra.ipynb is provided as an additional ILQL reference notebook.

---

Tips & Troubleshooting

• Out‑of‑memory (OOM): reduce batch sizes, sequence lengths, and number of rollouts; force CPU as a sanity check; close other GPU processes.
• Tokenizer warnings: GPT‑2 has no pad token by default; the scripts set pad_token to EOS and use left padding for generation.
• Checkpoints: make sure the path ckpts/best_checkpoint/hf_model exists after training before running inference/chat scripts.
• Windows PowerShell execution policy: if your venv activation is blocked, run PowerShell as Administrator and: Set-ExecutionPolicy RemoteSigned -Scope CurrentUser.

---

Acknowledgments

• TRLX: https://github.com/CarperAI/trlx
• Datasets and models via Hugging Face ecosystem

---

License

This repository contains example code built on top of TRLX and Hugging Face libraries. See the upstream projects for their respective licenses.