Step 6: Train with DPOTrainer

Section 7 of 19-~ 7 min read-Synced from Cuantum content

Now comes the core training step where you'll use the TRL library's DPOTrainer to align your model using the preference pairs you've carefully labeled. This is where your human judgments are translated into actual parameter updates that shift your model's behavior.

Understanding the DPO training architecture

DPO training requires two models working in tandem: a policy model (the one being trained) and a reference model (frozen for comparison). The reference model is typically initialized from the same base model checkpoint you're starting with—in this project, that's TinyLlama-1.1B-Chat-v1.0.

Why do we need both? The reference model serves as a stability anchor. During training, DPO compares the policy model's current outputs against what the reference model would have produced. This comparison prevents the policy model from drifting too far from its original capabilities while learning your preferences. Without this anchor, the model might overfit to your preference pairs in ways that degrade its general language modeling ability or introduce unexpected behavior on prompts outside your training distribution.

Think of the reference model as representing "what the model naturally wants to say" before alignment, while the policy model learns to adjust those tendencies based on your preference signal. The training objective mathematically balances two goals: (1) increase the probability of chosen responses relative to rejected ones, and (2) don't diverge too much from the reference model's probability distribution.

Setting up the training code

Here's the complete training script with both models initialized:

from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArgumentsfrom trl import DPOTrainer MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=True)policy_model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map="auto")ref_model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map="auto") training_args = TrainingArguments(    output_dir="outputs/ch3_dpo_chatbot",    per_device_train_batch_size=2,    gradient_accumulation_steps=8,    learning_rate=5e-5,    num_train_epochs=2,    logging_steps=10,    save_strategy="epoch",    report_to="none",    fp16=True) trainer = DPOTrainer(    model=policy_model,    ref_model=ref_model,    args=training_args,    beta=0.1,    train_dataset=dataset,    tokenizer=tokenizer) trainer.train()trainer.save_model("outputs/ch3_dpo_chatbot/final")tokenizer.save_pretrained("outputs/ch3_dpo_chatbot/final")

Breaking down the training arguments

The TrainingArguments control the mechanics of the training loop. Let's examine the key settings for this project:

per_device_train_batch_size=2: Processes 2 preference pairs at a time per GPU. This is kept small because each pair requires forward passes through both the policy and reference models, doubling memory usage compared to standard fine-tuning.
gradient_accumulation_steps=8: Accumulates gradients over 8 batches before updating weights, giving an effective batch size of 16. This provides more stable gradient estimates without requiring more GPU memory.
learning_rate=5e-5: A conservative learning rate appropriate for preference alignment. DPO is generally more sensitive to learning rate than supervised fine-tuning—too high and you'll see instability or divergence from the reference model; too low and alignment will be imperceptible.
num_train_epochs=2: Two passes through your preference dataset. For small datasets (100-300 pairs), this provides enough exposure without severe overfitting. If you have more data (1000+ pairs), you might reduce this to 1 epoch.
fp16=True: Enables mixed-precision training to reduce memory usage and speed up computation. Essential for running this on consumer GPUs.

The beta parameter: controlling alignment strength

The beta parameter in DPOTrainer is the single most important hyperparameter for preference alignment. It controls how aggressively the model learns from your preference judgments:

Lower beta (e.g., 0.05): Gentler alignment that stays closer to the reference model. The policy model will shift its behavior only slightly in the direction of your preferences. Use this when you want subtle refinements or when you're concerned about maintaining the model's existing capabilities.
Moderate beta (e.g., 0.1-0.2): Balanced alignment that provides clear behavior changes while maintaining reasonable stability. This is the recommended starting point for most projects. beta=0.1 typically produces noticeable improvements in alignment without introducing instability.
Higher beta (e.g., 0.3+): Stronger preference shaping that can produce dramatic behavior changes. The risk here is overfitting to your specific preference pairs or diverging so far from the reference model that you lose general capabilities. Reserve this for cases where you have high-quality, diverse preference data and need strong alignment.

The mathematical intuition: beta scales the KL divergence penalty between the policy and reference models. Lower beta means a stronger penalty for diverging from the reference, keeping changes conservative. Higher beta relaxes this penalty, allowing more aggressive preference learning.

What happens during training

When you call trainer.train(), the DPO algorithm processes each preference pair like this:

The policy model generates log probabilities for both the chosen and rejected responses
The reference model generates log probabilities for the same responses (but its parameters stay frozen)
DPO computes a loss that increases the policy model's probability of the chosen response relative to the rejected one, while penalizing divergence from the reference model based on beta
Gradients are backpropagated only through the policy model to update its parameters

This process repeats for every preference pair in your dataset, across the specified number of epochs. The logging output will show you the loss decreasing over time—this indicates the model is learning to distinguish chosen from rejected responses according to your labels.

Monitoring training progress

Watch for these signs of healthy training:

Loss should decrease steadily but not collapse to near-zero (which would indicate overfitting)
If you included a validation split, validation loss should track training loss without diverging significantly
Training should complete without out-of-memory errors (if you get OOM, reduce batch size)

Warning signs to watch for:

Loss increases or becomes erratic: your learning rate may be too high
Loss decreases to nearly zero very quickly: you're likely overfitting, especially if your dataset is small
Training is extremely slow: you may need to reduce sequence length or batch size, or enable gradient checkpointing

After training completes

The final lines save both your aligned model and tokenizer to outputs/ch3_dpo_chatbot/final. This directory now contains everything you need to load and use your aligned chatbot, just like any other HuggingFace model. The key difference is that this model's response distribution has been reshaped by your preference judgments—it should now be more likely to generate responses with the characteristics you marked as "chosen" and less likely to produce responses like those you marked as "rejected."

For most first runs with 100-300 preference pairs, beta=0.1 is a sensible starting point. You can always train a second version with adjusted beta based on your evaluation results in Step 7.