trl-training

TRL Training Skill

You are an expert at using the TRL (Transformers Reinforcement Learning) library to train and fine-tune large language models.

Overview

TRL provides CLI commands for post-training foundation models using state-of-the-art techniques:

• SFT (Supervised Fine-Tuning): Fine-tune models on instruction-following or conversational datasets
• DPO (Direct Preference Optimization): Align models using preference data
• GRPO (Group Relative Policy Optimization): Train models by ranking multiple sampled outputs relative to each other and optimizing based on their comparative rewards.
• RLOO (Reinforce Leave One Out): Online RL training with generation-based rewards
• Reward Model Training: Train reward models for RLHF

TRL is built on top of Hugging Face Transformers and Accelerate, providing seamless integration with the Hugging Face ecosystem.

Core Commands

trl sft - Supervised Fine-Tuning

Fine-tune language models on instruction-following or conversational datasets.

Full training:

trl sft \
  --model_name_or_path Qwen/Qwen2-0.5B \
  --dataset_name trl-lib/Capybara \
  --learning_rate 2.0e-5 \
  --num_train_epochs 1 \
  --packing \
  --per_device_train_batch_size 2 \

Train with LoRA adapters:

trl dpo - Direct Preference Optimization

Align models using preference data (chosen/rejected pairs).

Full training:

Train with LoRA adapters:

trl grpo - Group Relative Policy Optimization

Train models using reward functions or LLM-as-a-judge for evaluating generations and providing rewards.

Basic usage:

trl rloo - Reinforce Leave One Out

Online RL training where the model generates text and receives rewards based on custom criteria.

Basic usage:

trl reward - Reward Model Training

Train a reward model to score text quality for RLHF.

Full training:

Train with LoRA adapters:

Configuration Files

TRL supports YAML configuration files for reproducible training. All CLI arguments can be specified in a config file.

Example config (sft_config.yaml):

Launch with config:

Override config values:

Distributed Training

TRL integrates with Accelerate for multi-GPU and multi-node training.

Multi-GPU training:

Use predefined Accelerate configs:

TRL provides predefined configs: single_gpu, multi_gpu, fsdp1, fsdp2, zero1, zero2, zero3

Custom Accelerate config:

Fully Sharded Data Parallel (FSDP):

DeepSpeed ZeRO:

Troubleshooting

CUDA Out of Memory

• Reduce --per_device_train_batch_size and increase --gradient_accumulation_steps
• Enable --use_peft for LoRA training
• Use --gradient_checkpointing to save memory
• Try smaller model or longer sequence truncation

Dataset Loading Issues

• Verify dataset exists: check Hugging Face Hub or local path
• Check dataset format matches expected columns
• Use --dataset_config for multi-config datasets
• Inspect dataset: from datasets import load_dataset; ds = load_dataset(name)

Model Loading Issues

• Verify model exists on Hugging Face Hub
• Check if gated model requires authentication: hf auth login
• For local models, provide absolute path
• Ensure sufficient disk space and memory

Slow Training

• Enable dataset --packing for short sequences
• Use larger --per_device_train_batch_size if memory allows
• Enable --tf32 for faster computation on Ampere GPUs
• Use --bf16 on supported hardware
• Consider multi-GPU training with --num_processes

Generation Issues (GRPO/RLOO)

• Check prompt format in dataset
• Adjust --temperature and --top_p for generation
• Verify the reward function (for GRPO/RLOO)

Additional Resources

• Documentation: https://huggingface.co/docs/trl
• GitHub: https://github.com/huggingface/trl
• Examples: https://github.com/huggingface/trl/tree/main/examples

Best Practices

1. Start with SFT: Always fine-tune base models with SFT before preference alignment
2. Use LoRA for efficiency: Enable --use_peft for faster training and lower memory
3. Monitor training: Use --report_to trackio (or --report_to wandb or --report_to tensorboard) for tracking
4. Save checkpoints: TRL automatically saves checkpoints in --output_dir
5. Test on small datasets first: Verify pipeline works before full training
6. Use configuration files: Create YAML configs for reproducibility
7. Leverage Accelerate: Use multi-GPU training for faster iteration

When helping users with TRL:

• Always check which training method is appropriate for their use case
• Verify dataset format matches the expected schema
• Recommend starting with smaller models for testing
• Suggest LoRA for resource-constrained environments
• Point to specific documentation sections for advanced features