SkillCompiler/data/skills-bench/tasks/debug-trl-grpo/environment/skills/trl/references/trl-codebase.md

TRL Codebase — Detailed Module Guide

Table of Contents

• Trainer Layer
• Utility Module (trainer/utils.py)
• Models Layer
• Data Utilities
• Configuration Details

Trainer Layer

Common Trainer Pattern

Every TRL trainer follows this pattern:

1. Extends transformers.Trainer (inherits training loop, checkpointing, logging)
2. Defines a custom compute_loss method for its objective
3. Optionally overrides training_step for generation-based methods
4. Uses a corresponding *Config dataclass for hyperparameters

SFTTrainer

Standard supervised fine-tuning. Key features:

• Dataset packing (multiple examples per sequence for efficiency)
• Configurable max_seq_length and dataset_text_field
• Supports instruction-tuning formats via formatting_func

DPOTrainer

Direct Preference Optimization — learns from preference pairs without RL.

• Requires paired (chosen, rejected) examples
• beta controls KL constraint strength
• Multiple loss variants: sigmoid (default), hinge, ipo, kto_pair
• Can run reference-free (no frozen model needed)

GRPOTrainer

Group Relative Policy Optimization — see the grpo skill for algorithm details.

• Overrides training_step to add generation phase
• _generate_and_score_completions handles sampling + reward scoring
• compute_loss implements clipped surrogate + KL objective

KTOTrainer

Kahneman-Tversky Optimization — learns from unpaired binary feedback.

• Does not require paired preferences (just good/bad labels)
• desirable_weight / undesirable_weight control loss asymmetry
• Uses a KL term estimated from a reference model

OnlineDPOTrainer

Online variant of DPO that generates completions during training.

• Similar generation loop to GRPO
• Pairs completions by reward ranking for DPO-style loss
• Combines benefits of online generation with DPO's stability

Utility Module (trainer/utils.py)

selective_log_softmax

Memory-efficient per-token log-probability computation. Takes logits [B, T, V] and an index tensor [B, T], and returns per-token log-probabilities [B, T] — the same values F.log_softmax(logits, -1).gather(-1, index.unsqueeze(-1)).squeeze(-1) would produce, without materializing the vocab-sized intermediate tensor.

Used by: GRPOTrainer, DPOTrainer, OnlineDPOTrainer — any trainer needing per-token log probabilities.

Read trainer/utils.py directly before modifying. Key invariants the implementation must satisfy are listed in the rl-post-training skill.

decode_and_strip_padding

Converts token ID tensors to the text strings passed to the reward function. Handles padding, decoder artefacts, and any reasoning-block conventions the library supports.

Used by: GRPOTrainer, OnlineDPOTrainer — any trainer that generates and decodes completions.

Before changing behavior here, enumerate the completion shapes the model can emit and confirm each is handled; the current policy for complete, incomplete, and absent reasoning blocks is defined in the implementation.

Padding Utilities

• pad(tensors, padding_value, padding_side) — Pad list of variable-length tensors to uniform length
• pad_to_length(tensor, length, padding_value) — Pad or truncate to exact length

Models Layer

AutoModelForCausalLMWithValueHead

Wraps a causal LM with an additional linear head that outputs scalar values. Used by PPO-style trainers that need a critic (not used by GRPO, DPO, or SFT).

PreTrainedModelWrapper

Base class for model wrappers that need to modify forward pass behavior while preserving the underlying model's API.

Data Utilities

data_utils.py

• maybe_extract_prompt — Extracts prompt from conversation-format datasets
• apply_chat_template — Applies tokenizer chat template to conversation data
• Dataset formatting helpers for various training paradigms

Configuration Details

All configs inherit standard TrainingArguments fields (learning rate, batch size, gradient accumulation, etc.) and add trainer-specific parameters.

GRPOConfig Specifics

Parameter	Typical Value	Purpose
num_generations	4-16	Completions per prompt for advantage estimation
max_completion_length	256-1024	Max tokens per completion
beta	0.01-0.1	KL penalty coefficient
epsilon	0.1-0.2	Clipping range for surrogate objective
temperature	0.7-1.0	Sampling temperature for generation
reward_functions	list	Callable reward functions

DPOConfig Specifics

Parameter	Typical Value	Purpose
beta	0.1	KL constraint strength
loss_type	"sigmoid"	DPO loss variant
reference_free	False	Whether to skip reference model
label_smoothing	0.0	Smoothing for preference labels