trl

TRL: Transformer Reinforcement Learning

Overview

TRL provides a full stack of tools for training language models with reinforcement learning, including SFT (Supervised Fine-Tuning), Reward Modeling, RLHF (PPO), and DPO (Direct Preference Optimization). Apply this skill for instruction tuning, RLHF training, preference learning, and alignment training.

When to Use This Skill

This skill should be used when:

• Supervised fine-tuning (SFT) on instruction data
• Training with RLHF (Reinforcement Learning from Human Feedback)
• Implementing PPO training for LLM alignment
• Direct Preference Optimization (DPO)
• Reward model training
• Preference learning
• Chat model training
• Model alignment with human preferences

Quick Start

Basic Import and Setup

from trl import SFTTrainer, DPOTrainer, RewardTrainer
from trl.trainer import PPOTrainer, RLHFTrainer
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

Supervised Fine-Tuning (SFT)

from trl import SFTTrainer
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset

# Load model and tokenizer
model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")

# Load dataset
dataset = load_dataset("imdb", split="train[:1000]")

# Configure SFT trainer
trainer = SFTTrainer(
    model=model,
    train_dataset=dataset,
    tokenizer=tokenizer,
    dataset_text_field="text",
    max_seq_length=512,
)

# Train
trainer.train()

# Save
trainer.save_model("sft_output")

DPO (Direct Preference Optimization)

from trl import DPOTrainer
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset

# Load model
model = AutoModelForCausalLM.from_pretrained("gpt2")
ref_model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")

# Format dataset for DPO
# Dataset should have: prompt, chosen, rejected
dataset = load_dataset("Anthropic/hh-rlhf", split="train[:1000]")

def format_dpo(example):
    return {
        "prompt": example["prompt"],
        "chosen": example["chosen"],
        "rejected": example["rejected"]
    }

dataset = dataset.map(format_dpo)

# Configure DPO trainer
dpo_trainer = DPOTrainer(
    model=model,
    ref_model=ref_model,
    train_dataset=dataset,
    tokenizer=tokenizer,
    beta=0.1,  # DPO temperature
    max_length=512,
    max_prompt_length=128,
)

# Train
dpo_trainer.train()

# Save
dpo_trainer.save_model("dpo_output")