# Test configuration for reward model training
# Usage: python train_reward_model.py configs/reward_model_test.py

# Data paths (same data as DPO training)
data_dir = "/path/to/dpo/data"
out_dir = "/path/to/reward/model/checkpoints"

# Model checkpoint to start from (pre-trained GPT checkpoint)
preload_checkpoint = "/path/to/pretrained/model.pt"
preload_optimizer = False  # Don't load optimizer state, start fresh
preload_strict = False  # Allow missing reward head parameters

# Reward model specific
freeze_base_model = False  # Set to True to only train reward head
use_reward_head = True  # Enable reward head

# Training hyperparameters
batch_size = 8  # Must be even (for paired preferences)
gradient_accumulation_steps = 1
learning_rate = 1e-5  # Lower LR for reward model
min_lr = 1e-6
warmup_iters = 2_000
max_iters = 50_000
weight_decay = 0.01

# Evaluation
eval_interval = 1_000
eval_iters = 100
log_interval = 10

# System
device = "cuda"
dtype = "bfloat16"
compile = False
fsdp = False

# Logging
wandb_log = True
wandb_project = "suno-reward-model"
wandb_run_name = "reward-model-v1"

# Data settings
suppress_text = False
shuffle_data = False
local_shuffle_data = False

# Save checkpoints
checkpoint_save_old_format = True
step_save_iters = 10_000