rlhf_step2_reward.yaml

# RLHF Step 2: Train Reward Model
# Learn to score responses: is this output good or bad?
# Uses preference pairs (prompt / chosen / rejected).
#
# The bundled dataset is a 5-row FORMAT sample, not a training set.
# See examples/data/README.md.
#
# Usage:
#   soup train --config examples/configs/rlhf_step2_reward.yaml

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

data:
  train: examples/data/chat_preferences.jsonl  # must have chosen/rejected
  format: dpo
  # 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-4
  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

output: ./output_reward_model/