Direct Preference Optimization (DPO)¶
DPO (Direct Preference Optimization) is an elegant simplification of RLHF (Reinforcement Learning from Human Feedback) that makes preference learning more computationally efficient, especially for large language models.
The two key innovations are:
Eliminating the reward model: Instead of training a separate reward model to score outputs (which requires additional compute and memory), DPO directly optimizes the policy using preference data. It reparameterizes the reward function implicitly through the policy itself, deriving a closed-form solution for the optimal policy.
Preference-based optimization: DPO treats the preference learning problem as a classification task over pairs of responses. It maximizes the likelihood that preferred responses are ranked higher than rejected ones under the current policy, relative to a reference policy. This approach eliminates the need for sampling and reward model queries during training.
These changes are particularly valuable for LLM training because they reduce computational overhead by removing the need for a separate reward model and RL training loop, provide more stable training dynamics by avoiding the complexities of reinforcement learning, and they simplify implementation while achieving comparable or better performance than traditional RLHF.
Example¶
from agilerl.algorithms import DPO
from agilerl.llm_envs import DatasetEnv
from datasets import load_dataset
from peft import get_peft_model
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
# Instantiate the model and the associated tokenizer
model = AutoModelForCausalLM.from_pretrained(
"Qwen/Qwen2.5-3B",
torch_dtype=torch.bfloat16,
)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-3B")
# Load the dataset into a preference DatasetEnv
raw_dataset = load_dataset("HumanLLMs/Human-Like-DPO-Dataset", split="train").shuffle(seed=42)
train_test_split = raw_dataset.train_test_split(test_size=0.1)
train_dataset = train_test_split["train"]
test_dataset = train_test_split["test"]
env = DatasetEnv(
train_dataset=train_dataset,
test_dataset=test_dataset,
tokenizer=tokenizer,
objective="preference",
data_batch_size_per_gpu=16,
)
# Instantiate the agent
agent = DPO(
env.observation_space,
env.action_space,
actor_network=model,
pad_token_id=tokenizer.eos_token_id,
pad_token=tokenizer.eos_token,
device="cuda" if torch.cuda.is_available() else "cpu",
batch_size=32,
lr=0.000005,
beta=0.001,
update_epochs=1,
seed=42,
)
Training a DPO agent¶
To train a DPO agent on a single dataset environment, use the train_llm_dataset function:
from agilerl.training.llm import train_llm_dataset
train_llm_dataset(
[agent],
env,
num_epochs=1,
checkpoint_steps=250,
)
Saving and Loading Agents¶
To save an agent, use the save_llm_checkpoint function:
from agilerl.utils.utils import save_llm_checkpoint
save_llm_checkpoint(agent, "path/to/checkpoint")
To load a trained model, you must use the HuggingFace .from_pretrained method, AgileRL is compatible with HuggingFace and Peft models:
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import torch
base_model = AutoModelForCausalLM.from_pretrained(
"Qwen/Qwen2.5-3B",
torch_dtype=torch.bfloat16,
device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-3B")
model = PeftModel.from_pretrained(base_model, "path/to/model/directory")
Parameters¶
- class agilerl.algorithms.dpo.DPO(*args: Any, **kwargs: Any)¶
Direct Preference Optimization (DPO).
Paper: https://arxiv.org/pdf/2305.18290
- Parameters:
pad_token_id (int) – Pad token id
pad_token (str) – Pad token
model_name (str, optional) – Model name
actor_network (PreTrainedModel | PeftModel | None) – HuggingFace LLM
model_config (dict[str, Any] | None) – Model configuration, to be used when creating the model from a name or path.
hp_config (HyperparameterConfig, optional) – RL hyperparameter mutation configuration, defaults to None, whereby algorithm mutations are disabled.
index (int, optional) – Index to keep track of object instance during tournament selection and mutation, defaults to 0
batch_size (int, optional) – Batch size for training, defaults to 16
lr (float, optional) – Learning rate, defaults to 0.000005
beta (float, optional) – DPO beta parameter, defaults to 0.1
nll_alpha (float, optional) – Weight for the NLL loss on chosen responses (DPO + NLL), defaults to 1.0. Set to 0 to disable the NLL term entirely.
max_grad_norm (float, optional) – Maximum gradient norm, defaults to 0.1
update_epochs (int, optional) – Number of update epochs, defaults to 1
calc_position_embeddings (bool, optional) – Flag to indicate if position embeddings should be calculated, defaults to True
micro_batch_size_per_gpu (int, optional) – Micro batch size per GPU, defaults to None
mini_batch_size (int | None, optional) – Per-rank samples covered by one optimizer step.
Noneusesbatch_size / world_size.gradient_accumulation_stepsis derived asmini_batch_size / micro_batch_size_per_gpu.device (str, optional) – Device for accelerated computing, ‘cpu’ or ‘cuda’, defaults to ‘cpu’
lora_config (LoraConfig, optional) – Config for LoRA, defaults to None
fsdp_config (FSDPConfig | None, optional) – FSDP2 sharding settings for distributed runs, defaults to None
wrap (bool, optional) – Wrap models for distributed training upon creation, defaults to True
clone (bool, optional) – Flag to indicate if the instantiation is a cloning, defaults to False
seed (int, optional) – Seed for the random number generator, defaults to 42
gradient_checkpointing (bool, optional) – Flag to indicate if gradient checkpointing should be used, defaults to True
torch_compiler (str | None, optional) – Torch compile mode (e.g.
'default'), defaults to Noneuse_liger_loss (bool, optional) – Use Liger kernel for memory-efficient loss computation. Defaults to
True. PassFalseto opt out (requiresliger-kernelto be installed; warns and falls back toFalseotherwise). Whentraining=Falsethe standard path is always used regardless of this flag.chunk_rows (int | None, optional) – Primary chunk-size knob for fused logit tiles used by both standard and Liger paths.
cast_logprobs_to_fp32 (bool, optional) – When
True(default), run the per-token log-prob reduction (gather/logsumexp) in fp32 before casting back to the input dtype, for numerically stable log-probs.Falseruns it in the input dtype, saving a little memory at the cost of a per-token bf16 quantisation error that can bias importance-sampling ratios.use_separate_reference_adapter (bool, optional) – Keep a dedicated
referenceLoRA adapter whose weights are frozen snapshots of the actor used for the DPO log-probability baseline. WhenFalsethe reference log-probs are obtained by disabling the actor adapter at inference time. Defaults to True.quantization_config (BitsAndBytesConfig | None, optional) – Optional
transformers.BitsAndBytesConfigfor loading the base model in 4-/8-bit (QLoRA).lm_headis kept unquantized so the fused-linear-logprob path stays numerically exact.activation_offload (bool, optional) – When
True, run the training forward insidetorch.autograd.graph.save_on_cpuso tensors saved for backward live in pinned host RAM instead of GPU memory. Trades PCIe bandwidth for GPU memory (the win grows with sequence length); a no-op during rollout / reference forwards.moe_lora_recompute (bool | None, optional) – Recompute routed-expert LoRA activations in backward on frozen packed base weights.
None(default) recomputes only outside activation-checkpointed blocks.lora_target_scope (str | None, optional) – Optional PEFT LoRA path scope for multimodal models (e.g.
"language_model"). Passed toadapt_lora_config_for_model().
- clone(index: int | None = None, wrap: bool = True) Self¶
Create a clone of the algorithm.
QLoRA clones rebuild the base via
from_pretrainedand transfer only adapter (+ value head) weights. FSDP2 clones copy a rank-0 CPU full state dict onto a fresh CPU actor, then shard. The dense full model is never placed on GPU.- Parameters:
index (int | None, optional) – The index of the clone, defaults to None
wrap (bool, optional) – Unused. Clones always call
wrap_models(). Kept so tournament / multi-frequency can passwrap=False.
- Returns:
A clone of the algorithm
- Return type:
- configure_batch_size_per_process(batch_size: int, micro_batch_size_per_gpu: int | None, mini_batch_size: int | None, group_size: int = 1) None¶
Derive per-process batch sizes and gradient accumulation steps.
batch_sizeis the global collect size (prompt groups for GRPO-family). Each data-parallel replica holds(batch_size / dp_size) * group_sizesamples, wheredp_sizeis the process-group world folded by tensor-parallel degree (world_size / tp). Unsetmini_batch_sizeusesmicro_batch_size_per_gpuwhen the class default is"micro_batch"(RL rollout algorithms) and that is set, else the per-rank collect. Unsetmicro_batch_size_per_gpuuses the mini-batch.gradient_accumulation_stepsismini_batch_size / micro_batch_size_per_gpu.
- static copy_attributes(agent: IndividualT, clone: IndividualT, exclude: Iterable[str] = ()) IndividualT¶
Copy the non-evolvable attributes of the algorithm to a clone.
- Parameters:
agent (EvolvableAlgorithm) – The algorithm to copy attributes from.
clone (EvolvableAlgorithm) – The clone of the algorithm.
exclude (Iterable[str]) – Attribute names to leave on
clone/agent.
- Returns:
The clone of the algorithm.
- Return type:
- property current_lr_critic: float¶
Critic learning rate the next
learncall trains with;lris the peak withoutlr_critic.
- eval_policy_network_ids() set[int]¶
Return the id of every evaluation network in the agent’s policy group.
- evolvable_attributes(networks_only: bool = False) dict[str, Any]¶
Return the attributes related to the evolvable networks in the algorithm. Includes attributes that are either EvolvableModule or ModuleDict objects, as well as the optimizers associated with the networks.
- finalize_training_step(num_steps: int) None¶
Close the agent’s training block, storing any captured GraMa scores.
- Parameters:
num_steps (int) – Number of steps taken during the training step.
- Returns:
None.
- Return type:
None
- property fitness: list[float | ndarray[tuple[int, ...], dtype[_ScalarType_co]]]¶
Fitness history (scalars, or per-sub-agent rows for multi-agent).
- get_action(obs: ndarray[tuple[int, ...], dtype[_ScalarType_co]] | dict[str, ndarray[tuple[int, ...], dtype[_ScalarType_co]]] | tuple[ndarray[tuple[int, ...], dtype[_ScalarType_co]], ...] | Tensor | TensorDict | tuple[Tensor, ...] | dict[str, Tensor] | Number | list[RolloutPrompt] | RolloutPrompt | dict[str, ndarray[tuple[int, ...], dtype[_ScalarType_co]] | dict[str, ndarray[tuple[int, ...], dtype[_ScalarType_co]]] | tuple[ndarray[tuple[int, ...], dtype[_ScalarType_co]], ...] | Tensor | TensorDict | tuple[Tensor, ...] | dict[str, Tensor] | Number | list[RolloutPrompt] | RolloutPrompt], *args: Any, **kwargs: Any) NoReturn¶
Not implemented — DPO is an offline preference algorithm.
- Parameters:
obs (ObservationType | MultiAgentObservationType) – The observation of the agent
args (Any) – Additional arguments (unused; for base contract compatibility)
kwargs (Any) – Additional keyword arguments (e.g. training; unused)
- Raises:
NotImplementedError – Always.
- static get_action_dim(action_space: Space | list[Space] | dict[str, Space]) int | dict[str, int] | tuple[int | dict[str, int], ...]¶
Return the dimension of the action space as it pertains to the underlying networks (i.e. the output size of the networks).
- get_eval_modules(cloning: bool = True) tuple[dict[str, EvolvableModule], dict[str, EvolvableModule]]¶
Get the offsprings of all of the evaluation modules in the individual.
- Parameters:
cloning (bool, optional) – Whether to clone each evaluation module before returning it, defaults to True.
- Returns:
Tuple of offspring policy and the rest of the evaluation modules
- Return type:
tuple[dict[str, EvolvableModule], dict[str, EvolvableModule]]
- get_lr_names() list[str | tuple[str, str]]¶
Return the learning-rate attribute name(s) of each optimizer.
- get_policy() EvolvableModuleProtocol¶
Return the policy network of the algorithm.
- static get_state_dim(observation_space: Space | list[Space] | dict[str, Space]) tuple[int, ...] | dict[str, tuple[int, ...]] | tuple[tuple[int, ...] | dict[str, tuple[int, ...]], ...]¶
Return the dimension of the state space as it pertains to the underlying networks (i.e. the input size of the networks).
- property hp_config: HyperparameterConfig¶
Return the hyperparameter configuration for Evo-HPO mutations.
- init_training_step(capture_grama: bool = False) None¶
Open the agent’s training block: metrics tracking, and GraMa capture.
Hooks are registered afresh each cycle, so they follow the agent through architecture mutations, checkpoint reloads and accelerator re-wrapping. Opening a block implicitly closes one that an earlier call left open.
- Parameters:
capture_grama (bool) – Whether to register GraMa capture hooks for this training step. Defaults to False since the LLM finetuners never run ReGraMa.
- Returns:
None.
- Return type:
None
- static inspect_attributes(agent: EvolvableAlgorithmProtocol | AgentWrapperProtocol[Any], input_args_only: bool = False, exclude: Iterable[str] = ()) dict[str, Any]¶
Inspect and retrieve the attributes of the current object, excluding attributes related to the underlying evolvable networks (i.e. EvolvableModule, torch.optim.Optimizer) and with an option to include only the attributes that are input arguments to the constructor.
- Parameters:
input_args_only (bool) – If True, only include attributes that are input arguments to the constructor. Defaults to False.
exclude (Iterable[str], optional) – Extra attribute names to drop from the result, on top of the standard exclusions below. For a caller-specific reason to leave an attribute out of its own view.
- Returns:
A dictionary of attribute names and their values.
- Return type:
- learn(experiences: PreferencePrompts, training: bool = True) dict[str, float]¶
Update agent network parameters to learn from preference data.
- Parameters:
experiences (PreferencePrompts) – Batched chosen/rejected input ids and attention masks with prompt lengths, as produced by a preference
DatasetEnv.training (bool) – Whether the agent is training or not
- Returns:
Dict with the cross-rank means of
loss,chosen_rewardandrejected_reward, andlearn_phase_<phase>_swall seconds.- Return type:
- classmethod load(path: str, device: str | device = 'cpu', accelerator: Accelerator | None = None) Self¶
Load an algorithm from a checkpoint.
- Parameters:
path (string) – Location to load checkpoint from.
device (str, optional) – Device to load the algorithm on, defaults to ‘cpu’
accelerator (Accelerator | None, optional) – Accelerator object for distributed computing, defaults to None
- Returns:
An instance of the algorithm
- Return type:
- load_checkpoint(path: str, load_optimizer: bool = False, overwrite_reference_adapter: bool | None = None, overwrite_critic_adapter: bool = False, restore_config: bool = True, restore_hyperparameters: bool = True) None¶
Load adapter weights and algorithm state from a checkpoint directory.
Adapter roles restored on load:
actor— the trained policy. Always loaded.reference— loaded from the checkpoint’sreference/adapter when it has one; otherwise the checkpoint’sactoris copied ontoreferenceso SFT -> DPO -> GRPO chains work out of the box.critic— loaded from the checkpoint’scritic/adapter when it has one, otherwise left at its fresh LoRA init. Setoverwrite_critic_adapterto seed it from the actor.
The checkpoint’s LoRA config must match the live algorithm’s config; a mismatch raises
ValueError(re-create the agent with the checkpoint’s LoRA config to load it).The same flow applies to plain, DDP and FSDP2 runs:
- lora_only=T -> PEFT adapter dirs are loaded into the live
adapters.
- lora_only=F -> the full actor state_dict is restored from
attributes.pt.
When
load_optimizer=Truethe optimizer state is restored fromattributes.pt; if the checkpoint contains no optimizer state (saved withsave_optimizer=False), aUserWarningis emitted and a freshly-initialised optimizer is used. The LR schedule always resumes at the checkpoint’s learn step, with this instance’slr/lr_critic(after any hyperparameter restore) as its peaks.- Parameters:
path (str) – Directory containing a checkpoint written by
save_checkpoint().load_optimizer (bool) – If
Truealso load the optimizer state so training can resume.overwrite_reference_adapter (bool | None) – Copy the checkpoint’s
actorontoreferenceeven when it has areference/adapter.Nonecopies only when the checkpoint has no reference adapter.overwrite_critic_adapter (bool) – Seed
criticfrom the checkpoint’sactor.restore_config (bool) – If
False, keep this instance’s algorithm settings and restore only training state: step count, scores, fitness,reference_update_tracker,rngand, withrestore_hyperparameters, the registry’s mutable hyperparameters.restore_hyperparameters (bool) – If
False, keep this instance’s values for the registry’s mutable hyperparameters (e.g.lr). Only a run that mutates them needs the checkpoint’s values.
- load_weights(path: str, overwrite_reference_adapter: bool | None = None, overwrite_critic_adapter: bool = False) None¶
Load only the LoRA adapters (and value head) from a checkpoint directory.
- Parameters:
path (str) – Directory containing a checkpoint written by
save_checkpoint().overwrite_reference_adapter (bool | None) – See
load_checkpoint().overwrite_critic_adapter (bool) – See
load_checkpoint().
- classmethod population(size: int, device: str | device = 'cpu', resume_from_checkpoint: str | None = None, **kwargs: Any) list[Self]¶
Create a population of LLM algorithms.
Builds agent 0 fully (loading the model from disk), then clones for agents 1..N. Under FSDP2 / QLoRA this uses adapter-only
clone(); otherwise the actor is copied viaclone_llm().- Parameters:
- Returns:
A list of LLM algorithms.
- Return type:
list[LLMAlgorithm]
- preprocess_observation(observation: Tensor | TensorDict | tuple[Tensor, ...] | dict[str, Tensor]) Tensor | TensorDict | tuple[Tensor, ...] | dict[str, Tensor]¶
Preprocess observations (dummy) for forward pass through neural network.
- recompile() None¶
Recompile evolvable modules with
torch.compile.Iterates over
evolvable_attributesand compiles each one. Skipped for distributed runs, matching_initialize_actors().
- register_mutation_hook(hook: LambdaType | MethodType) None¶
Register a hook to be executed after a mutation is performed on the algorithm.
- Parameters:
hook (MutationHook) – The hook to be executed after mutation.
- register_network_group(group: NetworkGroup) None¶
Set the evaluation network for the algorithm.
- Parameters:
name (str) – The name of the evaluation network.
- reinit_optimizers(optimizer: OptimizerConfig | None = None) None¶
Reinitialize the optimizers of an algorithm. If no optimizer is passed, all optimizers are reinitialized.
- Parameters:
optimizer (OptimizerConfig | None, optional) – The optimizer to reinitialize, defaults to None, in which case all optimizers are reinitialized.
- save_checkpoint(path: str, lora_only: bool = True, save_optimizer: bool = True) None¶
Save adapter weights and algorithm state to a directory.
AgileRL never persists base-model weights when
lora_only=Truefor LLM algorithms: a checkpoint is a directory containing<adapter>/adapter_model.safetensors+adapter_config.json— one subdirectory per adapter inselected_adapters(alwaysactor, plusreference/criticwhen those adapters are configured). Written only whenlora_only=True.attributes.pt— algorithm hyperparameters, plus (optionally) the actor state dict and/or optimizer state dict. Always present.
The same format is written for plain, DDP and FSDP2 runs:
- lora_only=T, save_optimizer=T -> PEFT adapter dirs on disk +
optimizer state in
attributes.pt
lora_only=T, save_optimizer=F -> PEFT adapter dirs only lora_only=F, save_optimizer=T -> full actor state_dict +
optimizer state in
attributes.ptlora_only=F, save_optimizer=F -> full actor state_dict in
attributes.ptFSDP2-sharded parameters and optimizer state are gathered to full tensors before saving, so checkpoints are rank-count independent.
- Parameters:
path (str) – Directory to write the checkpoint into.
lora_only (bool) – If
True(default) only adapter weights are written to disk viasave_pretrained; the base model is shared across checkpoints and not serialised. IfFalse, the full actor state dict is persisted intoattributes.pt.save_optimizer (bool) – If
True(default) also persist the optimizer and LR scheduler state inattributes.ptso training can resume.
- property scores: list[float | list[float]]¶
Per-episode scores (per-group score rows for multi-agent metrics).
- select_adapter(adapter_name: str) Generator[None, None, None]¶
Temporarily switch adapter; restores the actor adapter on exit.
- Parameters:
adapter_name (str) – Name of the adapter to activate (“actor”, “critic”, “reference”).
- set_reference_policy(reference_update_tracker: int) None¶
Update the reference policy when the tracker advances past the stored value.
Base weights are immutable in AgileRL’s LoRA-only training: with
use_separate_reference_adapter=Truethe actor adapter is copied onto thereferenceadapter; without one the implicit reference (the base model with adapters disabled) cannot move, so the update request is acknowledged with a one-time warning and the KL anchor stays the initial policy.- Parameters:
reference_update_tracker (int) – The reference policy update tracker
- set_training_mode(training: bool) None¶
Set the training mode of the algorithm.
- Parameters:
training (bool) – If True, set the algorithm to training mode.
- snapshot_checkpoint(lora_only: bool = True, save_optimizer: bool = True) LLMCheckpointSnapshot¶
Copy the checkpoint of
save_checkpoint()into host memory.All ranks must call this together (FSDP2 gathers are collective). The snapshot holds the weights and optimizer state of this step: later training does not change it, so
LLMCheckpointSnapshot.write()may run on another thread while training continues.- Parameters:
lora_only (bool) – See
save_checkpoint().save_optimizer (bool) – See
save_checkpoint().
- Returns:
Host copy of the checkpoint; empty off the main process.
- Return type:
LLMCheckpointSnapshot
- test(env: DatasetEnv, loop: int = 1, *args: Any, **kwargs: Any) npt.NDArray¶
Return the fitness (test) score of the agent.
- Parameters:
env (DatasetEnv) – The environment to be tested in (
objective="preference")loop (int, optional) – Number of testing loops/episodes to complete. The returned score is the mean. Defaults to 1
- Returns:
Mean test score (numpy array)
- Return type:
npt.NDArray
- to_device(*experiences: Tensor | TensorDict | tuple[Tensor, ...] | dict[str, Tensor]) tuple[Tensor | TensorDict | tuple[Tensor, ...] | dict[str, Tensor], ...]¶
Move experiences to the device.
- unrolled_eval_networks() list[tuple[str | None, Module]]¶
Return the agent’s evaluation networks as (network_id, network) pairs.
- update_existing_adapter(checkpoint_dir: str, adapter_name: str) None¶
Overwrite weights of an existing adapter in-place without creating new parameters.
- Parameters:
checkpoint_dir (str) – Checkpoint directory
adapter_name (str.) – Adapter name
- Returns:
None
- Return type:
None
- static update_lr(optimizer: OptimizerWrapper, lr: float | tuple[float, float | None], scheduler_config: CosineLRScheduleConfig | None = None, schedule_step: int = 0) LambdaLR | None¶
Set the peak learning rate of each param group and rebuild the schedule.
- Parameters:
optimizer (OptimizerWrapper) – LLM optimizer.
lr (float | tuple[float, float | None]) – Learning rate value, or actor/critic pair; a
Nonecritic uses the actor rate.scheduler_config (CosineLRScheduleConfig | None) – Scheduler configuration;
Noneholdslrconstant.schedule_step (int) – Learn steps the schedule has already taken.
- Returns:
Scheduler at
schedule_stepwhenscheduler_configis set.- Return type:
LambdaLR | None
- use_adapter(adapter_name: str) None¶
Switch the active PEFT adapter, handling all side-effects.
For “reference”: switches adapter and freezes reference params (never trained). For all others: switches adapter and restores requires_grad=True on all training adapter LoRA params so distributed gradient hooks keep firing.
- Parameters:
adapter_name (str) – Name of the adapter to activate (“actor”, “critic”, “reference”).
- wrap_models() None¶
Prepare the actor for training.
Places the actor (dense
.to(device)or FSDP2 shard) then builds the optimizer and LR scheduler on those parameters. FSDP2 wraps each transformer block with activation checkpointing beforefully_shardwhengradient_checkpointingis on. Data-parallelprepare_actoruses HuggingFacegradient_checkpointing_enable.