RL4LMs: A Modular RL Library for Fine-tuning Language Models

Introduction

RL4LMs is an open-source, modular reinforcement learning (RL) library developed by AllenAI, specifically designed for fine-tuning language models (LMs) to align with human preferences. It provides a comprehensive toolkit for researchers and practitioners to experiment with various RL algorithms, reward functions, and metrics across diverse Natural Language Processing (NLP) tasks. The library is backed by extensive testing, with over 2000 experiments conducted on the GRUE benchmark, ensuring robustness and reliability. For more details, you can refer to their paper and website.

Installation

Getting started with RL4LMs is straightforward. You can clone the repository and install it locally:

git clone https://github.com/allenai/RL4LMs.git
cd RL4LMs
pip install -e .

For development using Docker containers with all dependencies pre-installed, use:

docker build . -t rl4lms

Optionally, for certain metric computations like SPICE, CoreNLP libraries are required:

cd rl4lms/envs/text_generation/caption_metrics/spice && bash get_stanford_models.sh

Examples

RL4LMs provides a simple training API that leverages YAML configuration files for quick experimentation. For instance, to train a T5-base model on CNN/DM summarization using PPO with Rouge-1 as the reward function, you can execute:

python scripts/training/train_text_generation.py --config_path scripts/training/task_configs/summarization/t5_ppo.yml

The configuration files allow you to define various building blocks such as datasets, tokenizers, reward functions, environments, and on-policy algorithms. For example, a dataset configuration might look like this:

datapool:
  id: cnn_daily_mail
  args:
    prompt_prefix: "Summarize: "

You can find config files for all supported tasks in the repository's scripts/training/task_configs directory.

Why Use RL4LMs

RL4LMs stands out due to its exceptional modularity and customizability. It offers:

• Comprehensive Building Blocks: Implementations of on-policy algorithms (PPO, A2C, TRPO, NLPO), various reward functions (lexical, semantic, task-specific), and metrics.
• Extensive Task Support: Benchmarked on 7 different NLP tasks, including summarization, machine translation, and dialogue generation.
• Flexible Policies: Supports both causal LMs (e.g., GPT-2/3) and seq2seq LMs (e.g., T5, BART) as actor-critic policies.
• Customizability: Users can easily add their own datasets, reward functions, evaluation metrics, on-policy algorithms, and actor-critic policies by sub-classing provided base classes and registering them.
• Robust Logging and Experiment Tracking: Integrates with WANDB and supports warm-starting, storing checkpoints and various training artifacts for reproducible research.

This makes RL4LMs an ideal framework for advancing research in reinforcement learning for natural language processing.

Links

• GitHub Repository: https://github.com/allenai/RL4LMs
• Paper: https://arxiv.org/abs/2210.01241
• Official Website: https://rl4lms.apps.allenai.org/
• Slack Channel: https://join.slack.com/t/slack-1sa3880/shared_invite/zt-1idqlnbnm-NIiZeMIOpYReXfX9uIT_PA