File size: 2,235 Bytes
dfb775d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | """RLHF (PPO) trainer — TRL PPOTrainer wrapper.
Online preference optimization with a learned reward model. Less efficient
than DPO for offline pairs; kept here as the canonical surface for online
RLHF runs.
"""
from __future__ import annotations
from pathlib import Path
from typing import Any
from mindxtrain.config.schema import XTrainConfig
def run_rlhf(cfg: XTrainConfig, out_dir: Path) -> Path:
"""Run a PPO/RLHF fine-tune; return the checkpoint directory."""
try:
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import PPOConfig, PPOTrainer
except ImportError as exc:
msg = "TRL stack not installed; run `uv sync --extra ml`."
raise RuntimeError(msg) from exc
method: Any = cfg.train.method
if getattr(method, "kind", "") != "rlhf":
msg = f"run_rlhf expects train.method.kind == 'rlhf'; got {getattr(method, 'kind', None)!r}"
raise ValueError(msg)
reward_model_path = getattr(method, "reward_model_path", "")
if not reward_model_path:
msg = "rlhf requires train.method.reward_model_path"
raise ValueError(msg)
out_dir = Path(out_dir)
out_dir.mkdir(parents=True, exist_ok=True)
tokenizer = AutoTokenizer.from_pretrained(cfg.model.name)
policy = AutoModelForCausalLM.from_pretrained(cfg.model.name)
ref_policy = AutoModelForCausalLM.from_pretrained(cfg.model.name)
reward_model = AutoModelForCausalLM.from_pretrained(reward_model_path)
train_ds = load_dataset(cfg.data.hf_id, split=getattr(cfg.data, "split", "train"))
ppo_cfg = PPOConfig(
output_dir=str(out_dir),
learning_rate=cfg.train.optim.learning_rate,
per_device_train_batch_size=cfg.train.micro_batch_size,
gradient_accumulation_steps=cfg.train.gradient_accumulation_steps,
num_train_epochs=cfg.train.num_epochs,
)
trainer = PPOTrainer(
args=ppo_cfg,
model=policy,
ref_model=ref_policy,
reward_model=reward_model,
train_dataset=train_ds,
processing_class=tokenizer,
)
trainer.train()
trainer.save_model(str(out_dir))
return out_dir
__all__ = ["run_rlhf"]
|