1365 lines
60 KiB
Python
1365 lines
60 KiB
Python
"""
|
|
2026.2.1
|
|
2026.2.1
|
|
4.57.6
|
|
0.24.0
|
|
__UNSLOTH_VERSIONING__
|
|
"""
|
|
|
|
# Unsloth auto generated code
|
|
# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
|
|
#
|
|
# This program is free software: you can redistribute it and/or modify
|
|
# it under the terms of the GNU Lesser General Public License as published by
|
|
# the Free Software Foundation, either version 3 of the License, or
|
|
# (at your option) any later version.
|
|
#
|
|
# This program is distributed in the hope that it will be useful,
|
|
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
# GNU General Public License for more details.
|
|
#
|
|
# You should have received a copy of the GNU Lesser General Public License
|
|
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
|
|
|
from torch import Tensor
|
|
import torch
|
|
import torch.nn as nn
|
|
from torch.nn import functional as F
|
|
from unsloth_zoo.temporary_patches.common import torch_compile
|
|
from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable
|
|
from trl.trainer.nash_md_trainer import (Any, BaseImageProcessor, BasePairwiseJudge, Callable, Dataset, EvalPrediction, F, FeatureExtractionMixin, GeometricMixtureWrapper, IterableDataset, NashMDConfig, NashMDTrainer, OnlineDPOTrainer, OptimizerNames, Optional, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, SIMPLE_CHAT_TEMPLATE, TrainerCallback, Union, empty_cache, get_reward, is_conversational, is_peft_available, jinja2, maybe_apply_chat_template, nn, selective_log_softmax, textwrap, torch, truncate_right, unwrap_model_for_generation)
|
|
|
|
|
|
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
|
|
import inspect
|
|
from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling as TransformersDataCollatorForLanguageModeling
|
|
from transformers.training_args import ParallelMode
|
|
from unsloth_zoo.device_type import DEVICE_TYPE, device_synchronize
|
|
|
|
# Wrap trainer with padding to right and enable training mode
|
|
# Also patches W&B since multiple runs must use wandb.finish()
|
|
import functools
|
|
from types import MethodType
|
|
try:
|
|
from unsloth_zoo.gradient_checkpointing import reset_unsloth_gradient_checkpointing_buffers
|
|
except:
|
|
def reset_unsloth_gradient_checkpointing_buffers(): pass
|
|
def prepare_for_training_mode(f):
|
|
@functools.wraps(f)
|
|
def wrapper(self, *args, **kwargs):
|
|
# Enable training mode
|
|
_was_training = None
|
|
# Get gradient checkpointing setting from training arguments
|
|
use_gc = getattr(self.args, 'gradient_checkpointing', True)
|
|
if hasattr(self, 'model') and hasattr(self.model, "training"):
|
|
_was_training = self.model.training
|
|
if hasattr(self, 'model') and hasattr(self.model, "for_training"):
|
|
self.model.for_training(use_gradient_checkpointing=use_gc)
|
|
output = f(self, *args, **kwargs)
|
|
# Restore previous mode when possible
|
|
if hasattr(self, 'model') and hasattr(self.model, "for_inference"):
|
|
if _was_training is False:
|
|
self.model.for_inference()
|
|
elif _was_training is True and hasattr(self.model, "for_training"):
|
|
self.model.for_training(use_gradient_checkpointing=use_gc)
|
|
# Reset gradient checkpointing buffers to free memory while staying ready for next run
|
|
try:
|
|
reset_unsloth_gradient_checkpointing_buffers()
|
|
except:
|
|
pass
|
|
# Patch W&B to enable logging on future runs, otherwise it'll overwrite the first run
|
|
try:
|
|
import wandb
|
|
wandb.finish()
|
|
except:
|
|
pass
|
|
return output
|
|
return wrapper
|
|
pass
|
|
|
|
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 chunked_hidden_states_selective_log_softmax(
|
|
hidden_states: torch.Tensor,
|
|
lm_head: torch.Tensor,
|
|
index: torch.Tensor,
|
|
chunks: int = 4,
|
|
logit_scale_multiply: float = 0.0,
|
|
logit_scale_divide: float = 0.0,
|
|
logit_softcapping: float = 0.0,
|
|
temperature: float = 1.0,
|
|
) -> torch.Tensor:
|
|
# All Unsloth Zoo code licensed under AGPL3
|
|
flat_hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1])
|
|
flat_index = index.reshape(-1)
|
|
|
|
chunked_hidden_states = torch.chunk(flat_hidden_states, chunks=chunks, dim=0)
|
|
chunked_index = torch.chunk(flat_index, chunks=chunks, dim=0)
|
|
|
|
all_per_token_logps = []
|
|
|
|
for chunk_hidden_states, chunk_index in zip(chunked_hidden_states, chunked_index):
|
|
chunk_logits = chunk_hidden_states.to(lm_head.dtype) @ lm_head.t()
|
|
|
|
if logit_scale_multiply != 0.0:
|
|
chunk_logits = chunk_logits * logit_scale_multiply
|
|
if logit_scale_divide != 0.0:
|
|
chunk_logits = chunk_logits / logit_scale_divide
|
|
if logit_softcapping != 0.0:
|
|
chunk_logits = chunk_logits * torch.tanh(chunk_logits / logit_softcapping)
|
|
|
|
chunk_logits = chunk_logits.to(torch.float32)
|
|
|
|
if temperature != 1.0:
|
|
chunk_logits = chunk_logits / temperature
|
|
|
|
selected_logits = torch.gather(chunk_logits, dim=-1, index=chunk_index.unsqueeze(-1)).squeeze(-1)
|
|
logsumexp_values = torch.logsumexp(chunk_logits, dim=-1)
|
|
per_token_logps = selected_logits - logsumexp_values
|
|
all_per_token_logps.append(per_token_logps)
|
|
|
|
all_per_token_logps = torch.concat(all_per_token_logps)
|
|
|
|
all_per_token_logps = all_per_token_logps.reshape((hidden_states.shape[0], hidden_states.shape[1]))
|
|
return all_per_token_logps
|
|
|
|
@torch.compile(dynamic = True, fullgraph = True, options = torch_compile_options,)
|
|
def chunked_selective_log_softmax(logits, index):
|
|
# Split into 4 chunks only
|
|
chunked_logits = torch.chunk(logits.reshape(-1, logits.shape[-1]), chunks = 4, dim = 0)
|
|
chunked_index = torch.chunk(index.reshape(-1), chunks = 4, dim = 0)
|
|
all_per_token_logps = []
|
|
# Below loop does the same as selective_log_softmax(chunk_logits, chunk_index)
|
|
for chunk_logits, chunk_index in zip(chunked_logits, chunked_index):
|
|
chunk_logits = chunk_logits.to(torch.float32)
|
|
selected_logits = torch.gather(chunk_logits, dim = -1, index = chunk_index.unsqueeze(-1)).squeeze(-1)
|
|
logsumexp_values = torch.logsumexp(chunk_logits, dim = -1)
|
|
per_token_logps = selected_logits - logsumexp_values
|
|
all_per_token_logps.append(per_token_logps)
|
|
pass
|
|
all_per_token_logps = torch.concat(all_per_token_logps)
|
|
all_per_token_logps = all_per_token_logps.reshape((logits.shape[0], logits.shape[1]))
|
|
return all_per_token_logps
|
|
|
|
def calculate_pad_tokens_in_prompt(
|
|
input_ids: torch.Tensor,
|
|
logits_to_keep: int,
|
|
pad_token_id: int
|
|
) -> torch.Tensor:
|
|
"""
|
|
Given prompt tensor, it returns all the left padded tokens in that sequence. so [pad, pad, pad, cat] = 3 tokens
|
|
"""
|
|
if logits_to_keep >= input_ids.shape[1]:
|
|
raise ValueError("logits_to_keep must be smaller than the sequence length.")
|
|
|
|
prompt_section = input_ids[:, :-logits_to_keep]
|
|
|
|
padding_mask = (prompt_section == pad_token_id)
|
|
|
|
pad_token_counts = padding_mask.sum(dim=1)
|
|
|
|
return pad_token_counts
|
|
|
|
def create_completion_attention_mask(
|
|
completion_input_ids: torch.Tensor,
|
|
left_pad_tokens_per_prompt: torch.Tensor,
|
|
max_left_pad: int,
|
|
pad_token_id: int
|
|
) -> torch.Tensor:
|
|
"""
|
|
Given that we have a sequence, [p,p,p,c,c,c,pad,pad,pad]
|
|
|
|
Where p are extra prompt tokens we got from slicing the torch tensor, c is completion tokens
|
|
and pad are pad tokens, this function would make a completion mask that would 0 out the pad
|
|
and p tokens. so in this example [0,0,0,1,1,1,0,0,0]
|
|
"""
|
|
batch_size, completion_len = completion_input_ids.shape
|
|
device = completion_input_ids.device
|
|
|
|
num_tokens_to_mask = max_left_pad - left_pad_tokens_per_prompt
|
|
|
|
indices = torch.arange(completion_len, device=device).unsqueeze(0)
|
|
shift_mask = indices >= num_tokens_to_mask.unsqueeze(1)
|
|
|
|
non_padding_mask = (completion_input_ids != pad_token_id)
|
|
|
|
final_mask = shift_mask & non_padding_mask
|
|
|
|
return final_mask
|
|
|
|
def left_pack_padding(tensor: torch.Tensor, pad_id: int) -> torch.Tensor:
|
|
"""
|
|
Moves all padding tokens in each sequence of a batch to the right.
|
|
"""
|
|
mask = (tensor != pad_id)
|
|
# Must do stable=True since binary mark is unordered
|
|
sorted_indices = torch.argsort(mask, dim=1, descending=True, stable=True)
|
|
packed_tensor = torch.gather(tensor, 1, sorted_indices)
|
|
return packed_tensor
|
|
|
|
def align_logprobs_with_mask(
|
|
logprob_tensor: torch.Tensor,
|
|
attention_mask: torch.Tensor,
|
|
pad_value: float = 0.0
|
|
) -> torch.Tensor:
|
|
"""
|
|
Aligns a log probability tensor with a given attention mask.
|
|
"""
|
|
|
|
device = logprob_tensor.device
|
|
batch_size, logprob_seq_len = logprob_tensor.shape
|
|
mask_seq_len = attention_mask.shape[1]
|
|
|
|
padded_logprobs = torch.full(
|
|
attention_mask.shape,
|
|
fill_value=pad_value,
|
|
dtype=logprob_tensor.dtype,
|
|
device=device
|
|
)
|
|
|
|
left_pad_counts = torch.argmax(attention_mask, dim=1)
|
|
|
|
cols = torch.arange(logprob_seq_len, device=device)
|
|
dest_indices = left_pad_counts.unsqueeze(1) + cols
|
|
|
|
# Create destination row indices
|
|
# Shape: [batch_size, logprob_seq_len]
|
|
row_indices = torch.arange(batch_size, device=device).unsqueeze(1).expand_as(dest_indices)
|
|
|
|
# --- 4. Filter out-of-bounds indices and perform assignment ---
|
|
# Create a mask to identify only the indices that are within the bounds
|
|
# of the target tensor's sequence length.
|
|
valid_mask = dest_indices < mask_seq_len
|
|
|
|
# Use this mask to select only the valid row indices, column indices,
|
|
# and the corresponding values from the logprob tensor.
|
|
# This flattens the selected elements into 1D tensors.
|
|
valid_rows = row_indices[valid_mask]
|
|
valid_cols = dest_indices[valid_mask]
|
|
valid_vals = logprob_tensor[valid_mask]
|
|
|
|
# Place the valid values into their correct positions in the padded tensor
|
|
# using a single, efficient advanced indexing operation.
|
|
padded_logprobs[valid_rows, valid_cols] = valid_vals
|
|
|
|
return padded_logprobs
|
|
|
|
def autotune_batch_and_chunks(
|
|
total_input_rows,
|
|
seq_len,
|
|
hidden_size,
|
|
vocab_size,
|
|
dtype_bytes=16,
|
|
multiplier=None
|
|
):
|
|
if multiplier is None:
|
|
final_m = max(4, seq_len // 4096)
|
|
else:
|
|
final_m = multiplier
|
|
|
|
if torch.cuda.is_available():
|
|
free_bytes, _ = torch.cuda.mem_get_info()
|
|
limit_gb = (free_bytes / (1024**3))*.80
|
|
elif hasattr(torch, "xpu") and torch.xpu.is_available():
|
|
# For XPU: estimate free memory from total - reserved
|
|
total_mem = torch.xpu.get_device_properties(0).total_memory
|
|
reserved_mem = torch.xpu.memory_reserved()
|
|
free_bytes = total_mem - reserved_mem
|
|
limit_gb = (free_bytes / (1024**3)) * 0.80
|
|
else:
|
|
# Fallback: assume 8GB available
|
|
limit_gb = 8.0
|
|
|
|
bytes_to_gb = 1024**3
|
|
|
|
b_vals = torch.arange(total_input_rows, 0, -1, device='cpu', dtype=torch.float32)
|
|
|
|
hidden_gb = (b_vals * seq_len * hidden_size * dtype_bytes) / bytes_to_gb
|
|
|
|
base_logits = ((b_vals/total_input_rows) * b_vals * seq_len * vocab_size * dtype_bytes) / bytes_to_gb
|
|
logits_gb = base_logits / final_m
|
|
|
|
total_mem_gb = hidden_gb + logits_gb
|
|
|
|
valid_mask = total_mem_gb <= limit_gb
|
|
valid_indices = torch.nonzero(valid_mask, as_tuple=False)
|
|
|
|
if valid_indices.shape[0] == 0:
|
|
#This means your GPU will OOM
|
|
return 4, final_m
|
|
|
|
best_idx = valid_indices[0].item()
|
|
final_b = int(b_vals[best_idx].item())
|
|
|
|
return final_b, final_m
|
|
@dataclass
|
|
class UnslothNashMDConfig(NashMDConfig):
|
|
"""
|
|
|
|
Configuration class for the [`NashMDTrainer`].
|
|
|
|
Subclass of [`OnlineDPOConfig`] we can use all its arguments and add the following:
|
|
|
|
Parameters:
|
|
mixture_coef (`float` or `list[float]`, *optional*, defaults to `0.5`):
|
|
Logit mixture coefficient for the model and reference model. If a list of floats is provided then the
|
|
mixture coefficient is selected for each new epoch and the last coefficient is used for the rest of the
|
|
epochs.
|
|
|
|
"""
|
|
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.'},
|
|
)
|
|
unsloth_logit_chunk_multiplier : Optional[int] = field(
|
|
default = None,
|
|
metadata = {'help': 'Multiplier for chunked logit computations.'},
|
|
)
|
|
unsloth_grpo_mini_batch : Optional[int] = field(
|
|
default = None,
|
|
metadata = {'help': 'Mini batch size for GRPO hidden state accumulation. Default is None unless user defines it.'},
|
|
)
|
|
max_seq_length : Optional[int] = field(
|
|
default = None,
|
|
metadata = {'help': 'Maximum sequence length to truncate to.'},
|
|
)
|
|
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',
|
|
lr_scheduler_kwargs = None,
|
|
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,
|
|
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 = None,
|
|
fsdp_min_num_params = 0,
|
|
fsdp_config = None,
|
|
fsdp_transformer_layer_cls_to_wrap = None,
|
|
accelerator_config = None,
|
|
parallelism_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',
|
|
project = 'huggingface',
|
|
trackio_space_id = 'trackio',
|
|
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,
|
|
hub_revision = None,
|
|
gradient_checkpointing = True,
|
|
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,
|
|
liger_kernel_config = None,
|
|
eval_use_gather_object = False,
|
|
average_tokens_across_devices = True,
|
|
reward_model_path = None,
|
|
judge = None,
|
|
max_new_tokens = 64,
|
|
max_length = 512,
|
|
temperature = 0.9,
|
|
top_p = 1.0,
|
|
top_k = None,
|
|
min_p = None,
|
|
repetition_penalty = 1.0,
|
|
generation_kwargs = {},
|
|
use_transformers_paged = False,
|
|
cache_implementation = None,
|
|
missing_eos_penalty = None,
|
|
loss_type = 'sigmoid',
|
|
disable_dropout = True,
|
|
use_vllm = False,
|
|
vllm_model_impl = 'vllm',
|
|
vllm_guided_decoding_regex = None,
|
|
vllm_gpu_memory_utilization = 0.55,
|
|
vllm_mode = 'colocate',
|
|
vllm_server_base_url = None,
|
|
vllm_server_host = '0.0.0.0',
|
|
vllm_server_port = 8000,
|
|
vllm_server_timeout = 240.0,
|
|
vllm_tensor_parallel_size = 1,
|
|
ds3_gather_for_generation = True,
|
|
model_init_kwargs = None,
|
|
reward_weights = None,
|
|
dataset_num_proc = None,
|
|
gpu_memory_utilization = None,
|
|
vllm_sampling_params = None,
|
|
unsloth_num_chunks = -1,
|
|
unsloth_logit_chunk_multiplier = None,
|
|
unsloth_grpo_mini_batch = None,
|
|
max_seq_length = None,
|
|
**kwargs,
|
|
):
|
|
if learning_rate < 1e-7: print(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: print(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 num_train_epochs is None:
|
|
num_train_epochs = 3.0 # Default to 3 epochs if None, max_steps will override
|
|
if output_dir is None and save_strategy == 'steps' and save_steps == 500:
|
|
output_dir = 'unsloth_training_checkpoints'
|
|
save_strategy = 'no'
|
|
import multiprocessing as _mp
|
|
if _mp.get_start_method() != 'fork':
|
|
dataset_num_proc = None
|
|
elif dataset_num_proc is None:
|
|
import psutil
|
|
dataset_num_proc = min(max((psutil.cpu_count() or 1)+4, 2), 64)
|
|
memory_gb_left = psutil.virtual_memory().available / (1024**3)
|
|
if memory_gb_left <= 2: dataset_num_proc = 1
|
|
else: dataset_num_proc = min(dataset_num_proc, int(memory_gb_left))
|
|
if temperature <= 0:
|
|
raise ValueError('Unsloth: Please set a positive non-zero temperature since your results will be wrong.')
|
|
elif temperature >= 10:
|
|
raise ValueError('Unsloth: Please set a positive non-zero temperature less than 10, since sampling will be quite erratic.')
|
|
|
|
|
|
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,
|
|
lr_scheduler_kwargs = lr_scheduler_kwargs,
|
|
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,
|
|
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,
|
|
parallelism_config = parallelism_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,
|
|
project = project,
|
|
trackio_space_id = trackio_space_id,
|
|
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,
|
|
hub_revision = hub_revision,
|
|
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,
|
|
liger_kernel_config = liger_kernel_config,
|
|
eval_use_gather_object = eval_use_gather_object,
|
|
average_tokens_across_devices = average_tokens_across_devices,
|
|
reward_model_path = reward_model_path,
|
|
judge = judge,
|
|
max_new_tokens = max_new_tokens,
|
|
max_length = max_length,
|
|
temperature = temperature,
|
|
top_p = top_p,
|
|
top_k = top_k,
|
|
min_p = min_p,
|
|
repetition_penalty = repetition_penalty,
|
|
generation_kwargs = generation_kwargs,
|
|
use_transformers_paged = use_transformers_paged,
|
|
cache_implementation = cache_implementation,
|
|
missing_eos_penalty = missing_eos_penalty,
|
|
loss_type = loss_type,
|
|
disable_dropout = disable_dropout,
|
|
use_vllm = use_vllm,
|
|
vllm_model_impl = vllm_model_impl,
|
|
vllm_guided_decoding_regex = vllm_guided_decoding_regex,
|
|
vllm_gpu_memory_utilization = vllm_gpu_memory_utilization,
|
|
vllm_mode = vllm_mode,
|
|
vllm_server_base_url = vllm_server_base_url,
|
|
vllm_server_host = vllm_server_host,
|
|
vllm_server_port = vllm_server_port,
|
|
vllm_server_timeout = vllm_server_timeout,
|
|
vllm_tensor_parallel_size = vllm_tensor_parallel_size,
|
|
ds3_gather_for_generation = ds3_gather_for_generation,
|
|
model_init_kwargs = model_init_kwargs,
|
|
reward_weights = reward_weights,
|
|
dataset_num_proc = dataset_num_proc,
|
|
gpu_memory_utilization = gpu_memory_utilization,**kwargs)
|
|
self.vllm_sampling_params = vllm_sampling_params
|
|
self.unsloth_num_chunks = unsloth_num_chunks
|
|
if unsloth_grpo_mini_batch is not None:
|
|
if self.generation_batch_size >= unsloth_grpo_mini_batch:
|
|
self.unsloth_grpo_mini_batch = unsloth_grpo_mini_batch
|
|
else:
|
|
raise ValueError(
|
|
f"Unsloth GRPO mini batch size needs to be less than or equal to the effective generation batch size, "
|
|
f"which is self.per_device_train_batch_size * gradient_accumulation_steps."
|
|
)
|
|
self.unsloth_logit_chunk_multiplier = unsloth_logit_chunk_multiplier
|
|
self.max_seq_length = max_seq_length
|
|
|
|
pass
|
|
|
|
class _UnslothNashMDTrainer(OnlineDPOTrainer):
|
|
""""""
|
|
|
|
_tag_names = ["trl", "nash-md"]
|
|
_name = "Nash-MD"
|
|
_paper = {
|
|
"title": "Nash Learning from Human Feedback",
|
|
"id": "2312.00886",
|
|
# docstyle-ignore
|
|
"citation": textwrap.dedent("""\
|
|
@inproceedings{munos2024nash,
|
|
title = {{Nash Learning from Human Feedback}},
|
|
author = {R{\'{e}}mi Munos and Michal Valko and Daniele Calandriello and Mohammad Gheshlaghi Azar and Mark Rowland and Zhaohan Daniel Guo and Yunhao Tang and Matthieu Geist and Thomas Mesnard and C{\\^{o}}me Fiegel and Andrea Michi and Marco Selvi and Sertan Girgin and Nikola Momchev and Olivier Bachem and Daniel J. Mankowitz and Doina Precup and Bilal Piot},
|
|
year = 2024,
|
|
booktitle = {Forty-first International Conference on Machine Learning, {ICML} 2024, Vienna, Austria, July 21-27, 2024},
|
|
publisher = {OpenReview.net},
|
|
url = {https://openreview.net/forum?id=Y5AmNYiyCQ}
|
|
}"""),
|
|
}
|
|
|
|
def __init__(
|
|
self,
|
|
model: Union[PreTrainedModel, nn.Module] = None,
|
|
ref_model: Union[PreTrainedModel, nn.Module] = None,
|
|
reward_funcs: Union[PreTrainedModel, nn.Module, None] = None,
|
|
judge: Optional[BasePairwiseJudge] = None,
|
|
args: Optional[NashMDConfig] = None,
|
|
data_collator: Optional[Callable] = None,
|
|
train_dataset: Optional[Union[Dataset, IterableDataset]] = None,
|
|
eval_dataset: Optional[Union[Dataset, dict[str, Dataset]]] = None,
|
|
processing_class: Optional[
|
|
Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin]
|
|
] = None,
|
|
peft_config: Optional[dict] = None,
|
|
compute_metrics: Optional[Callable[[EvalPrediction], 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,
|
|
# Deprecated parameters
|
|
reward_model: Optional[Union[PreTrainedModel, nn.Module]] = None,
|
|
) -> None:
|
|
super().__init__(
|
|
model=model,
|
|
ref_model=ref_model,
|
|
reward_funcs=reward_funcs,
|
|
judge=judge,
|
|
args=args,
|
|
data_collator=data_collator,
|
|
train_dataset=train_dataset,
|
|
eval_dataset=eval_dataset,
|
|
processing_class=processing_class,
|
|
reward_processing_classes=processing_class,
|
|
peft_config=peft_config,
|
|
compute_metrics=compute_metrics,
|
|
callbacks=callbacks,
|
|
optimizers=optimizers,
|
|
preprocess_logits_for_metrics=preprocess_logits_for_metrics,
|
|
reward_model=reward_model,
|
|
)
|
|
|
|
self._mixture_coef = self.args.mixture_coef
|
|
|
|
# Overwrite the stats dictionary to include NashMD specific statistics
|
|
self.stats = {
|
|
# Remove "non_score_reward", "rlhf_reward", "scores_margin"
|
|
# Add "mixture_coef"
|
|
"loss/kl": [],
|
|
"objective/entropy": [],
|
|
"loss/score": [],
|
|
"rewards/probabilities": [],
|
|
"rewards/accuracies": [],
|
|
"rewards/margins": [],
|
|
"logps/chosen": [],
|
|
"logps/rejected": [],
|
|
"val/model_contain_eos_token": [],
|
|
"val/ref_contain_eos_token": [],
|
|
"beta": [],
|
|
"mixture_coef": [],
|
|
}
|
|
if self.reward_funcs is not None:
|
|
if len(self.reward_funcs) != 1:
|
|
raise ValueError("NashMDTrainer only supports one reward function/model.")
|
|
self.reward_funcs = self.reward_funcs[0]
|
|
self.stats["rewards/chosen"] = []
|
|
self.stats["rewards/rejected"] = []
|
|
|
|
@property
|
|
def mixture_coef(self):
|
|
if isinstance(self._mixture_coef, list):
|
|
epoch = self.state.epoch
|
|
return self._mixture_coef[epoch] if epoch < len(self._mixture_coef) else self._mixture_coef[-1]
|
|
else:
|
|
return self._mixture_coef
|
|
|
|
def _generate_completions(self, model, prompts):
|
|
# Generate completions from the policy model.
|
|
with unwrap_model_for_generation(model, self.accelerator) as unwrapped_policy_for_gen_ctx:
|
|
model_output = unwrapped_policy_for_gen_ctx.generate(
|
|
input_ids=prompts["input_ids"],
|
|
attention_mask=prompts["attention_mask"],
|
|
generation_config=self.generation_config,
|
|
)
|
|
|
|
# Get the DDP/FSDP unwrapped version of the main model.
|
|
# This will be the policy model for GeometricMixtureWrapper (PEFT adapters active if PEFT is used).
|
|
policy_model_for_gmw = self.accelerator.unwrap_model(model)
|
|
|
|
# Determine the correct reference model for GeometricMixtureWrapper.
|
|
# This also needs to be DDP/FSDP unwrapped.
|
|
ref_model_for_gmw: torch.nn.Module
|
|
if self.ref_model is None:
|
|
# No explicit ref_model is provided.
|
|
# Use the base of the main `model` if it's a PEFT model.
|
|
# policy_model_for_gmw is already DDP-unwrapped.
|
|
if is_peft_available() and isinstance(policy_model_for_gmw, PeftModel):
|
|
ref_model_for_gmw = policy_model_for_gmw.get_base_model()
|
|
else:
|
|
# Not a PEFT model (or PEFT not available), or already a base model.
|
|
# Use the DDP-unwrapped policy model itself as the reference.
|
|
ref_model_for_gmw = policy_model_for_gmw
|
|
else:
|
|
# An explicit ref_model is provided. Unwrap it for DDP/FSDP.
|
|
ref_model_for_gmw = self.accelerator.unwrap_model(self.ref_model)
|
|
|
|
# Both models given to GeometricMixtureWrapper (policy_model_for_gmw and ref_model_for_gmw) are DDP-unwrapped.
|
|
with torch.no_grad(): # Ensure no_grad context for mixture model generation
|
|
mixture_model = GeometricMixtureWrapper(
|
|
model=policy_model_for_gmw,
|
|
ref_model=ref_model_for_gmw,
|
|
generation_config=self.generation_config,
|
|
mixture_coef=self.mixture_coef,
|
|
device=self.accelerator.device,
|
|
)
|
|
|
|
mixture_output = mixture_model.generate(
|
|
input_ids=prompts["input_ids"],
|
|
attention_mask=prompts["attention_mask"],
|
|
generation_config=self.generation_config,
|
|
)
|
|
|
|
return model_output, mixture_output
|
|
|
|
def _process_completions(self, model_output, mixture_output, prompts):
|
|
context_length = prompts["input_ids"].shape[1]
|
|
|
|
# Process model completions
|
|
model_completion_ids = model_output[:, context_length:]
|
|
model_completion_ids, model_completion_mask = truncate_right(
|
|
model_completion_ids, self.processing_class.eos_token_id, self.processing_class.pad_token_id
|
|
)
|
|
model_data = {
|
|
"input_ids": torch.cat((prompts["input_ids"], model_completion_ids), dim=1),
|
|
"attention_mask": torch.cat((prompts["attention_mask"], model_completion_mask), dim=1),
|
|
"raw": prompts["raw"],
|
|
}
|
|
|
|
# Process reference model completions
|
|
mixture_completion_ids = mixture_output[:, context_length:]
|
|
mixture_completion_ids, mixture_completion_mask = truncate_right(
|
|
mixture_completion_ids, self.processing_class.eos_token_id, self.processing_class.pad_token_id
|
|
)
|
|
mixture_data = {
|
|
"input_ids": torch.cat((prompts["input_ids"], mixture_completion_ids), dim=1),
|
|
"attention_mask": torch.cat((prompts["attention_mask"], mixture_completion_mask), dim=1),
|
|
"raw": prompts["raw"],
|
|
}
|
|
|
|
return model_data, mixture_data
|
|
|
|
def _compute_rewards(self, model_data, mixture_data, context_length):
|
|
with torch.no_grad():
|
|
_, model_scores, _ = get_reward(
|
|
self.reward_funcs, model_data["input_ids"], self.processing_class.pad_token_id, context_length
|
|
)
|
|
_, mixture_scores, _ = get_reward(
|
|
self.reward_funcs, mixture_data["input_ids"], self.processing_class.pad_token_id, context_length
|
|
)
|
|
|
|
# Apply EOS penalty if needed
|
|
if self.args.missing_eos_penalty is not None:
|
|
model_contain_eos = torch.any(model_data["input_ids"] == self.processing_class.eos_token_id, dim=-1)
|
|
mixture_contain_eos = torch.any(mixture_data["input_ids"] == self.processing_class.eos_token_id, dim=-1)
|
|
model_scores[~model_contain_eos] -= self.args.missing_eos_penalty
|
|
mixture_scores[~mixture_contain_eos] -= self.args.missing_eos_penalty
|
|
|
|
return model_scores, mixture_scores
|
|
|
|
def _compute_judge(self, model_data, mixture_data, context_length):
|
|
prompts = model_data["raw"]
|
|
model_data_completions = self.processing_class.batch_decode(
|
|
model_data["input_ids"][:, context_length:], skip_special_tokens=True
|
|
)
|
|
model_data_completions = [completion.strip() for completion in model_data_completions]
|
|
|
|
mixture_data_completions = self.processing_class.batch_decode(
|
|
mixture_data["input_ids"][:, context_length:], skip_special_tokens=True
|
|
)
|
|
mixture_data_completions = [completion.strip() for completion in mixture_data_completions]
|
|
if is_conversational({"prompt": prompts[0]}):
|
|
model_data_completions = [
|
|
[{"role": "assistant", "content": completion}] for completion in model_data_completions
|
|
]
|
|
environment = jinja2.Environment()
|
|
template = environment.from_string(SIMPLE_CHAT_TEMPLATE)
|
|
prompts = [template.render(messages=message) for message in prompts]
|
|
model_data_completions = [template.render(messages=completion) for completion in model_data_completions]
|
|
|
|
mixture_data_completions = [
|
|
[{"role": "assistant", "content": completion}] for completion in mixture_data_completions
|
|
]
|
|
mixture_data_completions = [
|
|
template.render(messages=completion) for completion in mixture_data_completions
|
|
]
|
|
|
|
probability = self.judge.judge(
|
|
prompts,
|
|
list(zip(model_data_completions, mixture_data_completions)),
|
|
return_scores=True,
|
|
)
|
|
return torch.tensor(probability, device=model_data["input_ids"].device)
|
|
|
|
def _compute_logprobs(self, model, model_data, context_length):
|
|
def compute_logprobs_for_data(m, data):
|
|
output = m(data["input_ids"], attention_mask=data["attention_mask"])
|
|
logits = output.logits[:, context_length - 1 : -1]
|
|
token_logprobs = selective_log_softmax(logits, data["input_ids"][:, context_length:])
|
|
return token_logprobs
|
|
|
|
# Compute logprobs for model completions under the model
|
|
model_logprobs_model_data = compute_logprobs_for_data(model, model_data)
|
|
|
|
# Compute logprobs of model completions under the reference model
|
|
with torch.no_grad():
|
|
if self.ref_model is None:
|
|
with model.disable_adapter():
|
|
ref_logprobs_model_data = compute_logprobs_for_data(model, model_data)
|
|
else:
|
|
ref_logprobs_model_data = compute_logprobs_for_data(self.ref_model, model_data)
|
|
|
|
# Mask padding tokens
|
|
model_padding_mask = model_data["attention_mask"][:, context_length:] == 0
|
|
model_logprobs_model_data = model_logprobs_model_data.masked_fill(model_padding_mask, 0.0)
|
|
ref_logprobs_model_data = ref_logprobs_model_data.masked_fill(model_padding_mask, 0.0)
|
|
|
|
return (model_logprobs_model_data, ref_logprobs_model_data)
|
|
|
|
def _compute_losses(
|
|
self,
|
|
model_logprobs_model_data,
|
|
ref_logprobs_model_data,
|
|
probability,
|
|
):
|
|
# reinforce score where 0.5 is a control variate
|
|
score = (probability - 0.5) * model_logprobs_model_data.sum(1)
|
|
|
|
# kl divergence via reinforce
|
|
with torch.no_grad():
|
|
log_ratio = model_logprobs_model_data - ref_logprobs_model_data
|
|
kl_div_log = log_ratio.sum(1)
|
|
kl_div_loss = (log_ratio * model_logprobs_model_data).sum(1)
|
|
|
|
# final loss
|
|
loss = self.beta * kl_div_loss - score
|
|
|
|
return loss.mean(), score, kl_div_log
|
|
|
|
def _log_statistics(
|
|
self,
|
|
model_data,
|
|
mixture_data,
|
|
model_logprobs_model_data,
|
|
ref_logprobs_model_data,
|
|
probability,
|
|
score,
|
|
kl_div,
|
|
context_length,
|
|
model_scores=None,
|
|
mixture_scores=None,
|
|
):
|
|
# Helper function to gather and compute mean
|
|
def gather_mean(tensor):
|
|
return self.accelerator.gather_for_metrics(tensor).mean().item()
|
|
|
|
# Log score
|
|
self.stats["loss/score"].append(gather_mean(score))
|
|
# Log KL divergence
|
|
self.stats["loss/kl"].append(gather_mean(kl_div))
|
|
|
|
# Log logprobs
|
|
model_logprobs_model_data_sum = model_logprobs_model_data.sum(1)
|
|
ref_logprobs_model_data_sum = ref_logprobs_model_data.sum(1)
|
|
|
|
self.stats["logps/chosen"].append(gather_mean(model_logprobs_model_data_sum))
|
|
self.stats["logps/rejected"].append(gather_mean(ref_logprobs_model_data_sum))
|
|
|
|
# Log rewards
|
|
if self.reward_funcs is not None:
|
|
self.stats["rewards/chosen"].append(gather_mean(model_scores))
|
|
self.stats["rewards/rejected"].append(gather_mean(mixture_scores))
|
|
|
|
# Log probabilities
|
|
self.stats["rewards/probabilities"].append(gather_mean(probability))
|
|
|
|
# Calculate entropy for model data
|
|
entropy_model_data = -model_logprobs_model_data.sum(1)
|
|
self.stats["objective/entropy"].append(gather_mean(entropy_model_data))
|
|
|
|
# Calculate margins
|
|
margin = model_logprobs_model_data_sum - ref_logprobs_model_data_sum
|
|
self.stats["rewards/margins"].append(gather_mean(margin))
|
|
|
|
# Calculate accuracy
|
|
accuracy = (margin > 0).float()
|
|
self.stats["rewards/accuracies"].append(gather_mean(accuracy))
|
|
|
|
# Log EOS token statistics
|
|
model_eos = (model_data["input_ids"][:, context_length:] == self.processing_class.eos_token_id).any(dim=1)
|
|
mixture_eos = (mixture_data["input_ids"][:, context_length:] == self.processing_class.eos_token_id).any(dim=1)
|
|
self.stats["val/model_contain_eos_token"].append(gather_mean(model_eos.float()))
|
|
self.stats["val/ref_contain_eos_token"].append(gather_mean(mixture_eos.float()))
|
|
|
|
# Log beta and mixture coef
|
|
self.stats["beta"].append(self.beta)
|
|
self.stats["mixture_coef"].append(self.mixture_coef)
|
|
|
|
def training_step(
|
|
self, model: nn.Module, inputs: dict[str, Union[torch.Tensor, Any]], num_items_in_batch: Optional[int] = None
|
|
) -> torch.Tensor:
|
|
model.train()
|
|
|
|
# Apply chat template and tokenize the input
|
|
batch_size = len(next(iter(inputs.values())))
|
|
prompts = inputs["prompt"]
|
|
inputs = [{k: v[i] for k, v in inputs.items()} for i in range(batch_size)]
|
|
inputs = [maybe_apply_chat_template(x, self.processing_class) for x in inputs]
|
|
inputs = [self.tokenize_row(x, self.model.config.is_encoder_decoder, self.processing_class) for x in inputs]
|
|
inputs = self.data_collator(inputs)
|
|
|
|
# need the prompt_ only
|
|
inputs = self._prepare_inputs(inputs)
|
|
context_length = inputs["prompt_input_ids"].shape[1]
|
|
prompts = {
|
|
"input_ids": inputs["prompt_input_ids"],
|
|
"attention_mask": inputs["prompt_attention_mask"],
|
|
"raw": prompts,
|
|
}
|
|
del inputs
|
|
|
|
# Sample completions from both the model and the reference model
|
|
model_output, mixture_output = self._generate_completions(model, prompts)
|
|
|
|
# Process model completions
|
|
model_data, mixture_data = self._process_completions(model_output, mixture_output, prompts)
|
|
|
|
# Compute rewards
|
|
if self.reward_funcs is not None:
|
|
model_scores, mixture_scores = self._compute_rewards(model_data, mixture_data, context_length)
|
|
# probability of the model data vs the mixture data
|
|
probability = F.sigmoid(model_scores - mixture_scores)
|
|
else:
|
|
model_scores, mixture_scores = None, None
|
|
probability = self._compute_judge(model_data, mixture_data, context_length)
|
|
|
|
# Compute logprobs
|
|
model_logprobs_model_data, ref_logprobs_model_data = self._compute_logprobs(model, model_data, context_length)
|
|
|
|
# Compute loss
|
|
loss, score, kl_div = self._compute_losses(model_logprobs_model_data, ref_logprobs_model_data, probability)
|
|
|
|
# Log everything
|
|
self._log_statistics(
|
|
model_data,
|
|
mixture_data,
|
|
model_logprobs_model_data.detach(),
|
|
ref_logprobs_model_data,
|
|
probability,
|
|
score.detach(),
|
|
kl_div.detach(),
|
|
context_length,
|
|
model_scores,
|
|
mixture_scores,
|
|
)
|
|
|
|
if (
|
|
self.args.torch_empty_cache_steps is not None
|
|
and self.state.global_step % self.args.torch_empty_cache_steps == 0
|
|
):
|
|
empty_cache()
|
|
|
|
kwargs = {}
|
|
# For LOMO optimizers you need to explicitly use the learning rate
|
|
if self.args.optim in [OptimizerNames.LOMO, OptimizerNames.ADALOMO]:
|
|
kwargs["learning_rate"] = self._get_learning_rate()
|
|
|
|
if self.args.n_gpu > 1:
|
|
loss = loss.mean() # mean() to average on multi-gpu parallel training
|
|
|
|
self.accelerator.backward(loss, **kwargs)
|
|
|
|
return loss.detach() / self.args.gradient_accumulation_steps
|
|
class UnslothNashMDTrainer(_UnslothNashMDTrainer):
|
|
"""
|
|
|
|
Trainer for the Nash-MD method.
|
|
|
|
It is implemented as a subclass of [`OnlineDPOTrainer`].
|
|
|
|
Args:
|
|
model ([`~transformers.PreTrainedModel`]):
|
|
The model to train, preferably an `AutoModelForCausalLM`.
|
|
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.
|
|
reward_funcs ([`~transformers.PreTrainedModel`]):
|
|
The reward model to score completions with, preferably an
|
|
[`~transformers.AutoModelForSequenceClassification`].
|
|
judge ([`BasePairwiseJudge`]):
|
|
The judge to use for pairwise comparison of model completions.
|
|
args ([`NashMDConfig`]):
|
|
The NashMD 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
|
|
([`DPODataCollatorWithPadding`]) 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 ([`~transformers.PreTrainedTokenizerBase`], [`~transformers.BaseImageProcessor`], [`~transformers.FeatureExtractionMixin`] or [`~transformers.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.
|
|
peft_config (`dict`):
|
|
The peft config to use for training.
|
|
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.
|
|
|
|
reward_model:
|
|
|
|
<Deprecated version="0.22.0">
|
|
|
|
This parameter is deprecated and will be removed in version 0.25.0. Use `reward_funcs` instead.
|
|
|
|
</Deprecated>
|
|
|
|
"""
|
|
def __init__(
|
|
self,
|
|
model = None,
|
|
ref_model = None,
|
|
reward_funcs = None,
|
|
judge = None,
|
|
args = None,
|
|
data_collator = None,
|
|
train_dataset = None,
|
|
eval_dataset = None,
|
|
processing_class = None,
|
|
peft_config = None,
|
|
compute_metrics = None,
|
|
callbacks = None,
|
|
preprocess_logits_for_metrics = None,
|
|
reward_model = None,
|
|
**kwargs
|
|
):
|
|
if args is None: args = UnslothNashMDConfig()
|
|
use_bf16 = getattr(args, 'bf16', False)
|
|
if type(use_bf16) is not bool: use_bf16 = False
|
|
use_fp16 = getattr(args, 'fp16', False)
|
|
if type(use_fp16) is not bool: use_fp16 = False
|
|
force_float32 = False
|
|
full_finetuning = os.environ.get('UNSLOTH_ENABLE_FULL_FINETUNING', '0') == '1'
|
|
if not full_finetuning and (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, 'dtype', None) or getattr(model.config, 'torch_dtype', None)
|
|
if dtype is None: dtype = model.get_input_embeddings().weight.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:
|
|
# Forced float32 training
|
|
args.fp16 = False
|
|
args.bf16 = False
|
|
os.environ['ACCELERATE_MIXED_PRECISION'] = 'no'
|
|
if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no'
|
|
# args.mixed_precision is a new argument which needs to be set now
|
|
elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32':
|
|
# Mixed precision training
|
|
args.fp16 = float16
|
|
args.bf16 = not float16
|
|
os.environ['ACCELERATE_MIXED_PRECISION'] = 'fp16' if float16 else 'bf16'
|
|
if hasattr(args, 'mixed_precision'): args.mixed_precision = 'fp16' if float16 else 'bf16'
|
|
# args.mixed_precision is a new argument which needs to be set now
|
|
elif mixed_precision_dtype == 'bfloat16':
|
|
# Both False since bfloat16 full finetuning doesn't do any autocasting.
|
|
args.fp16 = False
|
|
args.bf16 = False
|
|
os.environ['ACCELERATE_MIXED_PRECISION'] = 'no'
|
|
if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no'
|
|
# args.mixed_precision is a new argument which needs to be set now
|
|
|
|
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)
|
|
if type(fp16_full_eval) is not bool: fp16_full_eval = False
|
|
bf16_full_eval = getattr(args, 'bf16_full_eval', False)
|
|
if type(bf16_full_eval) is not bool: 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
|
|
elif args_max_seq_length is not None and model_max_seq_length is not None:
|
|
if args_max_seq_length > model_max_seq_length:
|
|
print('Unsloth: You set `max_seq_length` as ' + str(args_max_seq_length) + ' but '
|
|
'the maximum the model supports is ' + str(model_max_seq_length) + '. We shall reduce it.')
|
|
args.max_seq_length = model_max_seq_length
|
|
if model is not None and hasattr(model, 'for_training'):
|
|
model.for_training(use_gradient_checkpointing=getattr(args, 'gradient_checkpointing', True))
|
|
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,
|
|
pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None),
|
|
)
|
|
elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling) and 'labels' in train_dataset.column_names:
|
|
data_collator = DataCollatorForSeq2Seq(
|
|
__tokenizer,
|
|
pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None),
|
|
)
|
|
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,
|
|
pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None),
|
|
)
|
|
else:
|
|
data_collator = TransformersDataCollatorForLanguageModeling(
|
|
__tokenizer.tokenizer,
|
|
mlm = False,
|
|
mlm_probability = 0.0,
|
|
pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None),
|
|
)
|
|
other_metrics = []
|
|
|
|
from unsloth_zoo.logging_utils import PatchRLStatistics
|
|
PatchRLStatistics('nash_md_trainer', other_metrics)
|
|
|
|
# [TODO] Fix up DataParallel multiplying batch sizes
|
|
# [TODO] DDP works, but DP seems to not work? [TODO]
|
|
if getattr(args, "parallel_mode", None) == ParallelMode.NOT_DISTRIBUTED and args.n_gpu > 1:
|
|
if getattr(args, "_n_gpu", 1) != 1:
|
|
args._n_gpu = 1
|
|
if "model" in locals() and hasattr(model, "for_training"):
|
|
model.for_training(use_gradient_checkpointing=getattr(args, 'gradient_checkpointing', True))
|
|
super().__init__(
|
|
model = model,
|
|
ref_model = ref_model,
|
|
reward_funcs = reward_funcs,
|
|
judge = judge,
|
|
args = args,
|
|
data_collator = data_collator,
|
|
train_dataset = train_dataset,
|
|
eval_dataset = eval_dataset,
|
|
processing_class = processing_class,
|
|
peft_config = peft_config,
|
|
compute_metrics = compute_metrics,
|
|
callbacks = callbacks,
|
|
preprocess_logits_for_metrics = preprocess_logits_for_metrics,
|
|
reward_model = reward_model,**kwargs)
|
|
if "model" in locals() and hasattr(model, "for_inference"):
|
|
model.for_inference()
|
|
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
|
|
if hasattr(self, 'accelerator'):
|
|
scaler = self.accelerator.scaler
|
|
current_model = model
|
|
while hasattr(current_model, 'model'):
|
|
current_model.accelerator_scaler = scaler
|
|
current_model = current_model.model
|
|
current_model.accelerator_scaler = scaler
|
|
pass
|
|
if hasattr(self, 'train'):
|
|
self.train = MethodType(prepare_for_training_mode(self.__class__.train), self)
|
|
pass
|
|
if hasattr(self, 'llm') and self.llm is not None and hasattr(self.llm, 'get_tokenizer'):
|
|
_vllm_tok = self.llm.get_tokenizer()
|
|
_pc = getattr(self, 'processing_class', None) or getattr(self, 'tokenizer', None)
|
|
if _vllm_tok is not None and _pc is not None and getattr(_pc, 'chat_template', None) is not None and getattr(_vllm_tok, 'chat_template', None) is None:
|
|
_vllm_tok.chat_template = _pc.chat_template
|
|
pass
|
|
|
|
pass
|