rlhf_step3_ppo.yaml

# RLHF Step 3: PPO Training
# Optimize the policy with reinforcement learning against the reward model
# trained in step 2. Run rlhf_step2_reward.yaml first — `training.reward_model`
# points at its output directory.
#
# The bundled dataset is a 10-row FORMAT sample, not a training set.
# See examples/data/README.md.
#
# Usage:
#   soup train --config examples/configs/rlhf_step3_ppo.yaml

base: TinyLlama/TinyLlama-1.1B-Chat-v1.0
task: ppo
backend: transformers

data:
  train: examples/data/alpaca_tiny.jsonl
  format: alpaca
  # 512, not 2048: the longest row in any bundled fixture is ~233 tokens, so
  # nothing is truncated and there is 2x headroom. A larger value only inflates
  # the activation memory the VRAM pre-flight budgets for. Raise it when you
  # point this config at longer data of your own.
  max_length: 512

training:
  epochs: 2
  lr: 1e-5
  batch_size: 4
  gradient_accumulation_steps: 2
  quantization: none
  scheduler: cosine
  warmup_ratio: 0.1
  weight_decay: 0.01
  logging_steps: 10
  save_steps: 50
  ppo_epochs: 4
  ppo_clip_ratio: 0.2
  ppo_kl_penalty: 0.05
  reward_model: ./output_reward_model/
  lora:
    r: 64
    alpha: 128
    dropout: 0.05
    target_modules:
      - q_proj
      - v_proj

output: ./output_ppo/