2097 lines
107 KiB
Python
2097 lines
107 KiB
Python
"""
|
||
2025.6.1
|
||
2025.6.2
|
||
4.52.4
|
||
0.18.2
|
||
__UNSLOTH_VERSIONING__
|
||
"""
|
||
from torch import Tensor
|
||
import torch
|
||
import torch.nn as nn
|
||
from torch.nn import functional as F
|
||
from trl.trainer.dpo_trainer import (Any, AutoModelForCausalLM, BaseImageProcessor, Callable, DPOConfig, DPOTrainer, DataCollator, DataCollatorForPreference, DataLoader, Dataset, EvalLoopOutput, F, FDivergenceConstants, FDivergenceType, FeatureExtractionMixin, IterableDataset, Literal, MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES, Optional, PartialState, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, RunningMoments, SyncRefModelCallback, Trainer, TrainerCallback, Union, amp, cap_exp, contextmanager, create_reference_model, dataclass, defaultdict, disable_dropout_in_model, empty_cache, flush_left, flush_right, generate_model_card, get_comet_experiment_url, inspect, is_comet_available, is_peft_available, is_torch_xpu_available, is_wandb_available, log_table_to_comet_experiment, maybe_apply_chat_template, maybe_extract_prompt, nn, nullcontext, os, pad, pad_to_length, pd, peft_module_casting_to_bf16, prepare_deepspeed, prepare_fsdp, prepare_model_for_kbit_training, random, textwrap, torch, tqdm, transformers, version, warnings)
|
||
|
||
|
||
import os
|
||
from typing import *
|
||
from dataclasses import dataclass, field
|
||
from packaging.version import Version
|
||
import torch
|
||
import numpy as np
|
||
from contextlib import nullcontext
|
||
from torch.nn import functional as F
|
||
from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling as TransformersDataCollatorForLanguageModeling
|
||
|
||
torch_compile_options = {
|
||
"epilogue_fusion" : True,
|
||
"max_autotune" : False,
|
||
"shape_padding" : True,
|
||
"trace.enabled" : False,
|
||
"triton.cudagraphs" : False,
|
||
}
|
||
|
||
@torch.compile(dynamic = True, fullgraph = True, options = torch_compile_options,)
|
||
def selective_log_softmax(logits, index):
|
||
logits = logits.to(torch.float32)
|
||
selected_logits = torch.gather(logits, dim = -1, index = index.unsqueeze(-1)).squeeze(-1)
|
||
# loop to reduce peak mem consumption
|
||
# logsumexp_values = torch.stack([torch.logsumexp(lg, dim=-1) for lg in logits])
|
||
logsumexp_values = torch.logsumexp(logits, dim = -1)
|
||
per_token_logps = selected_logits - logsumexp_values # log_softmax(x_i) = x_i - logsumexp(x)
|
||
return per_token_logps
|
||
@dataclass
|
||
class UnslothDPOConfig(DPOConfig):
|
||
"""
|
||
|
||
Configuration class for the [`DPOTrainer`].
|
||
|
||
Using [`~transformers.HfArgumentParser`] we can turn this class into
|
||
[argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the
|
||
command line.
|
||
|
||
Parameters:
|
||
> Parameters that control the model and reference model
|
||
|
||
model_init_kwargs (`dict[str, Any]` or `None`, *optional*, defaults to `None`):
|
||
Keyword arguments for `AutoModelForCausalLM.from_pretrained`, used when the `model` argument of the
|
||
[`DPOTrainer`] is provided as a string.
|
||
ref_model_init_kwargs (`dict[str, Any]` or `None`, *optional*, defaults to `None`):
|
||
Keyword arguments for `AutoModelForCausalLM.from_pretrained`, used when the `ref_model` argument of the
|
||
[`DPOTrainer`] is provided as a string.
|
||
model_adapter_name (`str` or `None`, *optional*, defaults to `None`):
|
||
Name of the train target PEFT adapter, when using LoRA with multiple adapters.
|
||
ref_adapter_name (`str` or `None`, *optional*, defaults to `None`):
|
||
Name of the reference PEFT adapter, when using LoRA with multiple adapters.
|
||
force_use_ref_model (`bool`, *optional*, defaults to `False`):
|
||
If you provide a PEFT model as the active model and wish to use a different model for the `ref_model`, set
|
||
this flag to `True`.
|
||
disable_dropout (`bool`, *optional*, defaults to `True`):
|
||
Whether to disable dropout in the model and reference model.
|
||
use_logits_to_keep (`bool`, *optional*, defaults to `False`):
|
||
If `True`, only a specified number of logits are computed in the forward pass. This can be useful for
|
||
saving memory and speeding up training by not computing the logits for all tokens, especially in
|
||
scenarios when working with very long prompts where labels are ignored (-100).
|
||
|
||
> Parameters that control the data preprocessing
|
||
|
||
dataset_num_proc (`int` or `None`, *optional*, defaults to `None`):
|
||
Number of processes to use for processing the dataset.
|
||
padding_value (`int` or `None`, *optional*, defaults to `None`):
|
||
Padding value to use. If `None`, the padding value of the tokenizer is used.
|
||
label_pad_token_id (`int`, *optional*, defaults to `-100`):
|
||
Padding value to use for labels.
|
||
max_prompt_length (`int` or `None`, *optional*, defaults to `512`):
|
||
Maximum length of the prompt.
|
||
max_completion_length (`int` or `None`, *optional*, defaults to `None`):
|
||
Maximum length of the completion.
|
||
max_length (`int` or `None`, *optional*, defaults to `1024`):
|
||
Maximum length of the full sequence (prompt + completion).
|
||
truncation_mode (`str`, *optional*, defaults to `"keep_end"`):
|
||
Truncation mode to use when the sequence exceeds `max_length`. Possible values are `"keep_end"` and
|
||
`"keep_start"`.
|
||
padding_free (`bool`, *optional*, defaults to `False`):
|
||
Whether to perform forward passes without padding by flattening all sequences in the batch into a single
|
||
continuous sequence. This reduces memory usage by eliminating padding overhead. Currently, this is only
|
||
supported with the `flash_attention_2` attention implementation, which can efficiently handle the flattened
|
||
batch structure.
|
||
precompute_ref_log_probs (`bool`, *optional*, defaults to `False`):
|
||
Whether to precompute the log probabilities from the reference model. Setting this to `True` allows
|
||
training without needing the reference model during training, which can help reduce GPU memory usage. If
|
||
set to `False` (default), the reference model will be used during training to compute log probabilities
|
||
on-the-fly.
|
||
precompute_ref_batch_size (`int` or `None`, *optional*, defaults to `None`):
|
||
Batch size to use when precomputing reference model log probabilities. This can be set higher than the
|
||
training batch size to speed up preprocessing. If `None`, defaults to `per_device_train_batch_size` for
|
||
training and `per_device_eval_batch_size` for evaluation.
|
||
tools (`Optional[list[Union[dict, Callable]]]`, *optional*, defaults to `None`):
|
||
List of tools (callable functions) that will be accessible to the model.
|
||
If the template does not support function calling, this argument will have no effect.
|
||
|
||
> Parameters that control the training
|
||
|
||
learning_rate (`float`, *optional*, defaults to `1e-6`):
|
||
Initial learning rate for [`AdamW`] optimizer. The default value replaces that of
|
||
[`~transformers.TrainingArguments`].
|
||
loss_type (`str`, *optional*, defaults to `"sigmoid"`):
|
||
Type of loss to use. Possible values are:
|
||
|
||
- `"sigmoid"`: sigmoid loss from the original [DPO](https://huggingface.co/papers/2305.18290) paper.
|
||
- `"hinge"`: hinge loss on the normalized likelihood from the [SLiC](https://huggingface.co/papers/2305.10425) paper.
|
||
- `"ipo"`: IPO loss from the [IPO](https://huggingface.co/papers/2310.12036) paper.
|
||
- `"exo_pair"`: pairwise EXO loss from the [EXO](https://huggingface.co/papers/2402.00856) paper.
|
||
- `"nca_pair"`: pairwise NCA loss from the [NCA](https://huggingface.co/papers/2402.05369) paper.
|
||
- `"robust"`: unbiased estimate of the DPO loss that is robust to preference noise from the [Robust DPO](https://huggingface.co/papers/2403.00409) paper.
|
||
- `"bco_pair"`: pairwise BCO loss from the [BCO](https://huggingface.co/papers/2404.04656) paper.
|
||
- `"sppo_hard"`: SPPO loss with hard label from the [SPPO](https://huggingface.co/papers/2405.00675) paper.
|
||
- `"aot"`: AOT loss for paired datasets from the [AOT](https://huggingface.co/papers/2406.05882) paper.
|
||
- `"aot_pair"`: AOT loss for unpaired datasets from the [AOT](https://huggingface.co/papers/2406.05882) paper.
|
||
- `"discopop"`: DiscoPOP (a.k.a Log-Ratio Modulated Loss, LRML) loss from the [DiscoPOP](https://huggingface.co/papers/2406.08414) paper.
|
||
- `"apo_zero"`: APO-zero loss from the [APO](https://huggingface.co/papers/2408.06266) paper.
|
||
- `"apo_down"`: APO-down loss from the [APO](https://huggingface.co/papers/2408.06266) paper.
|
||
|
||
beta (`float`, *optional*, defaults to `0.1`):
|
||
Parameter controlling the deviation from the reference model. Higher β means less deviation from the
|
||
reference model. For the IPO loss (`loss_type="ipo"`), β is the regularization parameter denoted by τ in
|
||
the [paper](https://huggingface.co/papers/2310.12036).
|
||
f_divergence_type (`str`, *optional*, defaults to `FDivergenceType.REVERSE_KL`):
|
||
Type of f-divergence regularization function to compute divergence between policy and reference model.
|
||
f_alpha_divergence_coef (`float`, *optional*, defaults to `1.0`):
|
||
α coefficient in the α-divergence u^-α regularization function for DPO loss.
|
||
reference_free (`bool`, *optional*, defaults to `False`):
|
||
Whether to ignore the provided reference model and implicitly use a reference model that assigns equal
|
||
probability to all responses.
|
||
label_smoothing (`float`, *optional*, defaults to `0.0`):
|
||
Robust DPO label smoothing parameter from the [cDPO report](https://ericmitchell.ai/cdpo.pdf) and
|
||
[Robust DPO](https://huggingface.co/papers/2403.00409) paper that should be between `0.0` and `0.5`.
|
||
use_weighting (`bool`, *optional*, defaults to `False`):
|
||
Whether to weight the loss as done in the [WPO paper](https://huggingface.co/papers/2406.11827).
|
||
rpo_alpha (`float`, *optional*, defaults to `None`):
|
||
α parameter from the [RPO paper](https://huggingface.co/papers/2404.19733) (v3), which controls the
|
||
weighting of the NLL term in the loss. If `None`, no weighting is applied and the loss is the same as the
|
||
DPO loss. The paper recommends `rpo_alpha=1.0`.
|
||
ld_alpha (`float` or `None`, *optional*, defaults to `None`):
|
||
α parameter from the [LD-DPO paper](https://huggingface.co/papers/2409.06411), which controls the weighting
|
||
of the verbose token log-probabilities in responses. If `None`, no weighting is applied to the verbose
|
||
part, and the loss is equivalent to the standard DPO loss. The paper recommends setting `ld_alpha` between
|
||
`0.0` and `1.0`.
|
||
discopop_tau (`float`, *optional*, defaults to `0.05`):
|
||
τ/temperature parameter from the [DiscoPOP](https://huggingface.co/papers/2406.08414) paper, which controls
|
||
the shape of log ratio modulated loss. The paper recommends the default value `discopop_tau=0.05`.
|
||
sync_ref_model (`bool`, *optional*, defaults to `False`):
|
||
Whether to synchronize the reference model with the active model every `ref_model_sync_steps` steps, using
|
||
the `ref_model_mixup_alpha` parameter. This synchronization originites from the
|
||
[TR-DPO](https://huggingface.co/papers/2404.09656) paper.
|
||
ref_model_mixup_alpha (`float`, *optional*, defaults to `0.6`):
|
||
α parameter from the [TR-DPO](https://huggingface.co/papers/2404.09656) paper, which controls the mix
|
||
between the current policy and the previous reference policy during updates. The reference policy is
|
||
updated according to the equation: `π_ref = α * π_θ + (1 - α) * π_ref_prev`. To use this parameter, you
|
||
must set `sync_ref_model=True`.
|
||
ref_model_sync_steps (`int`, *optional*, defaults to `512`):
|
||
τ parameter from the [TR-DPO](https://huggingface.co/papers/2404.09656) paper, which determines how
|
||
frequently the current policy is synchronized with the reference policy. To use this parameter, you must
|
||
set `sync_ref_model=True`.
|
||
|
||
> Parameters that control the logging
|
||
|
||
generate_during_eval (`bool`, *optional*, defaults to `False`):
|
||
Whether to generate and log completions from both the model and the reference model to W&B or Comet during
|
||
evaluation.
|
||
|
||
"""
|
||
vllm_sampling_params: Optional[Any] = field(
|
||
default = None,
|
||
metadata = {'help': 'vLLM SamplingParams'},
|
||
)
|
||
unsloth_num_chunks : Optional[int] = field(
|
||
default = -1,
|
||
metadata = {'help': 'Chunk size to reduce memory usage. -1 is most efficient.'},
|
||
)
|
||
def __init__(
|
||
self,
|
||
output_dir = None,
|
||
overwrite_output_dir = None,
|
||
do_train = False,
|
||
do_eval = False,
|
||
do_predict = False,
|
||
eval_strategy = 'no',
|
||
prediction_loss_only = False,
|
||
per_device_train_batch_size = 4,
|
||
per_device_eval_batch_size = 4,
|
||
per_gpu_train_batch_size = None,
|
||
per_gpu_eval_batch_size = None,
|
||
gradient_accumulation_steps = 2,
|
||
eval_accumulation_steps = 2,
|
||
eval_delay = 0,
|
||
torch_empty_cache_steps = 250,
|
||
learning_rate = 5e-05,
|
||
weight_decay = 0.01,
|
||
adam_beta1 = 0.9,
|
||
adam_beta2 = 0.999,
|
||
adam_epsilon = 1e-08,
|
||
max_grad_norm = 1.0,
|
||
num_train_epochs = 3.0,
|
||
max_steps = -1,
|
||
lr_scheduler_type = 'linear',
|
||
warmup_ratio = 0.1,
|
||
warmup_steps = 0,
|
||
log_level = 'passive',
|
||
log_level_replica = 'warning',
|
||
log_on_each_node = True,
|
||
logging_dir = None,
|
||
logging_strategy = 'steps',
|
||
logging_first_step = False,
|
||
logging_steps = 1,
|
||
logging_nan_inf_filter = False,
|
||
save_strategy = 'steps',
|
||
save_steps = 500,
|
||
save_total_limit = None,
|
||
save_safetensors = True,
|
||
save_on_each_node = False,
|
||
save_only_model = False,
|
||
restore_callback_states_from_checkpoint = False,
|
||
no_cuda = False,
|
||
use_cpu = False,
|
||
use_mps_device = False,
|
||
seed = 3407,
|
||
data_seed = 3407,
|
||
jit_mode_eval = False,
|
||
use_ipex = False,
|
||
bf16 = False,
|
||
fp16 = False,
|
||
fp16_opt_level = 'O1',
|
||
half_precision_backend = 'auto',
|
||
bf16_full_eval = False,
|
||
fp16_full_eval = False,
|
||
tf32 = None,
|
||
local_rank = -1,
|
||
ddp_backend = None,
|
||
tpu_num_cores = None,
|
||
tpu_metrics_debug = False,
|
||
debug = '',
|
||
dataloader_drop_last = False,
|
||
eval_steps = None,
|
||
dataloader_num_workers = 0,
|
||
dataloader_prefetch_factor = None,
|
||
past_index = -1,
|
||
run_name = None,
|
||
disable_tqdm = None,
|
||
remove_unused_columns = True,
|
||
label_names = None,
|
||
load_best_model_at_end = False,
|
||
metric_for_best_model = None,
|
||
greater_is_better = None,
|
||
ignore_data_skip = False,
|
||
fsdp = '',
|
||
fsdp_min_num_params = 0,
|
||
fsdp_config = None,
|
||
fsdp_transformer_layer_cls_to_wrap = None,
|
||
accelerator_config = None,
|
||
deepspeed = None,
|
||
label_smoothing_factor = 0.0,
|
||
optim = 'adamw_8bit',
|
||
optim_args = None,
|
||
adafactor = False,
|
||
group_by_length = False,
|
||
length_column_name = 'length',
|
||
report_to = None,
|
||
ddp_find_unused_parameters = None,
|
||
ddp_bucket_cap_mb = None,
|
||
ddp_broadcast_buffers = None,
|
||
dataloader_pin_memory = True,
|
||
dataloader_persistent_workers = False,
|
||
skip_memory_metrics = True,
|
||
use_legacy_prediction_loop = False,
|
||
push_to_hub = False,
|
||
resume_from_checkpoint = None,
|
||
hub_model_id = None,
|
||
hub_strategy = 'every_save',
|
||
hub_token = None,
|
||
hub_private_repo = None,
|
||
hub_always_push = False,
|
||
gradient_checkpointing = False,
|
||
gradient_checkpointing_kwargs = None,
|
||
include_inputs_for_metrics = False,
|
||
eval_do_concat_batches = True,
|
||
fp16_backend = 'auto',
|
||
push_to_hub_model_id = None,
|
||
push_to_hub_organization = None,
|
||
push_to_hub_token = None,
|
||
mp_parameters = '',
|
||
auto_find_batch_size = False,
|
||
full_determinism = False,
|
||
torchdynamo = None,
|
||
ray_scope = 'last',
|
||
ddp_timeout = 1800,
|
||
torch_compile = False,
|
||
torch_compile_backend = None,
|
||
torch_compile_mode = None,
|
||
include_tokens_per_second = False,
|
||
include_num_input_tokens_seen = False,
|
||
neftune_noise_alpha = None,
|
||
optim_target_modules = None,
|
||
batch_eval_metrics = False,
|
||
eval_on_start = False,
|
||
use_liger_kernel = False,
|
||
eval_use_gather_object = False,
|
||
average_tokens_across_devices = False,
|
||
model_init_kwargs = None,
|
||
ref_model_init_kwargs = None,
|
||
model_adapter_name = None,
|
||
ref_adapter_name = None,
|
||
force_use_ref_model = False,
|
||
disable_dropout = True,
|
||
use_logits_to_keep = False,
|
||
dataset_num_proc = None,
|
||
padding_value = None,
|
||
label_pad_token_id = -100,
|
||
max_prompt_length = 512,
|
||
max_completion_length = None,
|
||
max_length = 1024,
|
||
truncation_mode = 'keep_end',
|
||
padding_free = False,
|
||
precompute_ref_log_probs = False,
|
||
precompute_ref_batch_size = None,
|
||
tools = None,
|
||
loss_type = 'sigmoid',
|
||
beta = 0.1,
|
||
f_alpha_divergence_coef = 1.0,
|
||
reference_free = False,
|
||
label_smoothing = 0.0,
|
||
use_weighting = False,
|
||
rpo_alpha = None,
|
||
ld_alpha = None,
|
||
discopop_tau = 0.05,
|
||
sync_ref_model = False,
|
||
ref_model_mixup_alpha = 0.6,
|
||
ref_model_sync_steps = 512,
|
||
generate_during_eval = False,
|
||
vllm_sampling_params = None,
|
||
unsloth_num_chunks = -1,
|
||
**kwargs,
|
||
):
|
||
if learning_rate < 1e-7: raise FloatingPointError(f'Unsloth: Your learning rate of `{learning_rate}` is too small and less than 1e-7! Consider increasing it, otherwise gradient updates will be close to 0!')
|
||
if learning_rate > 1: raise OverflowError(f'Unsloth: Your learning rate of `{learning_rate}` is way too larger > 1! Consider decreasing it to 1e-1, otherwise gradient updates will explode!')
|
||
if output_dir is None and save_strategy == 'steps' and save_steps == 500:
|
||
output_dir = 'unsloth_training_checkpoints'
|
||
save_strategy = 'no'
|
||
if dataset_num_proc is None:
|
||
from multiprocessing import cpu_count
|
||
dataset_num_proc = cpu_count()
|
||
|
||
super().__init__(
|
||
output_dir = output_dir,
|
||
overwrite_output_dir = overwrite_output_dir,
|
||
do_train = do_train,
|
||
do_eval = do_eval,
|
||
do_predict = do_predict,
|
||
eval_strategy = eval_strategy,
|
||
prediction_loss_only = prediction_loss_only,
|
||
per_device_train_batch_size = per_device_train_batch_size,
|
||
per_device_eval_batch_size = per_device_eval_batch_size,
|
||
per_gpu_train_batch_size = per_gpu_train_batch_size,
|
||
per_gpu_eval_batch_size = per_gpu_eval_batch_size,
|
||
gradient_accumulation_steps = gradient_accumulation_steps,
|
||
eval_accumulation_steps = eval_accumulation_steps,
|
||
eval_delay = eval_delay,
|
||
torch_empty_cache_steps = torch_empty_cache_steps,
|
||
learning_rate = learning_rate,
|
||
weight_decay = weight_decay,
|
||
adam_beta1 = adam_beta1,
|
||
adam_beta2 = adam_beta2,
|
||
adam_epsilon = adam_epsilon,
|
||
max_grad_norm = max_grad_norm,
|
||
num_train_epochs = num_train_epochs,
|
||
max_steps = max_steps,
|
||
lr_scheduler_type = lr_scheduler_type,
|
||
warmup_ratio = warmup_ratio,
|
||
warmup_steps = warmup_steps,
|
||
log_level = log_level,
|
||
log_level_replica = log_level_replica,
|
||
log_on_each_node = log_on_each_node,
|
||
logging_dir = logging_dir,
|
||
logging_strategy = logging_strategy,
|
||
logging_first_step = logging_first_step,
|
||
logging_steps = logging_steps,
|
||
logging_nan_inf_filter = logging_nan_inf_filter,
|
||
save_strategy = save_strategy,
|
||
save_steps = save_steps,
|
||
save_total_limit = save_total_limit,
|
||
save_safetensors = save_safetensors,
|
||
save_on_each_node = save_on_each_node,
|
||
save_only_model = save_only_model,
|
||
restore_callback_states_from_checkpoint = restore_callback_states_from_checkpoint,
|
||
no_cuda = no_cuda,
|
||
use_cpu = use_cpu,
|
||
use_mps_device = use_mps_device,
|
||
seed = seed,
|
||
data_seed = data_seed,
|
||
jit_mode_eval = jit_mode_eval,
|
||
use_ipex = use_ipex,
|
||
bf16 = bf16,
|
||
fp16 = fp16,
|
||
fp16_opt_level = fp16_opt_level,
|
||
half_precision_backend = half_precision_backend,
|
||
bf16_full_eval = bf16_full_eval,
|
||
fp16_full_eval = fp16_full_eval,
|
||
tf32 = tf32,
|
||
local_rank = local_rank,
|
||
ddp_backend = ddp_backend,
|
||
tpu_num_cores = tpu_num_cores,
|
||
tpu_metrics_debug = tpu_metrics_debug,
|
||
debug = debug,
|
||
dataloader_drop_last = dataloader_drop_last,
|
||
eval_steps = eval_steps,
|
||
dataloader_num_workers = dataloader_num_workers,
|
||
dataloader_prefetch_factor = dataloader_prefetch_factor,
|
||
past_index = past_index,
|
||
run_name = run_name,
|
||
disable_tqdm = disable_tqdm,
|
||
remove_unused_columns = remove_unused_columns,
|
||
label_names = label_names,
|
||
load_best_model_at_end = load_best_model_at_end,
|
||
metric_for_best_model = metric_for_best_model,
|
||
greater_is_better = greater_is_better,
|
||
ignore_data_skip = ignore_data_skip,
|
||
fsdp = fsdp,
|
||
fsdp_min_num_params = fsdp_min_num_params,
|
||
fsdp_config = fsdp_config,
|
||
fsdp_transformer_layer_cls_to_wrap = fsdp_transformer_layer_cls_to_wrap,
|
||
accelerator_config = accelerator_config,
|
||
deepspeed = deepspeed,
|
||
label_smoothing_factor = label_smoothing_factor,
|
||
optim = optim,
|
||
optim_args = optim_args,
|
||
adafactor = adafactor,
|
||
group_by_length = group_by_length,
|
||
length_column_name = length_column_name,
|
||
report_to = report_to,
|
||
ddp_find_unused_parameters = ddp_find_unused_parameters,
|
||
ddp_bucket_cap_mb = ddp_bucket_cap_mb,
|
||
ddp_broadcast_buffers = ddp_broadcast_buffers,
|
||
dataloader_pin_memory = dataloader_pin_memory,
|
||
dataloader_persistent_workers = dataloader_persistent_workers,
|
||
skip_memory_metrics = skip_memory_metrics,
|
||
use_legacy_prediction_loop = use_legacy_prediction_loop,
|
||
push_to_hub = push_to_hub,
|
||
resume_from_checkpoint = resume_from_checkpoint,
|
||
hub_model_id = hub_model_id,
|
||
hub_strategy = hub_strategy,
|
||
hub_token = hub_token,
|
||
hub_private_repo = hub_private_repo,
|
||
hub_always_push = hub_always_push,
|
||
gradient_checkpointing = gradient_checkpointing,
|
||
gradient_checkpointing_kwargs = gradient_checkpointing_kwargs,
|
||
include_inputs_for_metrics = include_inputs_for_metrics,
|
||
eval_do_concat_batches = eval_do_concat_batches,
|
||
fp16_backend = fp16_backend,
|
||
push_to_hub_model_id = push_to_hub_model_id,
|
||
push_to_hub_organization = push_to_hub_organization,
|
||
push_to_hub_token = push_to_hub_token,
|
||
mp_parameters = mp_parameters,
|
||
auto_find_batch_size = auto_find_batch_size,
|
||
full_determinism = full_determinism,
|
||
torchdynamo = torchdynamo,
|
||
ray_scope = ray_scope,
|
||
ddp_timeout = ddp_timeout,
|
||
torch_compile = torch_compile,
|
||
torch_compile_backend = torch_compile_backend,
|
||
torch_compile_mode = torch_compile_mode,
|
||
include_tokens_per_second = include_tokens_per_second,
|
||
include_num_input_tokens_seen = include_num_input_tokens_seen,
|
||
neftune_noise_alpha = neftune_noise_alpha,
|
||
optim_target_modules = optim_target_modules,
|
||
batch_eval_metrics = batch_eval_metrics,
|
||
eval_on_start = eval_on_start,
|
||
use_liger_kernel = use_liger_kernel,
|
||
eval_use_gather_object = eval_use_gather_object,
|
||
average_tokens_across_devices = average_tokens_across_devices,
|
||
model_init_kwargs = model_init_kwargs,
|
||
ref_model_init_kwargs = ref_model_init_kwargs,
|
||
model_adapter_name = model_adapter_name,
|
||
ref_adapter_name = ref_adapter_name,
|
||
force_use_ref_model = force_use_ref_model,
|
||
disable_dropout = disable_dropout,
|
||
use_logits_to_keep = use_logits_to_keep,
|
||
dataset_num_proc = dataset_num_proc,
|
||
padding_value = padding_value,
|
||
label_pad_token_id = label_pad_token_id,
|
||
max_prompt_length = max_prompt_length,
|
||
max_completion_length = max_completion_length,
|
||
max_length = max_length,
|
||
truncation_mode = truncation_mode,
|
||
padding_free = padding_free,
|
||
precompute_ref_log_probs = precompute_ref_log_probs,
|
||
precompute_ref_batch_size = precompute_ref_batch_size,
|
||
tools = tools,
|
||
loss_type = loss_type,
|
||
beta = beta,
|
||
f_alpha_divergence_coef = f_alpha_divergence_coef,
|
||
reference_free = reference_free,
|
||
label_smoothing = label_smoothing,
|
||
use_weighting = use_weighting,
|
||
rpo_alpha = rpo_alpha,
|
||
ld_alpha = ld_alpha,
|
||
discopop_tau = discopop_tau,
|
||
sync_ref_model = sync_ref_model,
|
||
ref_model_mixup_alpha = ref_model_mixup_alpha,
|
||
ref_model_sync_steps = ref_model_sync_steps,
|
||
generate_during_eval = generate_during_eval,**kwargs)
|
||
self.vllm_sampling_params = vllm_sampling_params
|
||
self.unsloth_num_chunks = unsloth_num_chunks
|
||
pass
|
||
|
||
class _UnslothDPOTrainer(Trainer):
|
||
r""""""
|
||
|
||
_tag_names = ["trl", "dpo"]
|
||
|
||
def __init__(
|
||
self,
|
||
model: Optional[Union[PreTrainedModel, nn.Module, str]] = None,
|
||
ref_model: Optional[Union[PreTrainedModel, nn.Module, str]] = None,
|
||
args: Optional[DPOConfig] = None,
|
||
data_collator: Optional[DataCollator] = None,
|
||
train_dataset: Optional[Dataset] = None,
|
||
eval_dataset: Optional[Union[Dataset, dict[str, Dataset]]] = None,
|
||
processing_class: Optional[
|
||
Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin]
|
||
] = None,
|
||
model_init: Optional[Callable[[], PreTrainedModel]] = None,
|
||
compute_metrics: Optional[Callable[[EvalLoopOutput], dict]] = None,
|
||
callbacks: Optional[list[TrainerCallback]] = None,
|
||
optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),
|
||
preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None,
|
||
peft_config: Optional[dict] = None,
|
||
):
|
||
if model is None:
|
||
raise ValueError("No model provided. Please provide a model to train.")
|
||
|
||
if not isinstance(model, str) and ref_model is model:
|
||
raise ValueError(
|
||
"`model` and `ref_model` cannot be the same object. If you want `ref_model` to be the "
|
||
"same as `model`, you must mass a copy of it, or `None` if you use peft."
|
||
)
|
||
|
||
if args.model_init_kwargs is None:
|
||
model_init_kwargs = {}
|
||
elif not isinstance(model, str):
|
||
raise ValueError(
|
||
"You passed model_init_kwargs to the DPOTrainer/DPOConfig, but your model is already instantiated."
|
||
)
|
||
else:
|
||
model_init_kwargs = args.model_init_kwargs
|
||
torch_dtype = model_init_kwargs.get("torch_dtype")
|
||
if torch_dtype is not None:
|
||
# Convert to `torch.dtype` if an str is passed
|
||
if isinstance(torch_dtype, str) and torch_dtype != "auto":
|
||
torch_dtype = getattr(torch, torch_dtype)
|
||
if torch_dtype != "auto" and not isinstance(torch_dtype, torch.dtype):
|
||
raise ValueError(
|
||
f"Invalid `torch_dtype` passed to the DPOConfig. Expected a string with either `torch.dtype` or 'auto', but got {torch_dtype}."
|
||
)
|
||
model_init_kwargs["torch_dtype"] = torch_dtype
|
||
|
||
if args.ref_model_init_kwargs is None:
|
||
ref_model_init_kwargs = {}
|
||
elif not isinstance(ref_model, str):
|
||
raise ValueError(
|
||
"You passed ref_model_init_kwargs to the DPOTrainer/DPOConfig, but your ref_model is already instantiated."
|
||
)
|
||
else:
|
||
ref_model_init_kwargs = args.ref_model_init_kwargs
|
||
torch_dtype = ref_model_init_kwargs.get("torch_dtype")
|
||
if torch_dtype is not None:
|
||
# Convert to `torch.dtype` if an str is passed
|
||
if isinstance(torch_dtype, str) and torch_dtype != "auto":
|
||
torch_dtype = getattr(torch, torch_dtype)
|
||
if torch_dtype != "auto" and not isinstance(torch_dtype, torch.dtype):
|
||
raise ValueError(
|
||
f"Invalid `torch_dtype` passed to the DPOConfig. Expected a string with either `torch.dtype` or 'auto', but got {torch_dtype}."
|
||
)
|
||
ref_model_init_kwargs["torch_dtype"] = torch_dtype
|
||
|
||
if isinstance(model, str):
|
||
model = AutoModelForCausalLM.from_pretrained(model, **model_init_kwargs)
|
||
|
||
if isinstance(ref_model, str):
|
||
ref_model = AutoModelForCausalLM.from_pretrained(ref_model, **ref_model_init_kwargs)
|
||
|
||
# Initialize this variable to False. This helps tracking the case when `peft_module_casting_to_bf16`
|
||
# has been called in order to properly call autocast if needed.
|
||
self._peft_has_been_casted_to_bf16 = False
|
||
|
||
if not is_peft_available() and peft_config is not None:
|
||
raise ValueError(
|
||
"PEFT is not installed and you passed a `peft_config` in the trainer's kwargs, please install it to use the PEFT models"
|
||
)
|
||
elif is_peft_available() and peft_config is not None:
|
||
# if model is a peft model and we have a peft_config, we merge and unload it first
|
||
if isinstance(model, PeftModel):
|
||
model = model.merge_and_unload()
|
||
|
||
if ref_model is not None and not args.force_use_ref_model:
|
||
raise ValueError(
|
||
"You passed both a ref_model and a peft_config. For training PEFT adapters with DPO there is no need to pass a reference"
|
||
" model. Please pass `ref_model=None` in case you want to train PEFT adapters, or pass a ref_model with `force_use_ref_model=True` in DPOTrainer's init."
|
||
" if you want to use a different ref_model."
|
||
)
|
||
|
||
if getattr(model, "is_loaded_in_8bit", False) or getattr(model, "is_loaded_in_4bit", False):
|
||
_support_gc_kwargs = hasattr(
|
||
args, "gradient_checkpointing_kwargs"
|
||
) and "gradient_checkpointing_kwargs" in list(
|
||
inspect.signature(prepare_model_for_kbit_training).parameters
|
||
)
|
||
|
||
prepare_model_kwargs = {"use_gradient_checkpointing": args.gradient_checkpointing}
|
||
|
||
if _support_gc_kwargs:
|
||
prepare_model_kwargs["gradient_checkpointing_kwargs"] = args.gradient_checkpointing_kwargs
|
||
|
||
model = prepare_model_for_kbit_training(model, **prepare_model_kwargs)
|
||
elif getattr(args, "gradient_checkpointing", False):
|
||
# For backward compatibility with older versions of transformers
|
||
if hasattr(model, "enable_input_require_grads"):
|
||
model.enable_input_require_grads()
|
||
else:
|
||
|
||
def make_inputs_require_grad(module, input, output):
|
||
output.requires_grad_(True)
|
||
|
||
model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)
|
||
|
||
# get peft model with the given config
|
||
model = model
|
||
if args.bf16 and getattr(model, "is_loaded_in_4bit", False):
|
||
peft_module_casting_to_bf16(model)
|
||
# If args.bf16 we need to explicitly call `generate` with torch amp autocast context manager
|
||
self._peft_has_been_casted_to_bf16 = True
|
||
|
||
# For models that use gradient_checkpointing, we need to attach a hook that enables input
|
||
# to explicitly have `requires_grad=True`, otherwise training will either silently
|
||
# fail or completely fail.
|
||
elif getattr(args, "gradient_checkpointing", False):
|
||
# For backward compatibility with older versions of transformers
|
||
if hasattr(model, "enable_input_require_grads"):
|
||
model.enable_input_require_grads()
|
||
else:
|
||
|
||
def make_inputs_require_grad(module, input, output):
|
||
output.requires_grad_(True)
|
||
|
||
model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)
|
||
|
||
if args.generate_during_eval and not (is_wandb_available() or is_comet_available()):
|
||
raise ValueError(
|
||
"`generate_during_eval=True` requires Weights and Biases or Comet to be installed."
|
||
" Please install `wandb` or `comet-ml` to resolve."
|
||
)
|
||
|
||
self.is_encoder_decoder = model.config.is_encoder_decoder
|
||
self.is_vision_model = model.config.model_type in MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES.keys()
|
||
self.is_peft_model = is_peft_available() and isinstance(model, PeftModel)
|
||
self.model_adapter_name = args.model_adapter_name
|
||
self.ref_adapter_name = args.ref_adapter_name
|
||
self.reference_free = args.reference_free
|
||
|
||
if ref_model:
|
||
self.ref_model = ref_model
|
||
elif self.is_peft_model or args.precompute_ref_log_probs:
|
||
# The `model` with adapters turned off will be used as the reference model
|
||
self.ref_model = None
|
||
else:
|
||
self.ref_model = create_reference_model(model)
|
||
|
||
if processing_class is None:
|
||
raise ValueError("processing_class must be specified to tokenize a DPO dataset.")
|
||
|
||
if args.padding_value is not None:
|
||
self.padding_value = args.padding_value
|
||
else:
|
||
if hasattr(processing_class, "pad_token_id") and processing_class.pad_token_id is not None:
|
||
self.padding_value = processing_class.pad_token_id
|
||
elif hasattr(processing_class, "tokenizer") and processing_class.tokenizer.pad_token_id is not None:
|
||
self.padding_value = processing_class.tokenizer.pad_token_id
|
||
else:
|
||
raise ValueError(
|
||
"`padding_value` is not specified in `DPOConfig`, and `pad_token_id` is missing in the "
|
||
"`processing_class`. Please either set the `padding_value` argument in `DPOConfig`, or set "
|
||
"`tokenizer.pad_token` (e.g., `tokenizer.pad_token = tokenizer.eos_token`) before instantiating "
|
||
"the trainer."
|
||
)
|
||
|
||
if data_collator is None:
|
||
data_collator = DataCollatorForPreference(pad_token_id=self.padding_value)
|
||
|
||
# Disable dropout in the model and reference model
|
||
if args.disable_dropout:
|
||
disable_dropout_in_model(model)
|
||
if self.ref_model is not None:
|
||
disable_dropout_in_model(self.ref_model)
|
||
|
||
self.generate_during_eval = args.generate_during_eval
|
||
self.label_pad_token_id = args.label_pad_token_id
|
||
self.max_prompt_length = args.max_prompt_length
|
||
self.max_completion_length = args.max_completion_length
|
||
self.max_length = args.max_length
|
||
self.truncation_mode = args.truncation_mode
|
||
self.precompute_ref_log_probs = args.precompute_ref_log_probs
|
||
self.use_logits_to_keep = args.use_logits_to_keep
|
||
|
||
if args.padding_free:
|
||
if model.config._attn_implementation != "flash_attention_2":
|
||
warnings.warn(
|
||
"Padding-free training is enabled, but the attention implementation is not set to "
|
||
"'flash_attention_2'. Padding-free training flattens batches into a single sequence, and "
|
||
"'flash_attention_2' is the only known attention mechanism that reliably supports this. Using "
|
||
"other implementations may lead to unexpected behavior. To ensure compatibility, set "
|
||
"`attn_implementation='flash_attention_2'` in the model configuration, or verify that your "
|
||
"attention mechanism can handle flattened sequences."
|
||
)
|
||
self.padding_free = args.padding_free
|
||
|
||
# Since ref_logs are precomputed on the first call to get_train/eval_dataloader
|
||
# keep track of first called to avoid computation of future calls
|
||
self._precomputed_train_ref_log_probs = False
|
||
self._precomputed_eval_ref_log_probs = False
|
||
|
||
if (
|
||
args.loss_type in ["hinge", "ipo", "bco_pair", "sppo_hard", "nca_pair", "apo_zero", "apo_down"]
|
||
and args.label_smoothing > 0
|
||
):
|
||
warnings.warn(
|
||
f"You are using the {args.loss_type} loss type that does not support label smoothing. The "
|
||
"`label_smoothing` parameter will be ignored. Set `label_smoothing` to `0.0` to remove this warning.",
|
||
UserWarning,
|
||
)
|
||
if args.loss_type == "kto_pair":
|
||
raise ValueError("Support for kto_pair has been removed in DPOTrainer. Please use KTOTrainer.")
|
||
|
||
self.beta = args.beta
|
||
self.label_smoothing = args.label_smoothing
|
||
self.loss_type = args.loss_type
|
||
self.aux_loss_enabled = getattr(model.config, "output_router_logits", False)
|
||
self.use_weighting = args.use_weighting
|
||
self.aux_loss_coef = getattr(model.config, "router_aux_loss_coef", 0.0)
|
||
if self.aux_loss_enabled and self.aux_loss_coef == 0.0:
|
||
warnings.warn(
|
||
"You set `output_router_logits` to `True` in the model config, but `router_aux_loss_coef` is set to "
|
||
"`0.0`, meaning the auxiliary loss will not be used. Either set `router_aux_loss_coef` to a value "
|
||
"greater than `0.0`, or set `output_router_logits` to `False` if you don't want to use the auxiliary "
|
||
"loss.",
|
||
UserWarning,
|
||
)
|
||
|
||
self._stored_metrics = defaultdict(lambda: defaultdict(list))
|
||
self.f_divergence_type = args.f_divergence_type
|
||
self.f_divergence_params = {FDivergenceConstants.ALPHA_DIVERGENCE_COEF_KEY: args.f_alpha_divergence_coef}
|
||
self.dataset_num_proc = args.dataset_num_proc
|
||
|
||
# The trainer estimates the number of FLOPs (floating-point operations) using the number of elements in the
|
||
# input tensor associated with the key "input_ids". However, in DPO, the sampled data does not include the
|
||
# "input_ids" key. Instead, the available keys are "prompt_input_ids", "chosen_input_ids", and
|
||
# "rejected_input_ids". As a result, the trainer issues the warning: "Could not estimate the number of tokens
|
||
# of the input, floating-point operations will not be computed." To suppress this warning, we set the
|
||
# "estimate_tokens" key in the model's "warnings_issued" dictionary to True. This acts as a flag to indicate
|
||
# that the warning has already been issued.
|
||
model.warnings_issued["estimate_tokens"] = True
|
||
|
||
# Dataset preparation
|
||
train_dataset = self._prepare_dataset(train_dataset, processing_class, args, "train")
|
||
if eval_dataset is not None:
|
||
if isinstance(eval_dataset, dict):
|
||
eval_dataset = {
|
||
key: self._prepare_dataset(dataset, processing_class, args, key)
|
||
for key, dataset in eval_dataset.items()
|
||
}
|
||
else:
|
||
eval_dataset = self._prepare_dataset(eval_dataset, processing_class, args, "eval")
|
||
|
||
super().__init__(
|
||
model=model,
|
||
args=args,
|
||
data_collator=data_collator,
|
||
train_dataset=train_dataset,
|
||
eval_dataset=eval_dataset,
|
||
processing_class=processing_class,
|
||
model_init=model_init,
|
||
compute_metrics=compute_metrics,
|
||
callbacks=callbacks,
|
||
optimizers=optimizers,
|
||
preprocess_logits_for_metrics=preprocess_logits_for_metrics,
|
||
)
|
||
|
||
# Gradient accumulation requires scaled loss. Normally, loss scaling in the parent class depends on whether the
|
||
# model accepts loss-related kwargs. Since we compute our own loss, this check is irrelevant. We set
|
||
# self.model_accepts_loss_kwargs to False to enable scaling.
|
||
self.model_accepts_loss_kwargs = False
|
||
|
||
# Add tags for models that have been loaded with the correct transformers version
|
||
if hasattr(self.model, "add_model_tags"):
|
||
self.model.add_model_tags(self._tag_names)
|
||
|
||
if not hasattr(self, "accelerator"):
|
||
raise AttributeError(
|
||
"Your `Trainer` does not have an `accelerator` object. Consider upgrading `transformers`."
|
||
)
|
||
|
||
# Deepspeed Zero-3 does not support precompute_ref_log_probs
|
||
if self.is_deepspeed_enabled:
|
||
if self.accelerator.state.deepspeed_plugin.zero_stage == 3 and self.precompute_ref_log_probs:
|
||
raise ValueError(
|
||
"You cannot use `precompute_ref_log_probs=True` with Deepspeed ZeRO-3. Please set `precompute_ref_log_probs=False`."
|
||
)
|
||
|
||
if self.ref_model is None:
|
||
if not (self.is_peft_model or self.precompute_ref_log_probs):
|
||
raise ValueError(
|
||
"No reference model and model is not a Peft model. Try setting `precompute_ref_log_probs=True`"
|
||
)
|
||
if args.sync_ref_model:
|
||
raise ValueError(
|
||
"You currently cannot use `ref_model=None` with TR-DPO method. Please provide `ref_model`."
|
||
)
|
||
else:
|
||
if self.is_deepspeed_enabled:
|
||
self.ref_model = prepare_deepspeed(self.ref_model, self.accelerator)
|
||
elif self.is_fsdp_enabled:
|
||
self.ref_model = prepare_fsdp(self.ref_model, self.accelerator)
|
||
else:
|
||
self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True)
|
||
|
||
if args.sync_ref_model:
|
||
if self.precompute_ref_log_probs:
|
||
raise ValueError(
|
||
"You cannot use `precompute_ref_log_probs=True` with TR-DPO method. Please set `precompute_ref_log_probs=False`."
|
||
)
|
||
|
||
self.add_callback(SyncRefModelCallback(ref_model=self.ref_model, accelerator=self.accelerator))
|
||
|
||
if self.loss_type == "bco_pair":
|
||
self.running = RunningMoments(self.accelerator)
|
||
|
||
def _prepare_dataset(
|
||
self,
|
||
dataset: Union[Dataset, IterableDataset],
|
||
processing_class: Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin],
|
||
args: DPOConfig,
|
||
dataset_name: str,
|
||
) -> Union[Dataset, IterableDataset]:
|
||
# Build the kwargs for the `map` function
|
||
map_kwargs = {"writer_batch_size": 10}
|
||
if isinstance(dataset, Dataset): # IterableDataset does not support num_proc
|
||
map_kwargs["num_proc"] = args.dataset_num_proc
|
||
|
||
with PartialState().main_process_first():
|
||
# Extract prompt if needed
|
||
if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc`
|
||
map_kwargs["desc"] = f"Extracting prompt in {dataset_name} dataset"
|
||
dataset = dataset.map(maybe_extract_prompt, **map_kwargs)
|
||
|
||
# Apply the chat template if needed
|
||
if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc`
|
||
map_kwargs["desc"] = f"Applying chat template to {dataset_name} dataset"
|
||
dataset = dataset.map(
|
||
maybe_apply_chat_template, fn_kwargs={"tokenizer": processing_class, "tools": args.tools}, **map_kwargs
|
||
)
|
||
|
||
# Tokenize the dataset
|
||
if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc`
|
||
map_kwargs["desc"] = f"Tokenizing {dataset_name} dataset"
|
||
|
||
dataset = dataset.map(
|
||
self.tokenize_row if not self.is_vision_model else self.process_row,
|
||
remove_columns=["chosen", "rejected"],
|
||
fn_kwargs={
|
||
"processing_class": processing_class,
|
||
"max_prompt_length": args.max_prompt_length,
|
||
"max_completion_length": args.max_completion_length,
|
||
# for enc-dec, we add the special tokens ([bos_token] + prompt + [eos_token]; completion + [eos_token])
|
||
"add_special_tokens": False,
|
||
},
|
||
**map_kwargs,
|
||
)
|
||
|
||
return dataset
|
||
|
||
@staticmethod
|
||
def tokenize_row(features, processing_class, max_prompt_length, max_completion_length, add_special_tokens):
|
||
"""
|
||
Tokenize a row of the dataset.
|
||
|
||
Args:
|
||
features (`dict[str, str]`):
|
||
Row of the dataset, should contain the keys `"prompt"`, `"chosen"`, and `"rejected"`.
|
||
processing_class (`PreTrainedTokenizerBase`):
|
||
Processing class used to process the data.
|
||
max_prompt_length (`int` or `None`):
|
||
Maximum length of the prompt sequence. If `None`, the prompt sequence is not truncated.
|
||
max_completion_length (`int` or `None`):
|
||
Maximum length of the completion sequences. If `None`, the completion sequences are not truncated.
|
||
add_special_tokens (`bool`):
|
||
Whether to add special tokens to the sequences. Typically used for encoder-decoder models. If `True`,
|
||
the prompt sequence will have a bos token prepended and an eos token appended. In any case, the
|
||
completion sequences will have an eos token appended.
|
||
|
||
Returns:
|
||
`dict[str, list[int]]`:
|
||
Tokenized sequences with the keys `"prompt_input_ids"`, `"chosen_input_ids"`, and
|
||
`"rejected_input_ids".
|
||
|
||
Example:
|
||
```python
|
||
>>> from transformers import GPT2Tokenizer
|
||
>>> tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
|
||
>>> features = {"prompt": "The sky is", "chosen": " blue", "rejected": " green"}
|
||
>>> DPOTrainer.tokenize_row(
|
||
... features, tokenizer, max_prompt_length=3, max_completion_length=3, add_special_tokens=False
|
||
... )
|
||
{'prompt_input_ids': [464, 6766, 318], 'chosen_input_ids': [4171, 50256], 'rejected_input_ids': [4077, 50256]}
|
||
```
|
||
"""
|
||
tokenizer = processing_class # the processing class is a tokenizer
|
||
prompt_input_ids = tokenizer(features["prompt"], add_special_tokens=False)["input_ids"]
|
||
chosen_input_ids = tokenizer(features["chosen"], add_special_tokens=False)["input_ids"]
|
||
rejected_input_ids = tokenizer(features["rejected"], add_special_tokens=False)["input_ids"]
|
||
|
||
# Add special tokens (typically for encoder-decoder models)
|
||
if add_special_tokens:
|
||
if tokenizer.bos_token_id is not None:
|
||
prompt_input_ids = [tokenizer.bos_token_id] + prompt_input_ids
|
||
if tokenizer.eos_token_id is not None:
|
||
prompt_input_ids = prompt_input_ids + [tokenizer.eos_token_id]
|
||
chosen_input_ids = chosen_input_ids + [tokenizer.eos_token_id]
|
||
rejected_input_ids = rejected_input_ids + [tokenizer.eos_token_id]
|
||
|
||
# Truncate prompt and completion sequences
|
||
if max_prompt_length is not None:
|
||
prompt_input_ids = prompt_input_ids[-max_prompt_length:]
|
||
if max_completion_length is not None:
|
||
chosen_input_ids = chosen_input_ids[:max_completion_length]
|
||
rejected_input_ids = rejected_input_ids[:max_completion_length]
|
||
|
||
return {
|
||
"prompt_input_ids": prompt_input_ids,
|
||
"chosen_input_ids": chosen_input_ids,
|
||
"rejected_input_ids": rejected_input_ids,
|
||
}
|
||
|
||
@staticmethod
|
||
def process_row(features, processing_class, max_prompt_length, max_completion_length, add_special_tokens):
|
||
"""
|
||
Same as `tokenize_row` but for vision models. Please refer to `tokenize_row` for more information.
|
||
"""
|
||
processor, tokenizer = processing_class, processing_class.tokenizer # the processing class is a processor
|
||
processed_features = processor(images=features["images"], text=features["prompt"], add_special_tokens=False)
|
||
|
||
prompt_input_ids = processed_features["input_ids"][0]
|
||
pixel_values = processed_features["pixel_values"][0]
|
||
chosen_input_ids = tokenizer(features["chosen"], add_special_tokens=False)["input_ids"]
|
||
rejected_input_ids = tokenizer(features["rejected"], add_special_tokens=False)["input_ids"]
|
||
|
||
# Add special tokens (typically for encoder-decoder models)
|
||
if add_special_tokens:
|
||
if tokenizer.bos_token_id is not None:
|
||
prompt_input_ids = [tokenizer.bos_token_id] + prompt_input_ids
|
||
if tokenizer.eos_token_id is not None:
|
||
prompt_input_ids = prompt_input_ids + [tokenizer.eos_token_id]
|
||
chosen_input_ids = chosen_input_ids + [tokenizer.eos_token_id]
|
||
rejected_input_ids = rejected_input_ids + [tokenizer.eos_token_id]
|
||
|
||
# Truncate prompt and completion sequences
|
||
if max_prompt_length is not None:
|
||
prompt_input_ids = prompt_input_ids[-max_prompt_length:]
|
||
if max_completion_length is not None:
|
||
chosen_input_ids = chosen_input_ids[:max_completion_length]
|
||
rejected_input_ids = rejected_input_ids[:max_completion_length]
|
||
|
||
output = {
|
||
"prompt_input_ids": prompt_input_ids,
|
||
"pixel_values": pixel_values,
|
||
"chosen_input_ids": chosen_input_ids,
|
||
"rejected_input_ids": rejected_input_ids,
|
||
}
|
||
|
||
if "pixel_attention_mask" in processed_features:
|
||
output["pixel_attention_mask"] = processed_features["pixel_attention_mask"][0]
|
||
if "image_sizes" in processed_features:
|
||
output["image_sizes"] = processed_features["image_sizes"][0]
|
||
|
||
return output
|
||
|
||
def _set_signature_columns_if_needed(self):
|
||
# If `self.args.remove_unused_columns` is True, non-signature columns are removed.
|
||
# By default, this method sets `self._signature_columns` to the model's expected inputs.
|
||
# In DPOTrainer, we preprocess data, so using the model's signature columns doesn't work.
|
||
# Instead, we set them to the columns expected by `DataCollatorForPreference`, hence the override.
|
||
if self._signature_columns is None:
|
||
self._signature_columns = [
|
||
"prompt_input_ids",
|
||
"chosen_input_ids",
|
||
"rejected_input_ids",
|
||
"image_sizes",
|
||
"ref_chosen_logps",
|
||
"ref_rejected_logps",
|
||
]
|
||
|
||
def get_train_dataloader(self) -> DataLoader:
|
||
"""
|
||
Returns the training [`~torch.utils.data.DataLoader`].
|
||
|
||
Subclass of transformers.src.transformers.trainer.get_train_dataloader to precompute `ref_log_probs`.
|
||
"""
|
||
|
||
if self.precompute_ref_log_probs and not self._precomputed_train_ref_log_probs:
|
||
batch_size = self.args.precompute_ref_batch_size or self.args.per_device_train_batch_size
|
||
dataloader_params = {
|
||
"batch_size": batch_size,
|
||
"collate_fn": self.data_collator,
|
||
"num_workers": self.args.dataloader_num_workers,
|
||
"pin_memory": self.args.dataloader_pin_memory,
|
||
"shuffle": False,
|
||
}
|
||
|
||
# prepare dataloader
|
||
data_loader = self.accelerator.prepare(DataLoader(self.train_dataset, **dataloader_params))
|
||
|
||
ref_chosen_logps = []
|
||
ref_rejected_logps = []
|
||
for padded_batch in tqdm(iterable=data_loader, desc="Train dataset reference log probs"):
|
||
ref_chosen_logp, ref_rejected_logp = self.compute_ref_log_probs(padded_batch)
|
||
ref_chosen_logp, ref_rejected_logp = self.accelerator.gather_for_metrics(
|
||
(ref_chosen_logp, ref_rejected_logp)
|
||
)
|
||
ref_chosen_logps.append(ref_chosen_logp.cpu())
|
||
ref_rejected_logps.append(ref_rejected_logp.cpu())
|
||
|
||
# Unnecessary cache clearing to avoid OOM
|
||
empty_cache()
|
||
self.accelerator.free_memory()
|
||
|
||
all_ref_chosen_logps = torch.cat(ref_chosen_logps).float().numpy()
|
||
all_ref_rejected_logps = torch.cat(ref_rejected_logps).float().numpy()
|
||
|
||
self.train_dataset = self.train_dataset.add_column(name="ref_chosen_logps", column=all_ref_chosen_logps)
|
||
self.train_dataset = self.train_dataset.add_column(
|
||
name="ref_rejected_logps", column=all_ref_rejected_logps
|
||
)
|
||
|
||
self._precomputed_train_ref_log_probs = True
|
||
|
||
return super().get_train_dataloader()
|
||
|
||
def get_eval_dataloader(self, eval_dataset: Optional[Dataset] = None) -> DataLoader:
|
||
"""
|
||
Returns the evaluation [`~torch.utils.data.DataLoader`].
|
||
|
||
Subclass of transformers.src.transformers.trainer.get_eval_dataloader to precompute `ref_log_probs`.
|
||
|
||
Args:
|
||
eval_dataset (`torch.utils.data.Dataset`, *optional*):
|
||
If provided, will override `self.eval_dataset`. If it is a [`~datasets.Dataset`], columns not accepted
|
||
by the `model.forward()` method are automatically removed. It must implement `__len__`.
|
||
"""
|
||
if eval_dataset is None and self.eval_dataset is None:
|
||
raise ValueError("Trainer: evaluation requires an eval_dataset.")
|
||
eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset
|
||
|
||
if self.precompute_ref_log_probs and not self._precomputed_eval_ref_log_probs:
|
||
batch_size = self.args.precompute_ref_batch_size or self.args.per_device_eval_batch_size
|
||
dataloader_params = {
|
||
"batch_size": batch_size,
|
||
"collate_fn": self.data_collator,
|
||
"num_workers": self.args.dataloader_num_workers,
|
||
"pin_memory": self.args.dataloader_pin_memory,
|
||
"shuffle": False,
|
||
}
|
||
|
||
# prepare dataloader
|
||
data_loader = self.accelerator.prepare(DataLoader(eval_dataset, **dataloader_params))
|
||
|
||
ref_chosen_logps = []
|
||
ref_rejected_logps = []
|
||
for padded_batch in tqdm(iterable=data_loader, desc="Eval dataset reference log probs"):
|
||
ref_chosen_logp, ref_rejected_logp = self.compute_ref_log_probs(padded_batch)
|
||
ref_chosen_logp, ref_rejected_logp = self.accelerator.gather_for_metrics(
|
||
(ref_chosen_logp, ref_rejected_logp)
|
||
)
|
||
ref_chosen_logps.append(ref_chosen_logp.cpu())
|
||
ref_rejected_logps.append(ref_rejected_logp.cpu())
|
||
|
||
all_ref_chosen_logps = torch.cat(ref_chosen_logps).float().numpy()
|
||
all_ref_rejected_logps = torch.cat(ref_rejected_logps).float().numpy()
|
||
|
||
eval_dataset = eval_dataset.add_column(name="ref_chosen_logps", column=all_ref_chosen_logps)
|
||
eval_dataset = eval_dataset.add_column(name="ref_rejected_logps", column=all_ref_rejected_logps)
|
||
|
||
# Save calculated ref_chosen_logps and ref_rejected_logps to the eval_dataset for subsequent runs
|
||
if self.eval_dataset is not None:
|
||
self.eval_dataset = eval_dataset
|
||
self._precomputed_eval_ref_log_probs = True
|
||
|
||
return super().get_eval_dataloader(eval_dataset=eval_dataset)
|
||
|
||
@contextmanager
|
||
def null_ref_context(self):
|
||
"""Context manager for handling null reference model (that is, peft adapter manipulation)."""
|
||
with (
|
||
self.accelerator.unwrap_model(self.model).disable_adapter()
|
||
if self.is_peft_model and not self.ref_adapter_name
|
||
else nullcontext()
|
||
):
|
||
if self.ref_adapter_name:
|
||
self.model.set_adapter(self.ref_adapter_name)
|
||
yield
|
||
if self.ref_adapter_name:
|
||
self.model.set_adapter(self.model_adapter_name or "default")
|
||
|
||
def compute_ref_log_probs(self, batch: dict[str, torch.LongTensor]) -> dict:
|
||
"""Computes log probabilities of the reference model for a single padded batch of a DPO specific dataset."""
|
||
device_type = "xpu" if is_torch_xpu_available() else "cuda"
|
||
compte_ref_context_manager = amp.autocast(device_type) if self._peft_has_been_casted_to_bf16 else nullcontext()
|
||
with torch.no_grad(), compte_ref_context_manager:
|
||
if self.ref_model is None:
|
||
with self.null_ref_context():
|
||
ref_model_output = self.concatenated_forward(self.model, batch, is_ref_model=True)
|
||
else:
|
||
ref_model_output = self.concatenated_forward(self.ref_model, batch, is_ref_model=True)
|
||
return ref_model_output["chosen_logps"], ref_model_output["rejected_logps"]
|
||
|
||
@staticmethod
|
||
def concatenated_inputs(
|
||
batch: dict[str, Union[list, torch.LongTensor]], padding_value: int
|
||
) -> dict[str, torch.LongTensor]:
|
||
"""
|
||
Concatenate the `chosen` and `rejected` inputs from the batch into a single tensor for both the prompt
|
||
and completion sequences.
|
||
|
||
Args:
|
||
batch (`dict[str, Union[list, torch.LongTensor]]`):
|
||
A batch of input data. The batch must contain the following keys:
|
||
|
||
- `"prompt_input_ids"`: Tensor of shape `(batch_size, prompt_length)` representing the prompt input IDs.
|
||
- `"chosen_input_ids"`: Tensor of shape `(batch_size, chosen_length)` representing the chosen completion input IDs.
|
||
- `"rejected_input_ids"`: Tensor of shape `(batch_size, rejected_length)` representing the rejected completion input IDs.
|
||
- `"prompt_pixel_values"` (optional): Tensor for pixel values, if available.
|
||
- `"prompt_pixel_attention_mask"` (optional): Tensor for pixel attention masks, if available.
|
||
|
||
padding_value (`int`):
|
||
The padding value to use for the concatenated completion sequences (`chosen_input_ids` and
|
||
`rejected_input_ids`).
|
||
|
||
Returns:
|
||
`dict[str, torch.LongTensor]`: A dictionary containing:
|
||
|
||
- `"prompt_input_ids"`: Concatenated prompt input IDs of shape `(2 * batch_size, prompt_length)`.
|
||
- `"completion_input_ids"`: Concatenated chosen and rejected completion input IDs of shape `(2 * batch_size, max_completion_length)`.
|
||
- `"prompt_attention_mask"`: Concatenated prompt attention masks of shape `(2 * batch_size, prompt_length)`.
|
||
- `"completion_attention_mask"`: Concatenated chosen and rejected attention masks of shape `(2 * batch_size, max_completion_length)`.
|
||
- `"pixel_values"` (optional): Concatenated pixel values if `"prompt_pixel_values"` are present.
|
||
- `"pixel_attention_mask"` (optional): Concatenated pixel attention masks if `"prompt_pixel_attention_mask"` are present.
|
||
|
||
Notes:
|
||
The completion input IDs and attention masks are padded to the maximum completion length of the chosen
|
||
or rejected sequences.
|
||
"""
|
||
output = {}
|
||
|
||
# For the prompt, the input_ids are the same for both the chosen and rejected responses
|
||
output["prompt_input_ids"] = torch.cat([batch["prompt_input_ids"], batch["prompt_input_ids"]], dim=0)
|
||
output["prompt_attention_mask"] = torch.cat(
|
||
[batch["prompt_attention_mask"], batch["prompt_attention_mask"]], dim=0
|
||
)
|
||
if "pixel_values" in batch:
|
||
output["pixel_values"] = torch.cat([batch["pixel_values"], batch["pixel_values"]], dim=0)
|
||
|
||
if "pixel_attention_mask" in batch:
|
||
output["pixel_attention_mask"] = torch.cat(
|
||
[batch["pixel_attention_mask"], batch["pixel_attention_mask"]], dim=0
|
||
)
|
||
if "image_sizes" in batch:
|
||
output["image_sizes"] = torch.cat([batch["image_sizes"], batch["image_sizes"]], dim=0)
|
||
|
||
# Concatenate the chosen and rejected completions
|
||
max_completion_length = max(batch["chosen_input_ids"].shape[1], batch["rejected_input_ids"].shape[1])
|
||
output["completion_input_ids"] = torch.cat(
|
||
(
|
||
pad_to_length(batch["chosen_input_ids"], max_completion_length, pad_value=padding_value),
|
||
pad_to_length(batch["rejected_input_ids"], max_completion_length, pad_value=padding_value),
|
||
),
|
||
)
|
||
output["completion_attention_mask"] = torch.cat(
|
||
(
|
||
pad_to_length(batch["chosen_attention_mask"], max_completion_length, pad_value=0),
|
||
pad_to_length(batch["rejected_attention_mask"], max_completion_length, pad_value=0),
|
||
),
|
||
)
|
||
|
||
return output
|
||
|
||
def dpo_loss(
|
||
self,
|
||
chosen_logps: torch.FloatTensor,
|
||
rejected_logps: torch.FloatTensor,
|
||
ref_chosen_logps: torch.FloatTensor,
|
||
ref_rejected_logps: torch.FloatTensor,
|
||
) -> tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
|
||
"""
|
||
Compute the DPO loss for a batch of policy and reference model log probabilities.
|
||
|
||
Args:
|
||
chosen_logps (`torch.FloatTensor`):
|
||
Log probabilities of the model for the chosen responses. Shape: `(batch_size,)`.
|
||
rejected_logps (`torch.FloatTensor`):
|
||
Log probabilities of the model for the rejected responses. Shape: `(batch_size,)`.
|
||
ref_chosen_logps (`torch.FloatTensor`):
|
||
Log probabilities of the reference model for the chosen responses. Shape: `(batch_size,)`.
|
||
ref_rejected_logps (`torch.FloatTensor`):
|
||
Log probabilities of the reference model for the rejected responses. Shape: `(batch_size,)`.
|
||
|
||
Returns:
|
||
A tuple of three tensors: `(losses, chosen_rewards, rejected_rewards)`.
|
||
The losses tensor contains the DPO loss for each example in the batch.
|
||
The `chosen_rewards` and `rejected_rewards` tensors contain the rewards for the chosen and rejected
|
||
responses, respectively.
|
||
"""
|
||
device = self.accelerator.device
|
||
|
||
# Get the log ratios for the chosen and rejected responses
|
||
chosen_logratios = chosen_logps.to(device) - (not self.reference_free) * ref_chosen_logps.to(device)
|
||
rejected_logratios = rejected_logps.to(device) - (not self.reference_free) * ref_rejected_logps.to(device)
|
||
|
||
if self.f_divergence_type == FDivergenceType.ALPHA_DIVERGENCE.value:
|
||
# The alpha-divergence formula: (1 - u^-alpha) / alpha
|
||
# The divergence difference between the chosen and rejected sample is:
|
||
# (1 - u[w]^-alpha) / alpha - (1 - u[l]^-alpha) / alpha
|
||
# = (u[l]^-alpha - u[w]^-alpha) / alpha
|
||
# where u[w] and u[l] are the policy/reference probability ratios
|
||
# for the chosen and rejected samples, respectively.
|
||
alpha_coef = FDivergenceConstants.ALPHA_DIVERGENCE_COEF_DEFAULT
|
||
if self.f_divergence_params and FDivergenceConstants.ALPHA_DIVERGENCE_COEF_KEY in self.f_divergence_params:
|
||
alpha_coef = float(self.f_divergence_params[FDivergenceConstants.ALPHA_DIVERGENCE_COEF_KEY])
|
||
logits = (cap_exp(rejected_logratios * -alpha_coef) - cap_exp(chosen_logratios * -alpha_coef)) / alpha_coef
|
||
else:
|
||
logratios = chosen_logps - rejected_logps
|
||
if self.reference_free:
|
||
ref_logratios = torch.tensor([0], dtype=logratios.dtype, device=logratios.device)
|
||
else:
|
||
ref_logratios = ref_chosen_logps - ref_rejected_logps
|
||
|
||
logratios = logratios.to(self.accelerator.device)
|
||
ref_logratios = ref_logratios.to(self.accelerator.device)
|
||
logits = logratios - ref_logratios
|
||
|
||
if self.f_divergence_type == FDivergenceType.JS_DIVERGENCE.value:
|
||
# The js-divergence formula: log(2 * u / (1 + u))
|
||
# The divergence difference between the chosen and rejected sample is:
|
||
# log(2 * u[w] / (1 + u[w])) - log(2 * u[l] / (1 + u[l]))
|
||
# = log(u[w]) - log(u[l]) - (log(1 + u[w]) - log(1 + u[l]))
|
||
# where u[w] and u[l] are the policy/reference probability ratios
|
||
# for the chosen and rejected samples, respectively.
|
||
logits -= F.softplus(chosen_logratios) - F.softplus(rejected_logratios)
|
||
|
||
# The beta is a temperature parameter for the DPO loss, typically something in the range of 0.1 to 0.5.
|
||
# We ignore the reference model as beta -> 0. The label_smoothing parameter encodes our uncertainty about the
|
||
# labels and calculates a conservative DPO loss.
|
||
if self.loss_type == "sigmoid":
|
||
losses = (
|
||
-F.logsigmoid(self.beta * logits) * (1 - self.label_smoothing)
|
||
- F.logsigmoid(-self.beta * logits) * self.label_smoothing
|
||
)
|
||
|
||
elif self.loss_type == "robust":
|
||
losses = (
|
||
-F.logsigmoid(self.beta * logits) * (1 - self.label_smoothing)
|
||
+ F.logsigmoid(-self.beta * logits) * self.label_smoothing
|
||
) / (1 - 2 * self.label_smoothing)
|
||
|
||
elif self.loss_type == "exo_pair":
|
||
# eqn (16) of the EXO paper: https://huggingface.co/papers/2402.00856
|
||
import math
|
||
|
||
if self.label_smoothing == 0:
|
||
self.label_smoothing = 1e-3
|
||
losses = (self.beta * logits).sigmoid() * (
|
||
F.logsigmoid(self.beta * logits) - math.log(1 - self.label_smoothing)
|
||
) + (-self.beta * logits).sigmoid() * (F.logsigmoid(-self.beta * logits) - math.log(self.label_smoothing))
|
||
|
||
elif self.loss_type == "hinge":
|
||
losses = torch.relu(1 - self.beta * logits)
|
||
|
||
elif self.loss_type == "ipo":
|
||
# eqn (17) of the paper where beta is the regularization parameter for the IPO loss, denoted by tau in the paper.
|
||
losses = (logits - 1 / (2 * self.beta)) ** 2
|
||
|
||
elif self.loss_type == "bco_pair":
|
||
chosen_logratios = chosen_logps - ref_chosen_logps
|
||
rejected_logratios = rejected_logps - ref_rejected_logps
|
||
chosen_rewards = self.beta * chosen_logratios
|
||
rejected_rewards = self.beta * rejected_logratios
|
||
rewards = torch.cat((chosen_rewards, rejected_rewards), 0).mean().detach()
|
||
self.running.update(rewards)
|
||
delta = self.running.mean
|
||
losses = -F.logsigmoid((self.beta * chosen_logratios) - delta) - F.logsigmoid(
|
||
-(self.beta * rejected_logratios - delta)
|
||
)
|
||
|
||
elif self.loss_type == "sppo_hard":
|
||
# In the paper (https://huggingface.co/papers/2405.00675), SPPO employs a soft probability approach,
|
||
# estimated using the PairRM score. The probability calculation is conducted outside of the trainer class.
|
||
# The version described here is the hard probability version, where P in Equation (4.7) of Algorithm 1 is
|
||
# set to 1 for the winner and 0 for the loser.
|
||
a = chosen_logps - ref_chosen_logps
|
||
b = rejected_logps - ref_rejected_logps
|
||
losses = (a - 0.5 / self.beta) ** 2 + (b + 0.5 / self.beta) ** 2
|
||
|
||
elif self.loss_type == "nca_pair":
|
||
chosen_rewards = (chosen_logps - ref_chosen_logps) * self.beta
|
||
rejected_rewards = (rejected_logps - ref_rejected_logps) * self.beta
|
||
losses = (
|
||
-F.logsigmoid(chosen_rewards)
|
||
- 0.5 * F.logsigmoid(-chosen_rewards)
|
||
- 0.5 * F.logsigmoid(-rejected_rewards)
|
||
)
|
||
|
||
elif self.loss_type == "aot_pair":
|
||
chosen_logratios = chosen_logps - ref_chosen_logps
|
||
rejected_logratios = rejected_logps - ref_rejected_logps
|
||
chosen_logratios_sorted, _ = torch.sort(chosen_logratios, dim=0)
|
||
rejected_logratios_sorted, _ = torch.sort(rejected_logratios, dim=0)
|
||
delta = chosen_logratios_sorted - rejected_logratios_sorted
|
||
losses = (
|
||
-F.logsigmoid(self.beta * delta) * (1 - self.label_smoothing)
|
||
- F.logsigmoid(-self.beta * delta) * self.label_smoothing
|
||
)
|
||
|
||
elif self.loss_type == "aot":
|
||
logratios = chosen_logps - rejected_logps
|
||
ref_logratios = ref_chosen_logps - ref_rejected_logps
|
||
logratios_sorted, _ = torch.sort(logratios, dim=0)
|
||
ref_logratios_sorted, _ = torch.sort(ref_logratios, dim=0)
|
||
delta = logratios_sorted - ref_logratios_sorted
|
||
losses = (
|
||
-F.logsigmoid(self.beta * delta) * (1 - self.label_smoothing)
|
||
- F.logsigmoid(-self.beta * delta) * self.label_smoothing
|
||
)
|
||
|
||
elif self.loss_type == "apo_zero":
|
||
# Eqn (7) of the APO paper (https://huggingface.co/papers/2408.06266)
|
||
# Use this loss when you believe the chosen outputs are better than your model's default output
|
||
losses_chosen = 1 - F.sigmoid(self.beta * chosen_logratios) # Increase chosen likelihood
|
||
losses_rejected = F.sigmoid(self.beta * rejected_logratios) # Decrease rejected likelihood
|
||
losses = losses_chosen + losses_rejected
|
||
|
||
elif self.loss_type == "apo_down":
|
||
# Eqn (8) of the APO paper (https://huggingface.co/papers/2408.06266)
|
||
# Use this loss when you believe the chosen outputs are worse than your model's default output.
|
||
# Decrease chosen likelihood and decrease rejected likelihood more
|
||
losses_chosen = F.sigmoid(self.beta * chosen_logratios)
|
||
losses_rejected = 1 - F.sigmoid(self.beta * (chosen_logratios - rejected_logratios))
|
||
losses = losses_chosen + losses_rejected
|
||
|
||
elif self.loss_type == "discopop":
|
||
# Eqn (5) of the DiscoPOP paper (https://huggingface.co/papers/2406.08414)
|
||
# This loss was discovered with LLM discovery
|
||
logratios = chosen_logps - rejected_logps
|
||
ref_logratios = ref_chosen_logps - ref_rejected_logps
|
||
logits = logratios - ref_logratios
|
||
logits = logits * self.beta
|
||
# Modulate the mixing coefficient based on the log ratio magnitudes
|
||
log_ratio_modulation = torch.sigmoid(logits / self.args.discopop_tau)
|
||
logistic_component = -F.logsigmoid(logits)
|
||
exp_component = torch.exp(-logits)
|
||
# Blend between logistic and exponential component based on log ratio modulation
|
||
losses = logistic_component * (1 - log_ratio_modulation) + exp_component * log_ratio_modulation
|
||
|
||
else:
|
||
raise ValueError(
|
||
f"Unknown loss type: {self.loss_type}. Should be one of ['sigmoid', 'hinge', 'ipo', 'exo_pair', "
|
||
"'nca_pair', 'robust', 'bco_pair', 'sppo_hard', 'aot', 'aot_pair', 'discopop', 'apo_zero', 'apo_down']"
|
||
)
|
||
|
||
chosen_rewards = self.beta * (chosen_logps.to(device) - ref_chosen_logps.to(device)).detach()
|
||
rejected_rewards = self.beta * (rejected_logps.to(device) - ref_rejected_logps.to(device)).detach()
|
||
|
||
return losses, chosen_rewards, rejected_rewards
|
||
|
||
def concatenated_forward(
|
||
self, model: nn.Module, batch: dict[str, Union[list, torch.LongTensor]], is_ref_model: bool = False
|
||
):
|
||
"""
|
||
Runs the given model on the given batch of inputs, concatenating the chosen and rejected inputs together.
|
||
|
||
We do this to avoid doing two forward passes, because it's faster for FSDP.
|
||
|
||
Args:
|
||
model:
|
||
Model to run the forward pass on.
|
||
batch:
|
||
Batch of input data.
|
||
is_ref_model:
|
||
Whether this method is being called for the reference model. If `True`, length desensitization is not
|
||
applied.
|
||
"""
|
||
num_examples = batch["prompt_input_ids"].shape[0]
|
||
|
||
concatenated_batch = self.concatenated_inputs(batch, padding_value=self.padding_value)
|
||
|
||
model_kwargs = {"use_cache": False}
|
||
if self.aux_loss_enabled:
|
||
model_kwargs["output_router_logits"] = True
|
||
|
||
# Add the pixel values and attention masks for vision models
|
||
if "pixel_values" in concatenated_batch:
|
||
model_kwargs["pixel_values"] = concatenated_batch["pixel_values"]
|
||
if "pixel_attention_mask" in concatenated_batch:
|
||
model_kwargs["pixel_attention_mask"] = concatenated_batch["pixel_attention_mask"]
|
||
if "image_sizes" in concatenated_batch:
|
||
model_kwargs["image_sizes"] = concatenated_batch["image_sizes"]
|
||
|
||
prompt_input_ids = concatenated_batch["prompt_input_ids"]
|
||
prompt_attention_mask = concatenated_batch["prompt_attention_mask"]
|
||
completion_input_ids = concatenated_batch["completion_input_ids"]
|
||
completion_attention_mask = concatenated_batch["completion_attention_mask"]
|
||
if self.is_encoder_decoder:
|
||
labels = completion_input_ids
|
||
labels[completion_attention_mask == 0] = self.label_pad_token_id
|
||
outputs = model(
|
||
input_ids=prompt_input_ids,
|
||
attention_mask=prompt_attention_mask,
|
||
labels=labels, # we need the labels for the logits to be returned
|
||
**model_kwargs,
|
||
)
|
||
logits = outputs.logits
|
||
loss_mask = completion_attention_mask.bool()
|
||
else:
|
||
# Concatenate the prompt and completion inputs
|
||
input_ids = torch.cat((prompt_input_ids, completion_input_ids), dim=1)
|
||
attention_mask = torch.cat((prompt_attention_mask, completion_attention_mask), dim=1)
|
||
# Mask the prompt but not the completion for the loss
|
||
loss_mask = torch.cat(
|
||
(torch.zeros_like(prompt_attention_mask), completion_attention_mask),
|
||
dim=1,
|
||
)
|
||
|
||
# Flush and truncate
|
||
if self.max_length is not None and self.max_length < attention_mask.size(1):
|
||
if self.truncation_mode == "keep_start":
|
||
# Flush left to reduce the memory usage
|
||
# [[0, 0, x, x, x, x], -> [[x, x, x, x],
|
||
# [0, x, x, x, 0, 0]] [x, x, x, 0]]
|
||
attention_mask, input_ids, loss_mask = flush_left(attention_mask, input_ids, loss_mask)
|
||
attention_mask = attention_mask[:, : self.max_length]
|
||
input_ids = input_ids[:, : self.max_length]
|
||
loss_mask = loss_mask[:, : self.max_length]
|
||
elif self.truncation_mode == "keep_end":
|
||
# Flush right before truncating left, then flush left
|
||
# [[0, 0, x, x, x, x], -> [[0, 0, x, x],
|
||
# [0, x, x, x, 0, 0]] [0, x, x, x]]
|
||
attention_mask, input_ids, loss_mask = flush_right(attention_mask, input_ids, loss_mask)
|
||
input_ids = input_ids[:, -self.max_length :]
|
||
attention_mask = attention_mask[:, -self.max_length :]
|
||
loss_mask = loss_mask[:, -self.max_length :]
|
||
attention_mask, input_ids, loss_mask = flush_left(attention_mask, input_ids, loss_mask)
|
||
else:
|
||
raise ValueError(
|
||
f"Unknown truncation mode: '{self.truncation_mode}'. Should be one of ['keep_end', "
|
||
"'keep_start']."
|
||
)
|
||
else:
|
||
# Flush left to reduce the memory usage
|
||
# [[0, 0, x, x, x, x], -> [[x, x, x, x],
|
||
# [0, x, x, x, 0, 0]] [x, x, x, 0]]
|
||
attention_mask, input_ids, loss_mask = flush_left(attention_mask, input_ids, loss_mask)
|
||
|
||
if self.use_logits_to_keep:
|
||
# Compute logits_to_keep based on loss_mask pattern:
|
||
# [[0, 0, 0, x, x, x, x],
|
||
# [0, 0, 0, x, x, x, 0]]
|
||
# ^ start computing logits from here ([:, -(7-3+1):])
|
||
first_compute_index = loss_mask.nonzero(as_tuple=True)[1].min()
|
||
logits_to_keep = (loss_mask.shape[1] - first_compute_index).item() + 1 # +1 for the first label
|
||
model_kwargs["logits_to_keep"] = logits_to_keep
|
||
|
||
if self.padding_free:
|
||
# Flatten the input_ids, position_ids, and loss_mask
|
||
# input_ids = [[a, b, c, 0], -> input_ids = [[a, b, c, d, e, f, g]]
|
||
# [d, e, f, g]] position_ids = [[0, 1, 2, 0, 1, 2, 3]]
|
||
input_ids = input_ids[attention_mask.bool()].unsqueeze(0)
|
||
loss_mask = loss_mask[attention_mask.bool()].unsqueeze(0)
|
||
position_ids = attention_mask.cumsum(1)[attention_mask.bool()].unsqueeze(0) - 1
|
||
model_kwargs["position_ids"] = position_ids
|
||
else:
|
||
model_kwargs["attention_mask"] = attention_mask
|
||
|
||
outputs = model(input_ids, **model_kwargs)
|
||
logits = outputs.logits
|
||
|
||
# Offset the logits by one to align with the labels
|
||
labels = torch.roll(input_ids, shifts=-1, dims=1)
|
||
loss_mask = torch.roll(loss_mask, shifts=-1, dims=1).bool()
|
||
|
||
if self.use_logits_to_keep:
|
||
# Align labels with logits
|
||
# logits: -, -, [x2, x3, x4, x5, x6]
|
||
# ^ --------- ^ after logits[:, :-1, :]
|
||
# labels: [y0, y1, y2, y3, y4, y5, y6]
|
||
# ^ --------- ^ with logits_to_keep=4, [:, -4:]
|
||
# loss_mask: [0, 0, 0, 1, 1, 1, 1]
|
||
labels = labels[:, -logits_to_keep:]
|
||
loss_mask = loss_mask[:, -logits_to_keep:]
|
||
|
||
if logits.shape[:2] != labels.shape[:2]:
|
||
# for llava, the returned logits include the image tokens (placed before the text tokens)
|
||
seq_len = labels.shape[1]
|
||
logits = logits[:, -seq_len:]
|
||
|
||
# Compute the log probabilities of the labels
|
||
labels[~loss_mask] = 0 # dummy token; we'll ignore the losses on these tokens later
|
||
per_token_logps = selective_log_softmax(logits, labels)
|
||
per_token_logps[~loss_mask] = 0
|
||
per_token_logps = torch.roll(per_token_logps, shifts=1, dims=1)
|
||
|
||
if self.padding_free:
|
||
# Unflatten the per_token_logps (shape: [1, sum_seq_len] -> [batch_size, seq_len])
|
||
batch_size, seq_len = attention_mask.shape
|
||
per_token_logps_ = torch.zeros(
|
||
batch_size, seq_len, device=outputs.logits.device, dtype=outputs.logits.dtype
|
||
)
|
||
per_token_logps_[attention_mask.bool()] = per_token_logps
|
||
per_token_logps = per_token_logps_
|
||
|
||
all_logps = per_token_logps.sum(-1)
|
||
|
||
output = {}
|
||
|
||
if self.use_weighting:
|
||
with torch.no_grad():
|
||
# Eq (2) of the WPO paper: https://huggingface.co/papers/2406.11827
|
||
logprobs = F.log_softmax(logits, dim=-1)
|
||
weights_adjustment_factor = torch.logsumexp(2 * logprobs, dim=-1) # same as sum(probs**2) in log space
|
||
per_token_logps_adjusted = per_token_logps - weights_adjustment_factor
|
||
all_weights = (per_token_logps_adjusted * loss_mask).sum(-1) / loss_mask.sum(-1)
|
||
chosen_weights = all_weights[:num_examples]
|
||
rejected_weights = all_weights[num_examples:]
|
||
output["policy_weights"] = torch.clamp(torch.exp(chosen_weights + rejected_weights), max=1)
|
||
|
||
if self.args.rpo_alpha is not None:
|
||
# Only use the chosen logits for the RPO loss
|
||
chosen_logits = logits[:num_examples]
|
||
chosen_labels = labels[:num_examples]
|
||
|
||
# Compute the log probabilities of the labels
|
||
output["nll_loss"] = F.cross_entropy(
|
||
torch.flatten(chosen_logits, end_dim=1), torch.flatten(chosen_labels, end_dim=1), ignore_index=0
|
||
)
|
||
|
||
if self.loss_type == "ipo":
|
||
all_logps = all_logps / loss_mask.sum(-1)
|
||
|
||
if self.args.ld_alpha is not None and not is_ref_model:
|
||
# Compute response lengths based on loss_mask
|
||
completion_lengths = loss_mask.sum(dim=1)
|
||
|
||
chosen_lengths = completion_lengths[:num_examples]
|
||
rejected_lengths = completion_lengths[num_examples:]
|
||
public_lengths = torch.min(chosen_lengths, rejected_lengths) # l_p in the paper
|
||
public_lengths = torch.cat([public_lengths, public_lengths], dim=0)
|
||
|
||
seq_len = per_token_logps.size(1)
|
||
position_ids = torch.arange(seq_len, device=per_token_logps.device).expand_as(per_token_logps)
|
||
|
||
ld_mask = position_ids < public_lengths.unsqueeze(1)
|
||
mask = position_ids < completion_lengths.unsqueeze(1)
|
||
|
||
front_mask = (ld_mask & mask).float()
|
||
rear_mask = (~ld_mask & mask).float()
|
||
front_logps = (per_token_logps * front_mask).sum(dim=1)
|
||
rear_logps = (per_token_logps * rear_mask).sum(dim=1)
|
||
|
||
all_logps = front_logps + self.args.ld_alpha * rear_logps
|
||
|
||
output["chosen_logps"] = all_logps[:num_examples]
|
||
output["rejected_logps"] = all_logps[num_examples:]
|
||
|
||
# Compute the mean logits
|
||
if self.padding_free:
|
||
# position_ids contains a sequence of range identifiers (e.g., [[0, 1, 2, 0, 1, 2, 3, ...]]).
|
||
# There are 2*num_examples ranges in total: the first half corresponds to the chosen tokens,
|
||
# and the second half to the rejected tokens.
|
||
# To find the start of the rejected tokens, we look for the num_examples+1-th zero in pos_id.
|
||
split_idx = (position_ids == 0).nonzero(as_tuple=True)[1][num_examples]
|
||
mean_chosen_logits = logits[0, :split_idx][loss_mask[0, :split_idx]].mean()
|
||
mean_rejected_logits = logits[0, split_idx:][loss_mask[0, split_idx:]].mean()
|
||
else:
|
||
mean_chosen_logits = logits[:num_examples][loss_mask[:num_examples]].mean()
|
||
mean_rejected_logits = logits[num_examples:][loss_mask[num_examples:]].mean()
|
||
|
||
output["mean_chosen_logits"] = mean_chosen_logits
|
||
output["mean_rejected_logits"] = mean_rejected_logits
|
||
|
||
if self.aux_loss_enabled:
|
||
output["aux_loss"] = outputs.aux_loss
|
||
|
||
return output
|
||
|
||
def get_batch_loss_metrics(
|
||
self,
|
||
model,
|
||
batch: dict[str, Union[list, torch.LongTensor]],
|
||
train_eval: Literal["train", "eval"] = "train",
|
||
):
|
||
"""Compute the DPO loss and other metrics for the given batch of inputs for train or test."""
|
||
metrics = {}
|
||
|
||
model_output = self.concatenated_forward(model, batch)
|
||
|
||
# if ref_chosen_logps and ref_rejected_logps in batch use them, otherwise use the reference model
|
||
if "ref_chosen_logps" in batch and "ref_rejected_logps" in batch:
|
||
ref_chosen_logps = batch["ref_chosen_logps"]
|
||
ref_rejected_logps = batch["ref_rejected_logps"]
|
||
else:
|
||
ref_chosen_logps, ref_rejected_logps = self.compute_ref_log_probs(batch)
|
||
|
||
losses, chosen_rewards, rejected_rewards = self.dpo_loss(
|
||
model_output["chosen_logps"], model_output["rejected_logps"], ref_chosen_logps, ref_rejected_logps
|
||
)
|
||
reward_accuracies = (chosen_rewards > rejected_rewards).float()
|
||
|
||
if self.args.rpo_alpha is not None:
|
||
losses = losses + self.args.rpo_alpha * model_output["nll_loss"] # RPO loss from V3 of the paper
|
||
|
||
if self.use_weighting:
|
||
losses = losses * model_output["policy_weights"]
|
||
|
||
if self.aux_loss_enabled:
|
||
losses = losses + self.aux_loss_coef * model_output["aux_loss"]
|
||
|
||
prefix = "eval_" if train_eval == "eval" else ""
|
||
metrics[f"{prefix}rewards/chosen"] = self.accelerator.gather_for_metrics(chosen_rewards).mean().item()
|
||
metrics[f"{prefix}rewards/rejected"] = self.accelerator.gather_for_metrics(rejected_rewards).mean().item()
|
||
metrics[f"{prefix}rewards/accuracies"] = self.accelerator.gather_for_metrics(reward_accuracies).mean().item()
|
||
metrics[f"{prefix}rewards/margins"] = (
|
||
self.accelerator.gather_for_metrics(chosen_rewards - rejected_rewards).mean().item()
|
||
)
|
||
metrics[f"{prefix}logps/chosen"] = (
|
||
self.accelerator.gather_for_metrics(model_output["chosen_logps"]).detach().mean().item()
|
||
)
|
||
metrics[f"{prefix}logps/rejected"] = (
|
||
self.accelerator.gather_for_metrics(model_output["rejected_logps"]).detach().mean().item()
|
||
)
|
||
metrics[f"{prefix}logits/chosen"] = (
|
||
self.accelerator.gather_for_metrics(model_output["mean_chosen_logits"]).detach().mean().item()
|
||
)
|
||
metrics[f"{prefix}logits/rejected"] = (
|
||
self.accelerator.gather_for_metrics(model_output["mean_rejected_logits"]).detach().mean().item()
|
||
)
|
||
if self.args.rpo_alpha is not None:
|
||
metrics[f"{prefix}nll_loss"] = (
|
||
self.accelerator.gather_for_metrics(model_output["nll_loss"]).detach().mean().item()
|
||
)
|
||
if self.aux_loss_enabled:
|
||
metrics[f"{prefix}aux_loss"] = (
|
||
self.accelerator.gather_for_metrics(model_output["aux_loss"]).detach().mean().item()
|
||
)
|
||
|
||
return losses.mean(), metrics
|
||
|
||
def compute_loss(
|
||
self,
|
||
model: Union[PreTrainedModel, nn.Module],
|
||
inputs: dict[str, Union[torch.Tensor, Any]],
|
||
return_outputs=False,
|
||
num_items_in_batch=None,
|
||
) -> Union[torch.Tensor, tuple[torch.Tensor, dict[str, torch.Tensor]]]:
|
||
device_type = "xpu" if is_torch_xpu_available() else "cuda"
|
||
compute_loss_context_manager = (
|
||
amp.autocast(device_type) if self._peft_has_been_casted_to_bf16 else nullcontext()
|
||
)
|
||
with compute_loss_context_manager:
|
||
loss, metrics = self.get_batch_loss_metrics(model, inputs, train_eval="train")
|
||
|
||
# Make sure to move the loss to the device the original accumulating loss is at back in the `Trainer` class:
|
||
loss = loss.to(self.args.device)
|
||
# force log the metrics
|
||
self.store_metrics(metrics, train_eval="train")
|
||
|
||
if return_outputs:
|
||
return loss, metrics
|
||
|
||
return loss
|
||
|
||
def generate_from_model_and_ref(self, model, batch: dict[str, torch.LongTensor]) -> tuple[str, str]:
|
||
"""Generate samples from the model and reference model for the given batch of inputs."""
|
||
|
||
# If one uses `generate_during_eval` with peft + bf16, we need to explicitly call generate with
|
||
# the torch amp context manager as some hidden states are silently casted to full precision.
|
||
device_type = "xpu" if is_torch_xpu_available() else "cuda"
|
||
generate_context_manager = amp.autocast(device_type) if self._peft_has_been_casted_to_bf16 else nullcontext()
|
||
|
||
with generate_context_manager:
|
||
policy_output = model.generate(
|
||
input_ids=batch["prompt_input_ids"],
|
||
attention_mask=batch["prompt_attention_mask"],
|
||
max_length=self.max_length,
|
||
do_sample=True,
|
||
pad_token_id=self.padding_value,
|
||
)
|
||
|
||
# if ref_output in batch use that otherwise use the reference model
|
||
if "ref_output" in batch:
|
||
ref_output = batch["ref_output"]
|
||
else:
|
||
if self.ref_model is None:
|
||
with self.null_ref_context():
|
||
ref_output = self.model.generate(
|
||
input_ids=batch["prompt_input_ids"],
|
||
attention_mask=batch["prompt_attention_mask"],
|
||
max_length=self.max_length,
|
||
do_sample=True,
|
||
pad_token_id=self.padding_value,
|
||
)
|
||
else:
|
||
ref_output = self.ref_model.generate(
|
||
input_ids=batch["prompt_input_ids"],
|
||
attention_mask=batch["prompt_attention_mask"],
|
||
max_length=self.max_length,
|
||
do_sample=True,
|
||
pad_token_id=self.padding_value,
|
||
)
|
||
|
||
policy_output = pad_to_length(policy_output, self.max_length, self.padding_value)
|
||
policy_output_decoded = self.processing_class.batch_decode(policy_output, skip_special_tokens=True)
|
||
|
||
ref_output = pad_to_length(ref_output, self.max_length, self.padding_value)
|
||
ref_output_decoded = self.processing_class.batch_decode(ref_output, skip_special_tokens=True)
|
||
|
||
return policy_output_decoded, ref_output_decoded
|
||
|
||
def prediction_step(
|
||
self,
|
||
model: Union[PreTrainedModel, nn.Module],
|
||
inputs: dict[str, Union[torch.Tensor, Any]],
|
||
prediction_loss_only: bool,
|
||
ignore_keys: Optional[list[str]] = None,
|
||
):
|
||
if ignore_keys is None:
|
||
if hasattr(model, "config"):
|
||
ignore_keys = getattr(model.config, "keys_to_ignore_at_inference", [])
|
||
else:
|
||
ignore_keys = []
|
||
|
||
device_type = "xpu" if is_torch_xpu_available() else "cuda"
|
||
prediction_context_manager = amp.autocast(device_type) if self._peft_has_been_casted_to_bf16 else nullcontext()
|
||
|
||
with torch.no_grad(), prediction_context_manager:
|
||
loss, metrics = self.get_batch_loss_metrics(model, inputs, train_eval="eval")
|
||
|
||
# force log the metrics
|
||
self.store_metrics(metrics, train_eval="eval")
|
||
|
||
if prediction_loss_only:
|
||
return loss.detach(), None, None
|
||
|
||
# logits for the chosen and rejected samples from model
|
||
logits_dict = {
|
||
"eval_logits/chosen": metrics["eval_logits/chosen"],
|
||
"eval_logits/rejected": metrics["eval_logits/rejected"],
|
||
}
|
||
logits = [v for k, v in logits_dict.items() if k not in ignore_keys]
|
||
logits = torch.tensor(logits, device=self.accelerator.device)
|
||
labels = torch.zeros(logits.shape[0], device=self.accelerator.device)
|
||
|
||
return (loss.detach(), logits, labels)
|
||
|
||
def store_metrics(self, metrics: dict[str, float], train_eval: Literal["train", "eval"] = "train") -> None:
|
||
for key, value in metrics.items():
|
||
self._stored_metrics[train_eval][key].append(value)
|
||
|
||
def evaluation_loop(
|
||
self,
|
||
dataloader: DataLoader,
|
||
description: str,
|
||
prediction_loss_only: Optional[bool] = None,
|
||
ignore_keys: Optional[list[str]] = None,
|
||
metric_key_prefix: str = "eval",
|
||
) -> EvalLoopOutput:
|
||
"""
|
||
Overriding built-in evaluation loop to store metrics for each batch.
|
||
Prediction/evaluation loop, shared by `Trainer.evaluate()` and `Trainer.predict()`.
|
||
|
||
Works both with or without labels.
|
||
"""
|
||
|
||
# Sample and save to game log if requested (for one batch to save time)
|
||
if self.generate_during_eval:
|
||
# Generate random indices within the range of the total number of samples
|
||
num_samples = len(dataloader.dataset)
|
||
random_indices = random.sample(range(num_samples), k=self.args.eval_batch_size)
|
||
|
||
# Use dataloader.dataset.select to get the random batch without iterating over the DataLoader
|
||
random_batch_dataset = dataloader.dataset.select(random_indices)
|
||
random_batch = self.data_collator(random_batch_dataset)
|
||
random_batch = self._prepare_inputs(random_batch)
|
||
|
||
policy_output_decoded, ref_output_decoded = self.generate_from_model_and_ref(self.model, random_batch)
|
||
|
||
table = pd.DataFrame(
|
||
columns=["Prompt", "Policy", "Ref Model"],
|
||
data=[
|
||
[prompt, pol[len(prompt) :], ref[len(prompt) :]]
|
||
for prompt, pol, ref in zip(
|
||
random_batch_dataset["prompt"], policy_output_decoded, ref_output_decoded
|
||
)
|
||
],
|
||
)
|
||
if "wandb" in self.args.report_to and self.accelerator.is_main_process:
|
||
wandb.log({"game_log": wandb.Table(data=table)})
|
||
|
||
if "comet_ml" in self.args.report_to:
|
||
log_table_to_comet_experiment(
|
||
name="game_log.csv",
|
||
table=table,
|
||
)
|
||
|
||
# Base evaluation
|
||
initial_output = super().evaluation_loop(
|
||
dataloader, description, prediction_loss_only, ignore_keys, metric_key_prefix
|
||
)
|
||
|
||
return initial_output
|
||
|
||
def log(self, logs: dict[str, float], start_time: Optional[float] = None) -> None:
|
||
"""
|
||
Log `logs` on the various objects watching training, including stored metrics.
|
||
|
||
Args:
|
||
logs (`dict[str, float]`):
|
||
The values to log.
|
||
start_time (`float` or `None`, *optional*, defaults to `None`):
|
||
Start time of the training.
|
||
"""
|
||
# logs either has 'loss' or 'eval_loss'
|
||
train_eval = "train" if "loss" in logs else "eval"
|
||
# Add averaged stored metrics to logs
|
||
for key, metrics in self._stored_metrics[train_eval].items():
|
||
logs[key] = torch.tensor(metrics).mean().item()
|
||
del self._stored_metrics[train_eval]
|
||
|
||
if version.parse(transformers.__version__) >= version.parse("4.47.0.dev0"):
|
||
return super().log(logs, start_time)
|
||
else: # transformers<=4.46
|
||
return super().log(logs)
|
||
|
||
def create_model_card(
|
||
self,
|
||
model_name: Optional[str] = None,
|
||
dataset_name: Optional[str] = None,
|
||
tags: Union[str, list[str], None] = None,
|
||
):
|
||
"""
|
||
Creates a draft of a model card using the information available to the `Trainer`.
|
||
|
||
Args:
|
||
model_name (`str` or `None`, *optional*, defaults to `None`):
|
||
Name of the model.
|
||
dataset_name (`str` or `None`, *optional*, defaults to `None`):
|
||
Name of the dataset used for training.
|
||
tags (`str`, `list[str]` or `None`, *optional*, defaults to `None`):
|
||
Tags to be associated with the model card.
|
||
"""
|
||
if not self.is_world_process_zero():
|
||
return
|
||
|
||
if hasattr(self.model.config, "_name_or_path") and not os.path.isdir(self.model.config._name_or_path):
|
||
base_model = self.model.config._name_or_path
|
||
else:
|
||
base_model = None
|
||
|
||
tags = tags or []
|
||
if isinstance(tags, str):
|
||
tags = [tags]
|
||
|
||
if hasattr(self.model.config, "unsloth_version"):
|
||
tags.append("unsloth")
|
||
|
||
citation = textwrap.dedent(
|
||
"""\
|
||
@inproceedings{rafailov2023direct,
|
||
title = {{Direct Preference Optimization: Your Language Model is Secretly a Reward Model}},
|
||
author = {Rafael Rafailov and Archit Sharma and Eric Mitchell and Christopher D. Manning and Stefano Ermon and Chelsea Finn},
|
||
year = 2023,
|
||
booktitle = {Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 - 16, 2023},
|
||
url = {http://papers.nips.cc/paper_files/paper/2023/hash/a85b405ed65c6477a4fe8302b5e06ce7-Abstract-Conference.html},
|
||
editor = {Alice Oh and Tristan Naumann and Amir Globerson and Kate Saenko and Moritz Hardt and Sergey Levine},
|
||
}"""
|
||
)
|
||
|
||
model_card = generate_model_card(
|
||
base_model=base_model,
|
||
model_name=model_name,
|
||
hub_model_id=self.hub_model_id,
|
||
dataset_name=dataset_name,
|
||
tags=tags,
|
||
wandb_url=wandb.run.get_url() if is_wandb_available() and wandb.run is not None else None,
|
||
comet_url=get_comet_experiment_url(),
|
||
trainer_name="DPO",
|
||
trainer_citation=citation,
|
||
paper_title="Direct Preference Optimization: Your Language Model is Secretly a Reward Model",
|
||
paper_id="2305.18290",
|
||
)
|
||
|
||
model_card.save(os.path.join(self.args.output_dir, "README.md"))
|
||
class UnslothDPOTrainer(_UnslothDPOTrainer):
|
||
"""
|
||
|
||
Initialize DPOTrainer.
|
||
|
||
Args:
|
||
model (`transformers.PreTrainedModel`):
|
||
The model to train, preferably an `AutoModelForSequenceClassification`.
|
||
ref_model (`PreTrainedModelWrapper`):
|
||
Hugging Face transformer model with a casual language modelling head. Used for implicit reward computation and loss. If no
|
||
reference model is provided, the trainer will create a reference model with the same architecture as the model to be optimized.
|
||
args (`DPOConfig`):
|
||
The DPO config arguments to use for training.
|
||
data_collator (`transformers.DataCollator`):
|
||
The data collator to use for training. If None is specified, the default data collator (`DataCollatorForPreference`) will be used
|
||
which will pad the sequences to the maximum length of the sequences in the batch, given a dataset of paired sequences.
|
||
train_dataset (`datasets.Dataset`):
|
||
The dataset to use for training.
|
||
eval_dataset (`datasets.Dataset`):
|
||
The dataset to use for evaluation.
|
||
processing_class (`PreTrainedTokenizerBase` or `BaseImageProcessor` or `FeatureExtractionMixin` or `ProcessorMixin`, *optional*):
|
||
Processing class used to process the data. If provided, will be used to automatically process the inputs
|
||
for the model, and it will be saved along the model to make it easier to rerun an interrupted training or
|
||
reuse the fine-tuned model.
|
||
model_init (`Callable[[], transformers.PreTrainedModel]`):
|
||
The model initializer to use for training. If None is specified, the default model initializer will be used.
|
||
compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*):
|
||
The function to use to compute the metrics. Must take a `EvalPrediction` and return
|
||
a dictionary string to metric values.
|
||
callbacks (`list[transformers.TrainerCallback]`):
|
||
The callbacks to use for training.
|
||
optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`):
|
||
The optimizer and scheduler to use for training.
|
||
preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`):
|
||
The function to use to preprocess the logits before computing the metrics.
|
||
peft_config (`dict`, defaults to `None`):
|
||
The PEFT configuration to use for training. If you pass a PEFT configuration, the model will be wrapped in a PEFT model.
|
||
|
||
"""
|
||
def __init__(
|
||
self,
|
||
model = None,
|
||
ref_model = None,
|
||
args = None,
|
||
data_collator = None,
|
||
train_dataset = None,
|
||
eval_dataset = None,
|
||
processing_class = None,
|
||
model_init = None,
|
||
compute_metrics = None,
|
||
callbacks = None,
|
||
preprocess_logits_for_metrics = None,
|
||
peft_config = None,
|
||
**kwargs
|
||
):
|
||
if args is None: args = UnslothDPOConfig()
|
||
use_bf16 = getattr(args, 'bf16', False)
|
||
use_fp16 = getattr(args, 'fp16', False)
|
||
force_float32 = False
|
||
if os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1':
|
||
print('Unsloth: Switching to float32 training since model cannot work with float16')
|
||
force_float32 = True
|
||
mixed_precision_dtype = os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32')
|
||
dtype = getattr(model.config, 'torch_dtype', None)
|
||
if dtype is None: dtype = model.get_input_embeddings().dtype
|
||
from unsloth_zoo.utils import _get_dtype
|
||
dtype = _get_dtype(dtype)
|
||
float16 = dtype == torch.float16
|
||
if not force_float32 and (float16 and use_bf16): raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`')
|
||
if not force_float32 and (not float16 and use_fp16): raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`')
|
||
if force_float32:
|
||
args.fp16 = False
|
||
args.bf16 = False
|
||
os.environ['ACCELERATE_MIXED_PRECISION'] = 'no'
|
||
elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32':
|
||
args.fp16 = float16
|
||
args.bf16 = not float16
|
||
os.environ['ACCELERATE_MIXED_PRECISION'] = 'fp16' if float16 else 'bf16'
|
||
if getattr(args, 'eval_dataset', None) is not None and getattr(args, 'eval_strategy', 'no') == 'no':
|
||
args.eval_strategy = 'steps'
|
||
if getattr(args, 'eval_steps', None) is None: args.eval_steps = 0.1
|
||
ga_steps = getattr(args, 'gradient_accumulation_steps', None)
|
||
if ga_steps is not None and ga_steps > 1:
|
||
from transformers import __version__ as transformers_version
|
||
if Version(transformers_version) <= Version('4.45.2'):
|
||
print('**** Unsloth: Please use our fixed gradient_accumulation_steps by updating transformers, TRL and Unsloth!\n'
|
||
'`pip install --upgrade --no-cache-dir --force-reinstall --no-deps unsloth transformers trl unsloth_zoo`')
|
||
if getattr(args, 'eval_strategy', 'no') != 'no':
|
||
eval_bsz = getattr(args, 'per_device_eval_batch_size', 8)
|
||
if eval_bsz == 8 and args.per_device_train_batch_size < eval_bsz: args.per_device_eval_batch_size = args.per_device_train_batch_size
|
||
if getattr(args, 'eval_accumulation_steps', None) is None and ga_steps is not None: args.eval_accumulation_steps = ga_steps
|
||
fp16_full_eval = getattr(args, 'fp16_full_eval', False)
|
||
bf16_full_eval = getattr(args, 'bf16_full_eval', False)
|
||
if args.fp16 and bf16_full_eval: args.bf16_full_eval = False; args.fp16_full_eval = True
|
||
if args.bf16 and fp16_full_eval: args.bf16_full_eval = True; args.fp16_full_eval = False
|
||
if force_float32:
|
||
args.bf16_full_eval = False
|
||
args.fp16_full_eval = False
|
||
elif os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') == 'bfloat16':
|
||
args.bf16_full_eval = True
|
||
args.fp16_full_eval = False
|
||
elif not bf16_full_eval and not fp16_full_eval:
|
||
args.bf16_full_eval = args.bf16
|
||
args.fp16_full_eval = args.fp16
|
||
_output_logits = False
|
||
if locals().get('compute_metrics', None) is not None: _output_logits = True
|
||
if locals().get('preprocess_logits_for_metrics', None) is not None: _output_logits = True
|
||
if _output_logits:
|
||
os.environ['UNSLOTH_RETURN_LOGITS'] = '1'
|
||
if 'max_seq_length' not in locals() and not hasattr(args, 'max_seq_length'):
|
||
pass
|
||
else:
|
||
model_max_seq_length = getattr(model, 'max_seq_length', None)
|
||
args_max_seq_length = getattr(args, 'max_seq_length', None)
|
||
if args_max_seq_length is None and model_max_seq_length is not None:
|
||
max_seq_length = model.max_seq_length
|
||
if hasattr(args, 'max_seq_length'): args.max_seq_length = max_seq_length
|
||
if model is not None and hasattr(model, 'for_training'):
|
||
model.for_training()
|
||
if 'tokenizer' in locals() and hasattr(tokenizer, 'padding_side'): tokenizer.padding_side = 'right'
|
||
if 'processing_class' in locals():
|
||
if hasattr(processing_class, 'padding_side'): processing_class.padding_side = 'right'
|
||
if hasattr(processing_class, 'tokenizer') and hasattr(processing_class.tokenizer, 'padding_side'): processing_class.tokenizer.padding_side = 'right'
|
||
__tokenizer = processing_class if 'processing_class' in locals() else tokenizer
|
||
from unsloth_zoo.vision_utils import UnslothVisionDataCollator
|
||
if not isinstance(data_collator, UnslothVisionDataCollator):
|
||
if isinstance(data_collator, DataCollatorForSeq2Seq) and 'labels' not in train_dataset.column_names:
|
||
data_collator = TransformersDataCollatorForLanguageModeling(__tokenizer, mlm = False, mlm_probability = 0.0)
|
||
elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling) and 'labels' in train_dataset.column_names:
|
||
data_collator = DataCollatorForSeq2Seq(__tokenizer)
|
||
else:
|
||
if hasattr(args, 'remove_unused_columns'): args.remove_unused_columns = False
|
||
if hasattr(args, 'dataset_text_field'): args.dataset_text_field = ''
|
||
if hasattr(args, 'dataset_kwargs'): args.dataset_kwargs = {'skip_prepare_dataset': True}
|
||
if not isinstance(data_collator, UnslothVisionDataCollator):
|
||
if not hasattr(__tokenizer, 'pad') and hasattr(__tokenizer, 'tokenizer'):
|
||
if isinstance(data_collator, DataCollatorForSeq2Seq):
|
||
data_collator = DataCollatorForSeq2Seq(__tokenizer.tokenizer)
|
||
else:
|
||
data_collator = TransformersDataCollatorForLanguageModeling(__tokenizer.tokenizer, mlm = False, mlm_probability = 0.0)
|
||
other_metrics = []
|
||
|
||
from unsloth_zoo.logging_utils import PatchRLStatistics
|
||
PatchRLStatistics('dpo_trainer', other_metrics)
|
||
if hasattr(train_dataset, 'column_names'):
|
||
column_names = set(train_dataset.column_names)
|
||
check = ['chosen', 'rejected', 'prompt', 'chosen_input_ids', 'chosen_attention_mask',
|
||
'chosen_labels', 'rejected_input_ids', 'rejected_attention_mask', 'rejected_labels',
|
||
'prompt_input_ids', 'prompt_attention_mask']
|
||
if all(x in column_names for x in check):
|
||
train_dataset = train_dataset.remove_columns(['chosen', 'rejected', 'prompt'])
|
||
del check, column_names
|
||
|
||
super().__init__(
|
||
model = model,
|
||
ref_model = ref_model,
|
||
args = args,
|
||
data_collator = data_collator,
|
||
train_dataset = train_dataset,
|
||
eval_dataset = eval_dataset,
|
||
processing_class = processing_class,
|
||
model_init = model_init,
|
||
compute_metrics = compute_metrics,
|
||
callbacks = callbacks,
|
||
preprocess_logits_for_metrics = preprocess_logits_for_metrics,
|
||
peft_config = peft_config,**kwargs)
|
||
if hasattr(self, 'neftune_hook_handle'):
|
||
self.neftune_hook_handle.remove()
|
||
if hasattr(self, 'neftune_hook_handle'): del self.neftune_hook_handle
|
||
if getattr(args, 'neftune_noise_alpha', None) is not None:
|
||
model.get_input_embeddings().neftune_noise_alpha = self.neftune_noise_alpha
|
||
pass
|
||
|
||
pass
|