Implementation:Hiyouga LLaMA Factory PPO Trainer
| Knowledge Sources | |
|---|---|
| Domains | Machine Learning, RLHF, Reinforcement Learning |
| Last Updated | 2026-02-06 19:00 GMT |
Overview
Custom PPO trainer implementing the complete online RLHF training loop with response generation and reward computation in LLaMA-Factory.
Description
CustomPPOTrainer inherits from both TRL's PPOTrainer and HuggingFace's Trainer to implement a full Proximal Policy Optimization training loop for reinforcement learning from human feedback. Unlike DPO/KTO which use offline preference data, PPO performs online learning by generating responses from the policy model, scoring them with a reward model, and optimizing the policy against the reward signal. The ppo_train method implements the complete training loop: iterating over batches, calling get_inputs to generate model responses, get_rewards to compute reward scores (via local reward model with LoRA adapter switching, separate full reward model, or external API), and then performing PPO optimization steps. It supports DeepSpeed, custom optimizers, bf16 training via autocast context, configurable experience buffers, and proper checkpoint saving with value head separation.
Usage
Instantiated by the PPO training workflow when stage="ppo" is set in FinetuningArguments. Requires a reward_model to be specified. Supports reward models of type "lora" (shares base model, switches adapters), "full" (separate model), or "api" (external reward service). Requires TRL version >=0.8.6 and <=0.9.6.
Code Reference
Source Location
- Repository: Hiyouga_LLaMA_Factory
- File: src/llamafactory/train/ppo/trainer.py
- Lines: 1-518
Signature
class CustomPPOTrainer(PPOTrainer, Trainer):
def __init__(
self,
model_args: "ModelArguments",
training_args: "Seq2SeqTrainingArguments",
finetuning_args: "FinetuningArguments",
generating_args: "GeneratingArguments",
callbacks: Optional[list["TrainerCallback"]],
model: "AutoModelForCausalLMWithValueHead",
reward_model: Optional["AutoModelForCausalLMWithValueHead"],
ref_model: Optional["AutoModelForCausalLMWithValueHead"],
tokenizer: "PreTrainedTokenizer",
processor: Optional["ProcessorMixin"],
data_collator: "DataCollatorWithPadding",
train_dataset: Optional["Dataset"] = None,
eval_dataset: Optional["Dataset"] = None,
) -> None: ...
def ppo_train(self, resume_from_checkpoint: Optional[str] = None) -> None:
"""Implement training loop for the PPO stage."""
def create_optimizer(self, model, training_args, finetuning_args) -> "torch.optim.Optimizer": ...
def create_scheduler(self, training_args, num_training_steps, optimizer) -> "torch.optim.lr_scheduler.LRScheduler": ...
@torch.no_grad()
def get_inputs(self, batch) -> tuple[list["torch.Tensor"], list["torch.Tensor"]]:
"""Generate model responses given queries."""
@torch.no_grad()
def get_rewards(self, queries, responses) -> list["torch.Tensor"]:
"""Compute scores using given reward model."""
def batched_forward_pass(
self, model, queries, responses, model_inputs,
return_logits=False, response_masks=None,
) -> tuple["torch.Tensor", Optional["torch.Tensor"], "torch.Tensor", "torch.Tensor"]: ...
def save_model(self, output_dir: Optional[str] = None) -> None: ...
Import
from llamafactory.train.ppo.trainer import CustomPPOTrainer
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| model | AutoModelForCausalLMWithValueHead | Yes | Policy model with value head for PPO optimization |
| reward_model | AutoModelForCausalLMWithValueHead or None | Conditional | Reward model (required unless reward_model_type="api") |
| ref_model | AutoModelForCausalLMWithValueHead or None | No | Reference model for KL penalty computation |
| model_args | ModelArguments | Yes | Model configuration (upcast_layernorm, infer_backend, etc.) |
| finetuning_args | FinetuningArguments | Yes | PPO-specific: ppo_epochs, ppo_buffer_size, ppo_target, ppo_score_norm, reward_model_type |
| generating_args | GeneratingArguments | Yes | Generation parameters (temperature, top_p, max_new_tokens, etc.) |
| tokenizer | PreTrainedTokenizer | Yes | Tokenizer for encoding/decoding |
| train_dataset | Dataset | Yes | Training dataset providing prompt input_ids |
Outputs
| Name | Type | Description |
|---|---|---|
| ppo_train | None | Side effect: trains the model in-place, saves checkpoints, logs metrics |
| get_inputs | tuple[list[Tensor], list[Tensor]] | (queries, responses) generated by the policy model |
| get_rewards | list[Tensor] | Reward scores for each query-response pair (float32) |
| logged metrics | dict | loss, reward, learning_rate, epoch logged at each logging step |
Usage Examples
from llamafactory.train.ppo.trainer import CustomPPOTrainer
# Instantiated by the PPO workflow (simplified)
trainer = CustomPPOTrainer(
model_args=model_args,
training_args=training_args,
finetuning_args=finetuning_args,
generating_args=generating_args,
callbacks=callbacks,
model=policy_model_with_value_head,
reward_model=reward_model_with_value_head,
ref_model=reference_model,
tokenizer=tokenizer,
processor=processor,
data_collator=data_collator,
train_dataset=train_dataset,
)
# Run PPO training loop
trainer.ppo_train()
# The training loop:
# 1. Generate responses from policy model (get_inputs)
# 2. Score responses with reward model (get_rewards)
# 3. Compute PPO loss and update policy (self.step)
# 4. Log metrics and save checkpoints periodically