rlhf

name	rlhf
description	Set up and run the full RLHF pipeline (SFT, reward model training, RL from reward model) using the Tinker API. Use when the user wants to do RLHF, train a reward model, or run the full preference-based RL pipeline.
argument-hint	[model-name]

RL from Human Feedback (RLHF) Pipeline

Help the user set up and run the full 3-stage RLHF pipeline using the Tinker API.

Overview

RLHF is a multi-stage pipeline:

1. SFT Stage — Fine-tune base model on instruction data
2. Reward Model (RM) Stage — Train a reward model on preference comparisons
3. RL Stage — Optimize the SFT policy using the reward model

Step 1: Understand the request

Ask the user (if not already specified):

• Base model: Which model to start from (e.g., meta-llama/Llama-3.2-3B)
• Preference data: Which comparison dataset (HHH, HelpSteer3, UltraFeedback, or custom)
• Which stages to run: All 3, or skip SFT/RM if checkpoints exist
• LoRA rank: Typically 64 for RLHF

Step 2: Reference existing recipes

Read these files:

• tinker_cookbook/recipes/preference/rlhf/rlhf_pipeline.py — Complete 3-stage pipeline
• tinker_cookbook/rl/preference_envs.py — Preference-based RL environments
• tinker_cookbook/preference/types.py — PreferenceModelBuilder
• tinker_cookbook/preference/comparison_policy_evaluator.py — RM evaluation
• docs/preferences/rlhf-example.mdx — RLHF guide

Step 3: Configure each stage

Stage 1: SFT

Standard supervised fine-tuning (see /sft skill). Key settings:

• Dataset: NoRobots or similar instruction data
• sft_learning_rate: 2e-4 (LoRA)
• train_on_what: TrainOnWhat.ALL_ASSISTANT_MESSAGES

Stage 2: Reward Model

Train on preference comparisons:

• Uses ChatDatasetBuilderFromComparisons with a comparison builder (e.g., HHHComparisonBuilder)
• rm_learning_rate: 3e-4
• Produces a reward model checkpoint used in Stage 3

Stage 3: RL from Reward Model

Optimize SFT policy using RM scores:

• Load SFT checkpoint as starting policy
• Load RM weights for scoring
• PreferenceModelBuilderFromChatRenderer wraps the RM
• PairwisePreferenceRLDatasetBuilder creates the RL environment
• rl_learning_rate: 1e-5 (much lower than SFT)
• tournament_pattern: ALL_PAIRS_BOTH_WAYS for pairwise comparison

Typical Hyperparameters

@chz.chz
class CLIConfig:
    base_model: str = "meta-llama/Llama-3.2-3B"
    lora_rank: int = 64
    batch_size: int = 256
    max_length: int = 16384
    sft_learning_rate: float = 2e-4
    rm_learning_rate: float = 3e-4
    rl_learning_rate: float = 1e-5
    rl_max_tokens: int = 1024
    rl_group_size: int = 4

Step 4: Write the training script

Follow the pipeline pattern from rlhf_pipeline.py:

import asyncio
import os
import chz
from tinker_cookbook import checkpoint_utils, model_info
from tinker_cookbook.preference.types import PreferenceModelBuilderFromChatRenderer
from tinker_cookbook.rl import preference_envs, train
from tinker_cookbook.supervised import train as supervised_train

# Stage 1: SFT
def sft_stage(log_path, base_model, ...):
    # Standard SFT config + supervised_train.main()
    ...

# Stage 2: Reward Model
def train_rm(log_path, base_model, ...):
    # Train on preference comparisons
    ...

# Stage 3: RL
async def train_rl(log_path, sft_log_path, rm_log_path, base_model, ...):
    sft_checkpoint = checkpoint_utils.get_last_checkpoint(sft_log_path)["state_path"]
    rm_weights = checkpoint_utils.get_last_checkpoint(rm_log_path)["sampler_path"]

    preference_model_builder = PreferenceModelBuilderFromChatRenderer(
        renderer_name=renderer_name,
        model_name=base_model,
        rm_weights_path=rm_weights,
    )
    rl_dataset_builder = preference_envs.PairwisePreferenceRLDatasetBuilder(
        comparison_builder=comparison_builder,
        preference_model_builder=preference_model_builder,
        batch_size=batch_size,
        group_size=group_size,
        tournament_pattern=preference_envs.TournamentPattern.ALL_PAIRS_BOTH_WAYS,
        ...
    )
    config = train.Config(
        load_checkpoint_path=sft_checkpoint,
        dataset_builder=rl_dataset_builder,
        learning_rate=1e-5,
        ...
    )
    await train.main(config)

Step 5: Run

# Full pipeline
python -m tinker_cookbook.recipes.preference.rlhf.rlhf_pipeline

# Skip SFT (already have checkpoint)
python -m tinker_cookbook.recipes.preference.rlhf.rlhf_pipeline run_sft=False

# Skip SFT and RM
python -m tinker_cookbook.recipes.preference.rlhf.rlhf_pipeline run_sft=False run_rm=False

Step 6: Add tests

If you created a new RLHF recipe, add a smoke test:

# tests/recipes/test_recipe_<name>.py
import pytest
from tests.helpers import run_recipe

@pytest.mark.integration
def test_<recipe_name>():
    run_recipe(
        "tinker_cookbook.recipes.<recipe_name>.train",
        ["behavior_if_log_dir_exists=delete"],
    )

run_recipe() automatically passes max_steps=2 so the recipe runs 2 training steps and exits. See tests/recipes/test_recipe_rlhf_pipeline.py for the existing example.

Common pitfalls

• RL learning rate must be much lower than SFT (1e-5 vs 2e-4)
• Checkpoints flow between stages: SFT → RL policy init, RM → RL reward scoring
• Use checkpoint_utils.get_last_checkpoint() to find checkpoints from previous stages
• RM quality directly impacts RL — validate RM before running Stage 3
• group_size in RL stage affects variance of reward estimates