From 9dc75f007f74f7ea5204f867baedb9927f9a8f58 Mon Sep 17 00:00:00 2001 From: haotian <2421912570@qq.com> Date: Thu, 10 Jul 2025 17:09:39 +0800 Subject: [PATCH] =?UTF-8?q?001=E6=B5=8B=E8=AF=95unsloth,002=E4=B8=8B?= =?UTF-8?q?=E8=BD=BD=E6=95=B0=E6=8D=AE=E9=9B=86=20=E8=BF=90=E8=A1=8C?= =?UTF-8?q?=E6=88=90=E5=8A=9F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- 001测试unsloth.py | 145 + 002下载数据集.py | 6 + .../UnslothAlignPropTrainer.py | 653 +++++ unsloth_compiled_cache/UnslothBCOTrainer.py | 1823 ++++++++++++ unsloth_compiled_cache/UnslothCPOTrainer.py | 1596 +++++++++++ unsloth_compiled_cache/UnslothDDPOTrainer.py | 889 ++++++ unsloth_compiled_cache/UnslothDPOTrainer.py | 2443 ++++++++++++++++ unsloth_compiled_cache/UnslothGKDTrainer.py | 860 ++++++ unsloth_compiled_cache/UnslothGRPOTrainer.py | 2497 +++++++++++++++++ .../UnslothIterativeSFTTrainer.py | 947 +++++++ unsloth_compiled_cache/UnslothKTOTrainer.py | 2034 ++++++++++++++ .../UnslothNashMDTrainer.py | 993 +++++++ unsloth_compiled_cache/UnslothORPOTrainer.py | 1547 ++++++++++ .../UnslothOnlineDPOTrainer.py | 1260 +++++++++ unsloth_compiled_cache/UnslothPPOTrainer.py | 1300 +++++++++ unsloth_compiled_cache/UnslothPRMTrainer.py | 826 ++++++ unsloth_compiled_cache/UnslothRLOOTrainer.py | 1155 ++++++++ .../UnslothRewardTrainer.py | 844 ++++++ unsloth_compiled_cache/UnslothSFTTrainer.py | 1206 ++++++++ unsloth_compiled_cache/UnslothXPOTrainer.py | 1036 +++++++ 20 files changed, 24060 insertions(+) create mode 100644 001测试unsloth.py create mode 100644 002下载数据集.py create mode 100644 unsloth_compiled_cache/UnslothAlignPropTrainer.py create mode 100644 unsloth_compiled_cache/UnslothBCOTrainer.py create mode 100644 unsloth_compiled_cache/UnslothCPOTrainer.py create mode 100644 unsloth_compiled_cache/UnslothDDPOTrainer.py create mode 100644 unsloth_compiled_cache/UnslothDPOTrainer.py create mode 100644 unsloth_compiled_cache/UnslothGKDTrainer.py create mode 100644 unsloth_compiled_cache/UnslothGRPOTrainer.py create mode 100644 unsloth_compiled_cache/UnslothIterativeSFTTrainer.py create mode 100644 unsloth_compiled_cache/UnslothKTOTrainer.py create mode 100644 unsloth_compiled_cache/UnslothNashMDTrainer.py create mode 100644 unsloth_compiled_cache/UnslothORPOTrainer.py create mode 100644 unsloth_compiled_cache/UnslothOnlineDPOTrainer.py create mode 100644 unsloth_compiled_cache/UnslothPPOTrainer.py create mode 100644 unsloth_compiled_cache/UnslothPRMTrainer.py create mode 100644 unsloth_compiled_cache/UnslothRLOOTrainer.py create mode 100644 unsloth_compiled_cache/UnslothRewardTrainer.py create mode 100644 unsloth_compiled_cache/UnslothSFTTrainer.py create mode 100644 unsloth_compiled_cache/UnslothXPOTrainer.py diff --git a/001测试unsloth.py b/001测试unsloth.py new file mode 100644 index 0000000..9af3ffd --- /dev/null +++ b/001测试unsloth.py @@ -0,0 +1,145 @@ +from unsloth import FastLanguageModel +import torch +from datasets import load_dataset + + +fourbit_models = [ + "unsloth/Qwen3-1.7B-unsloth-bnb-4bit", # Qwen 14B 2x faster + "unsloth/Qwen3-4B-unsloth-bnb-4bit", + "unsloth/Qwen3-8B-unsloth-bnb-4bit", + "unsloth/Qwen3-14B-unsloth-bnb-4bit", + "unsloth/Qwen3-32B-unsloth-bnb-4bit", + + # 4bit dynamic quants for superior accuracy and low memory use + "unsloth/gemma-3-12b-it-unsloth-bnb-4bit", + "unsloth/Phi-4", + "unsloth/Llama-3.1-8B", + "unsloth/Llama-3.2-3B", + "unsloth/orpheus-3b-0.1-ft-unsloth-bnb-4bit" # [NEW] We support TTS models! +] + + +# 加载预训练模型 +model, tokenizer = FastLanguageModel.from_pretrained( + model_name = "unsloth/Qwen3-1.7B-unsloth-bnb-4bit", + max_seq_length = 2048, # Context length - can be longer, but uses more memory + load_in_4bit = True, # 4bit uses much less memory , 启用QLoRA + load_in_8bit = False, # A bit more accurate, uses 2x memory + full_finetuning = False, # We have full finetuning now! + # token = "hf_...", # use one if using gated models +) + +model = FastLanguageModel.get_peft_model( + model, + r = 16, # Choose any number > 0! Suggested 8, 16, 32, 64, 128 + target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", + "gate_proj", "up_proj", "down_proj",], + lora_alpha = 16, # Best to choose alpha = rank or rank*2 + lora_dropout = 0, # Supports any, but = 0 is optimized + bias = "none", # Supports any, but = "none" is optimized + # [NEW] "unsloth" uses 30% less VRAM, fits 2x larger batch sizes! + use_gradient_checkpointing = "unsloth", # True or "unsloth" for very long context + random_state = 3407, + use_rslora = False, # We support rank stabilized LoRA + loftq_config = None, # And LoftQ +) + +non_reasoning_dataset = load_dataset("mlabonne/FineTome-100k", split = "train") + +def generate_conversation(examples): + problems = examples["problem"] + solutions = examples["generated_solution"] + conversations = [] + for problem, solution in zip(problems, solutions): + conversations.append([ + {"role" : "user", "content" : problem}, + {"role" : "assistant", "content" : solution}, + ]) + return { "conversations": conversations, } + + +from unsloth.chat_templates import standardize_sharegpt +dataset = standardize_sharegpt(non_reasoning_dataset) + +# Apply chat template to each conversation in the dataset +non_reasoning_conversations = [ + tokenizer.apply_chat_template( + conversation, + tokenize = False, + ) for conversation in dataset["conversations"] +] + + +chat_percentage = 0.25 + +import pandas as pd +non_reasoning_subset = pd.Series(non_reasoning_conversations) +non_reasoning_subset = non_reasoning_subset.sample( + int(len(non_reasoning_conversations)*(chat_percentage/(1 - chat_percentage))), + random_state = 2407, +) + + +data = pd.concat([ + # pd.Series(reasoning_conversations), + pd.Series(non_reasoning_subset) +]) +data.name = "text" + +from datasets import Dataset +combined_dataset = Dataset.from_pandas(pd.DataFrame(data)) +combined_dataset = combined_dataset.shuffle(seed = 3407) + + +from trl import SFTTrainer, SFTConfig +trainer = SFTTrainer( + model = model, + tokenizer = tokenizer, + train_dataset = combined_dataset, + eval_dataset = None, # Can set up evaluation! + args = SFTConfig( + dataset_text_field = "text", + per_device_train_batch_size = 2, + gradient_accumulation_steps = 4, # Use GA to mimic batch size! + warmup_steps = 5, + # num_train_epochs = 1, # Set this for 1 full training run. + max_steps = 30, + learning_rate = 2e-4, # Reduce to 2e-5 for long training runs + logging_steps = 1, + optim = "adamw_8bit", + weight_decay = 0.01, + lr_scheduler_type = "linear", + seed = 3407, + report_to = "none", # Use this for WandB etc + ), +) + + +gpu_stats = torch.cuda.get_device_properties(0) +start_gpu_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3) +max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3) +print(f"GPU = {gpu_stats.name}. Max memory = {max_memory} GB.") +print(f"{start_gpu_memory} GB of memory reserved.") + + +trainer_stats = trainer.train() + + + +messages = [ + {"role" : "user", "content" : "Solve (x + 2)^2 = 0."} +] +text = tokenizer.apply_chat_template( + messages, + tokenize = False, + add_generation_prompt = True, # Must add for generation + enable_thinking = False, # Disable thinking +) + +from transformers import TextStreamer +_ = model.generate( + **tokenizer(text, return_tensors = "pt").to("cuda"), + max_new_tokens = 256, # Increase for longer outputs! + temperature = 0.7, top_p = 0.8, top_k = 20, # For non thinking + streamer = TextStreamer(tokenizer, skip_prompt = True), +) \ No newline at end of file diff --git a/002下载数据集.py b/002下载数据集.py new file mode 100644 index 0000000..ec38e83 --- /dev/null +++ b/002下载数据集.py @@ -0,0 +1,6 @@ +from datasets import load_dataset + + +# ds = load_dataset("Open-Orca/OpenOrca") + +non_reasoning_dataset = load_dataset("mlabonne/FineTome-100k", split = "train") \ No newline at end of file diff --git a/unsloth_compiled_cache/UnslothAlignPropTrainer.py b/unsloth_compiled_cache/UnslothAlignPropTrainer.py new file mode 100644 index 0000000..4bc0c2a --- /dev/null +++ b/unsloth_compiled_cache/UnslothAlignPropTrainer.py @@ -0,0 +1,653 @@ +""" +2025.7.1 +2025.7.1 +4.53.1 +0.19.1 +__UNSLOTH_VERSIONING__ +""" +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from trl.trainer.alignprop_trainer import (Accelerator, AlignPropConfig, AlignPropTrainer, Any, Callable, DDPOStableDiffusionPipeline, Optional, Path, ProjectConfiguration, PyTorchModelHubMixin, Union, defaultdict, generate_model_card, get_comet_experiment_url, is_wandb_available, logger, os, set_seed, textwrap, torch, warn) + + +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 UnslothAlignPropConfig(AlignPropConfig): + """ + + Configuration class for the [`AlignPropTrainer`]. + + 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: + exp_name (`str`, *optional*, defaults to `os.path.basename(sys.argv[0])[: -len(".py")]`): + Name of this experiment (defaults to the file name without the extension). + run_name (`str`, *optional*, defaults to `""`): + Name of this run. + seed (`int`, *optional*, defaults to `0`): + Random seed for reproducibility. + log_with (`str` or `None`, *optional*, defaults to `None`): + Log with either `"wandb"` or `"tensorboard"`. Check + [tracking](https://huggingface.co/docs/accelerate/usage_guides/tracking) for more details. + log_image_freq (`int`, *optional*, defaults to `1`): + Frequency for logging images. + tracker_kwargs (`dict[str, Any]`, *optional*, defaults to `{}`): + Keyword arguments for the tracker (e.g., `wandb_project`). + accelerator_kwargs (`dict[str, Any]`, *optional*, defaults to `{}`): + Keyword arguments for the accelerator. + project_kwargs (`dict[str, Any]`, *optional*, defaults to `{}`): + Keyword arguments for the accelerator project config (e.g., `logging_dir`). + tracker_project_name (`str`, *optional*, defaults to `"trl"`): + Name of project to use for tracking. + logdir (`str`, *optional*, defaults to `"logs"`): + Top-level logging directory for checkpoint saving. + num_epochs (`int`, *optional*, defaults to `100`): + Number of epochs to train. + save_freq (`int`, *optional*, defaults to `1`): + Number of epochs between saving model checkpoints. + num_checkpoint_limit (`int`, *optional*, defaults to `5`): + Number of checkpoints to keep before overwriting old ones. + mixed_precision (`str`, *optional*, defaults to `"fp16"`): + Mixed precision training. + allow_tf32 (`bool`, *optional*, defaults to `True`): + Allow `tf32` on Ampere GPUs. + resume_from (`str`, *optional*, defaults to `""`): + Path to resume training from a checkpoint. + sample_num_steps (`int`, *optional*, defaults to `50`): + Number of sampler inference steps. + sample_eta (`float`, *optional*, defaults to `1.0`): + Eta parameter for the DDIM sampler. + sample_guidance_scale (`float`, *optional*, defaults to `5.0`): + Classifier-free guidance weight. + train_batch_size (`int`, *optional*, defaults to `1`): + Batch size for training. + train_use_8bit_adam (`bool`, *optional*, defaults to `False`): + Whether to use the 8bit Adam optimizer from `bitsandbytes`. + train_learning_rate (`float`, *optional*, defaults to `1e-3`): + Learning rate. + train_adam_beta1 (`float`, *optional*, defaults to `0.9`): + Beta1 for Adam optimizer. + train_adam_beta2 (`float`, *optional*, defaults to `0.999`): + Beta2 for Adam optimizer. + train_adam_weight_decay (`float`, *optional*, defaults to `1e-4`): + Weight decay for Adam optimizer. + train_adam_epsilon (`float`, *optional*, defaults to `1e-8`): + Epsilon value for Adam optimizer. + train_gradient_accumulation_steps (`int`, *optional*, defaults to `1`): + Number of gradient accumulation steps. + train_max_grad_norm (`float`, *optional*, defaults to `1.0`): + Maximum gradient norm for gradient clipping. + negative_prompts (`str` or `None`, *optional*, defaults to `None`): + Comma-separated list of prompts to use as negative examples. + truncated_backprop_rand (`bool`, *optional*, defaults to `True`): + If `True`, randomized truncation to different diffusion timesteps is used. + truncated_backprop_timestep (`int`, *optional*, defaults to `49`): + Absolute timestep to which the gradients are backpropagated. Used only if `truncated_backprop_rand=False`. + truncated_rand_backprop_minmax (`tuple[int, int]`, *optional*, defaults to `(0, 50)`): + Range of diffusion timesteps for randomized truncated backpropagation. + push_to_hub (`bool`, *optional*, defaults to `False`): + Whether to push the final model to the Hub. + + """ + 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, + exp_name = '001测试unsloth', + run_name = '', + seed = 3407, + log_with = None, + log_image_freq = 1, + tracker_project_name = 'trl', + logdir = 'logs', + num_epochs = 100, + save_freq = 1, + num_checkpoint_limit = 5, + mixed_precision = 'fp16', + allow_tf32 = True, + resume_from = '', + sample_num_steps = 50, + sample_eta = 1.0, + sample_guidance_scale = 5.0, + train_batch_size = 1, + train_use_8bit_adam = False, + train_learning_rate = 5e-05, + train_adam_beta1 = 0.9, + train_adam_beta2 = 0.999, + train_adam_weight_decay = 0.01, + train_adam_epsilon = 1e-08, + train_gradient_accumulation_steps = 2, + train_max_grad_norm = 1.0, + negative_prompts = None, + truncated_backprop_rand = True, + truncated_backprop_timestep = 49, + push_to_hub = False, + vllm_sampling_params = None, + unsloth_num_chunks = -1, + **kwargs, + ): + + super().__init__( + exp_name = exp_name, + run_name = run_name, + seed = seed, + log_with = log_with, + log_image_freq = log_image_freq, + tracker_project_name = tracker_project_name, + logdir = logdir, + num_epochs = num_epochs, + save_freq = save_freq, + num_checkpoint_limit = num_checkpoint_limit, + mixed_precision = mixed_precision, + allow_tf32 = allow_tf32, + resume_from = resume_from, + sample_num_steps = sample_num_steps, + sample_eta = sample_eta, + sample_guidance_scale = sample_guidance_scale, + train_batch_size = train_batch_size, + train_use_8bit_adam = train_use_8bit_adam, + train_learning_rate = train_learning_rate, + train_adam_beta1 = train_adam_beta1, + train_adam_beta2 = train_adam_beta2, + train_adam_weight_decay = train_adam_weight_decay, + train_adam_epsilon = train_adam_epsilon, + train_gradient_accumulation_steps = train_gradient_accumulation_steps, + train_max_grad_norm = train_max_grad_norm, + negative_prompts = negative_prompts, + truncated_backprop_rand = truncated_backprop_rand, + truncated_backprop_timestep = truncated_backprop_timestep, + push_to_hub = push_to_hub,**kwargs) + self.vllm_sampling_params = vllm_sampling_params + self.unsloth_num_chunks = unsloth_num_chunks +pass + +class _UnslothAlignPropTrainer(PyTorchModelHubMixin): + """""" + + _tag_names = ["trl", "alignprop"] + + def __init__( + self, + config: AlignPropConfig, + reward_function: Callable[[torch.Tensor, tuple[str], tuple[Any]], torch.Tensor], + prompt_function: Callable[[], tuple[str, Any]], + sd_pipeline: DDPOStableDiffusionPipeline, + image_samples_hook: Optional[Callable[[Any, Any, Any], Any]] = None, + ): + if image_samples_hook is None: + warn("No image_samples_hook provided; no images will be logged") + + self.prompt_fn = prompt_function + self.reward_fn = reward_function + self.config = config + self.image_samples_callback = image_samples_hook + + accelerator_project_config = ProjectConfiguration(**self.config.project_kwargs) + + if self.config.resume_from: + self.config.resume_from = os.path.normpath(os.path.expanduser(self.config.resume_from)) + if "checkpoint_" not in os.path.basename(self.config.resume_from): + # get the most recent checkpoint in this directory + checkpoints = list( + filter( + lambda x: "checkpoint_" in x, + os.listdir(self.config.resume_from), + ) + ) + if len(checkpoints) == 0: + raise ValueError(f"No checkpoints found in {self.config.resume_from}") + checkpoint_numbers = sorted([int(x.split("_")[-1]) for x in checkpoints]) + self.config.resume_from = os.path.join( + self.config.resume_from, + f"checkpoint_{checkpoint_numbers[-1]}", + ) + + accelerator_project_config.iteration = checkpoint_numbers[-1] + 1 + + self.accelerator = Accelerator( + log_with=self.config.log_with, + mixed_precision=self.config.mixed_precision, + project_config=accelerator_project_config, + # we always accumulate gradients across timesteps; we want config.train.gradient_accumulation_steps to be the + # number of *samples* we accumulate across, so we need to multiply by the number of training timesteps to get + # the total number of optimizer steps to accumulate across. + gradient_accumulation_steps=self.config.train_gradient_accumulation_steps, + **self.config.accelerator_kwargs, + ) + + is_using_tensorboard = config.log_with is not None and config.log_with == "tensorboard" + + if self.accelerator.is_main_process: + self.accelerator.init_trackers( + self.config.tracker_project_name, + config=dict(alignprop_trainer_config=config.to_dict()) + if not is_using_tensorboard + else config.to_dict(), + init_kwargs=self.config.tracker_kwargs, + ) + + logger.info(f"\n{config}") + + set_seed(self.config.seed, device_specific=True) + + self.sd_pipeline = sd_pipeline + + self.sd_pipeline.set_progress_bar_config( + position=1, + disable=not self.accelerator.is_local_main_process, + leave=False, + desc="Timestep", + dynamic_ncols=True, + ) + + # For mixed precision training we cast all non-trainable weights [vae, non-lora text_encoder and non-lora unet] to half-precision + # as these weights are only used for inference, keeping weights in full precision is not required. + if self.accelerator.mixed_precision == "fp16": + inference_dtype = torch.float16 + elif self.accelerator.mixed_precision == "bf16": + inference_dtype = torch.bfloat16 + else: + inference_dtype = torch.float32 + + self.sd_pipeline.vae.to(self.accelerator.device, dtype=inference_dtype) + self.sd_pipeline.text_encoder.to(self.accelerator.device, dtype=inference_dtype) + self.sd_pipeline.unet.to(self.accelerator.device, dtype=inference_dtype) + + trainable_layers = self.sd_pipeline.get_trainable_layers() + + self.accelerator.register_save_state_pre_hook(self._save_model_hook) + self.accelerator.register_load_state_pre_hook(self._load_model_hook) + + # Enable TF32 for faster training on Ampere GPUs, + # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices + if self.config.allow_tf32: + torch.backends.cuda.matmul.allow_tf32 = True + + self.optimizer = self._setup_optimizer( + trainable_layers.parameters() if not isinstance(trainable_layers, list) else trainable_layers + ) + + self.neg_prompt_embed = self.sd_pipeline.text_encoder( + self.sd_pipeline.tokenizer( + [""] if self.config.negative_prompts is None else self.config.negative_prompts, + return_tensors="pt", + padding="max_length", + truncation=True, + max_length=self.sd_pipeline.tokenizer.model_max_length, + ).input_ids.to(self.accelerator.device) + )[0] + + # NOTE: for some reason, autocast is necessary for non-lora training but for lora training it isn't necessary and it uses + # more memory + self.autocast = self.sd_pipeline.autocast or self.accelerator.autocast + + if hasattr(self.sd_pipeline, "use_lora") and self.sd_pipeline.use_lora: + unet, self.optimizer = self.accelerator.prepare(trainable_layers, self.optimizer) + self.trainable_layers = list(filter(lambda p: p.requires_grad, unet.parameters())) + else: + self.trainable_layers, self.optimizer = self.accelerator.prepare(trainable_layers, self.optimizer) + + if config.resume_from: + logger.info(f"Resuming from {config.resume_from}") + self.accelerator.load_state(config.resume_from) + self.first_epoch = int(config.resume_from.split("_")[-1]) + 1 + else: + self.first_epoch = 0 + + def compute_rewards(self, prompt_image_pairs): + reward, reward_metadata = self.reward_fn( + prompt_image_pairs["images"], prompt_image_pairs["prompts"], prompt_image_pairs["prompt_metadata"] + ) + return reward + + def step(self, epoch: int, global_step: int): + """ + Perform a single step of training. + + Args: + epoch (int): The current epoch. + global_step (int): The current global step. + + Side Effects: + - Model weights are updated + - Logs the statistics to the accelerator trackers. + - If `self.image_samples_callback` is not None, it will be called with the prompt_image_pairs, global_step, + and the accelerator tracker. + + Returns: + global_step (int): The updated global step. + """ + info = defaultdict(list) + + self.sd_pipeline.unet.train() + + for _ in range(self.config.train_gradient_accumulation_steps): + with self.accelerator.accumulate(self.sd_pipeline.unet), self.autocast(), torch.enable_grad(): + prompt_image_pairs = self._generate_samples( + batch_size=self.config.train_batch_size, + ) + + rewards = self.compute_rewards(prompt_image_pairs) + + prompt_image_pairs["rewards"] = rewards + + rewards_vis = self.accelerator.gather(rewards).detach().cpu().numpy() + + loss = self.calculate_loss(rewards) + + self.accelerator.backward(loss) + + if self.accelerator.sync_gradients: + self.accelerator.clip_grad_norm_( + self.trainable_layers.parameters() + if not isinstance(self.trainable_layers, list) + else self.trainable_layers, + self.config.train_max_grad_norm, + ) + + self.optimizer.step() + self.optimizer.zero_grad() + + info["reward_mean"].append(rewards_vis.mean()) + info["reward_std"].append(rewards_vis.std()) + info["loss"].append(loss.item()) + + # Checks if the accelerator has performed an optimization step behind the scenes + if self.accelerator.sync_gradients: + # log training-related stuff + info = {k: torch.mean(torch.tensor(v)) for k, v in info.items()} + info = self.accelerator.reduce(info, reduction="mean") + info.update({"epoch": epoch}) + self.accelerator.log(info, step=global_step) + global_step += 1 + info = defaultdict(list) + else: + raise ValueError( + "Optimization step should have been performed by this point. Please check calculated gradient accumulation settings." + ) + # Logs generated images + if self.image_samples_callback is not None and global_step % self.config.log_image_freq == 0: + self.image_samples_callback(prompt_image_pairs, global_step, self.accelerator.trackers[0]) + + if epoch != 0 and epoch % self.config.save_freq == 0 and self.accelerator.is_main_process: + self.accelerator.save_state() + + return global_step + + def calculate_loss(self, rewards): + """ + Calculate the loss for a batch of an unpacked sample + + Args: + rewards (torch.Tensor): + Differentiable reward scalars for each generated image, shape: [batch_size] + + Returns: + loss (torch.Tensor) (all of these are of shape (1,)) + """ + # Loss is specific to Aesthetic Reward function used in AlignProp (https://huggingface.co/papers/2310.03739) + loss = 10.0 - (rewards).mean() + return loss + + def loss( + self, + advantages: torch.Tensor, + clip_range: float, + ratio: torch.Tensor, + ): + unclipped_loss = -advantages * ratio + clipped_loss = -advantages * torch.clamp( + ratio, + 1.0 - clip_range, + 1.0 + clip_range, + ) + return torch.mean(torch.maximum(unclipped_loss, clipped_loss)) + + def _setup_optimizer(self, trainable_layers_parameters): + if self.config.train_use_8bit_adam: + import bitsandbytes + + optimizer_cls = bitsandbytes.optim.AdamW8bit + else: + optimizer_cls = torch.optim.AdamW + + return optimizer_cls( + trainable_layers_parameters, + lr=self.config.train_learning_rate, + betas=(self.config.train_adam_beta1, self.config.train_adam_beta2), + weight_decay=self.config.train_adam_weight_decay, + eps=self.config.train_adam_epsilon, + ) + + def _save_model_hook(self, models, weights, output_dir): + self.sd_pipeline.save_checkpoint(models, weights, output_dir) + weights.pop() # ensures that accelerate doesn't try to handle saving of the model + + def _load_model_hook(self, models, input_dir): + self.sd_pipeline.load_checkpoint(models, input_dir) + models.pop() # ensures that accelerate doesn't try to handle loading of the model + + def _generate_samples(self, batch_size, with_grad=True, prompts=None): + """ + Generate samples from the model + + Args: + batch_size (int): Batch size to use for sampling + with_grad (bool): Whether the generated RGBs should have gradients attached to it. + + Returns: + prompt_image_pairs (dict[Any]) + """ + prompt_image_pairs = {} + + sample_neg_prompt_embeds = self.neg_prompt_embed.repeat(batch_size, 1, 1) + + if prompts is None: + prompts, prompt_metadata = zip(*[self.prompt_fn() for _ in range(batch_size)]) + else: + prompt_metadata = [{} for _ in range(batch_size)] + + prompt_ids = self.sd_pipeline.tokenizer( + prompts, + return_tensors="pt", + padding="max_length", + truncation=True, + max_length=self.sd_pipeline.tokenizer.model_max_length, + ).input_ids.to(self.accelerator.device) + + prompt_embeds = self.sd_pipeline.text_encoder(prompt_ids)[0] + + if with_grad: + sd_output = self.sd_pipeline.rgb_with_grad( + prompt_embeds=prompt_embeds, + negative_prompt_embeds=sample_neg_prompt_embeds, + num_inference_steps=self.config.sample_num_steps, + guidance_scale=self.config.sample_guidance_scale, + eta=self.config.sample_eta, + truncated_backprop_rand=self.config.truncated_backprop_rand, + truncated_backprop_timestep=self.config.truncated_backprop_timestep, + truncated_rand_backprop_minmax=self.config.truncated_rand_backprop_minmax, + output_type="pt", + ) + else: + sd_output = self.sd_pipeline( + prompt_embeds=prompt_embeds, + negative_prompt_embeds=sample_neg_prompt_embeds, + num_inference_steps=self.config.sample_num_steps, + guidance_scale=self.config.sample_guidance_scale, + eta=self.config.sample_eta, + output_type="pt", + ) + + images = sd_output.images + + prompt_image_pairs["images"] = images + prompt_image_pairs["prompts"] = prompts + prompt_image_pairs["prompt_metadata"] = prompt_metadata + + return prompt_image_pairs + + def train(self, epochs: Optional[int] = None): + """ + Train the model for a given number of epochs + """ + global_step = 0 + if epochs is None: + epochs = self.config.num_epochs + for epoch in range(self.first_epoch, epochs): + global_step = self.step(epoch, global_step) + + def _save_pretrained(self, save_directory): + self.sd_pipeline.save_pretrained(save_directory) + self.create_model_card() + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) + + 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 + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + citation = textwrap.dedent("""\ + @article{prabhudesai2024aligning, + title = {{Aligning Text-to-Image Diffusion Models with Reward Backpropagation}}, + author = {Mihir Prabhudesai and Anirudh Goyal and Deepak Pathak and Katerina Fragkiadaki}, + year = 2024, + eprint = {arXiv:2310.03739} + }""") + + 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="AlignProp", + trainer_citation=citation, + paper_title="Aligning Text-to-Image Diffusion Models with Reward Backpropagation", + paper_id="2310.03739", + ) + + model_card.save(os.path.join(self.args.output_dir, "README.md")) +class UnslothAlignPropTrainer(_UnslothAlignPropTrainer): + """ + + The AlignPropTrainer uses Deep Diffusion Policy Optimization to optimise diffusion models. Note, this trainer is + heavily inspired by the work here: https://github.com/mihirp1998/AlignProp/ As of now only Stable Diffusion based + pipelines are supported + + Attributes: + config (`AlignPropConfig`): + Configuration object for AlignPropTrainer. Check the documentation of `PPOConfig` for more details. + reward_function (`Callable[[torch.Tensor, tuple[str], tuple[Any]], torch.Tensor]`): + Reward function to be used + prompt_function (`Callable[[], tuple[str, Any]]`): + Function to generate prompts to guide model + sd_pipeline (`DDPOStableDiffusionPipeline`): + Stable Diffusion pipeline to be used for training. + image_samples_hook (`Optional[Callable[[Any, Any, Any], Any]]`): + Hook to be called to log images + + """ + def __init__( + self, + config, + reward_function, + prompt_function, + sd_pipeline, + image_samples_hook = None, + **kwargs + ): + if args is None: args = UnslothAlignPropConfig() + other_metrics = [] + + from unsloth_zoo.logging_utils import PatchRLStatistics + PatchRLStatistics('alignprop_trainer', other_metrics) + + super().__init__( + config = config, + reward_function = reward_function, + prompt_function = prompt_function, + sd_pipeline = sd_pipeline, + image_samples_hook = image_samples_hook,**kwargs) + +pass diff --git a/unsloth_compiled_cache/UnslothBCOTrainer.py b/unsloth_compiled_cache/UnslothBCOTrainer.py new file mode 100644 index 0000000..565b69b --- /dev/null +++ b/unsloth_compiled_cache/UnslothBCOTrainer.py @@ -0,0 +1,1823 @@ +""" +2025.7.1 +2025.7.1 +4.53.1 +0.19.1 +__UNSLOTH_VERSIONING__ +""" +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from trl.trainer.bco_trainer import (Any, AutoModelForCausalLM, BCOConfig, BCOTrainer, BaseImageProcessor, CLF_NAME, Callable, DPODataCollatorWithPadding, DataCollator, DataLoader, Dataset, EvalLoopOutput, F, FeatureExtractionMixin, Literal, Optional, PartialState, Path, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, RUNNING_NAME, RunningMoments, SequentialSampler, Trainer, TrainerCallback, TrainingArguments, Union, _process_tokens, _tokenize, autocast, contextmanager, create_reference_model, defaultdict, disable_dropout_in_model, generate_model_card, get_comet_experiment_url, has_length, inspect, is_comet_available, is_joblib_available, is_peft_available, is_sklearn_available, is_wandb_available, itemgetter, log_table_to_comet_experiment, logger, maybe_apply_chat_template, nn, np, nullcontext, os, pad_to_length, pd, peft_module_casting_to_bf16, prepare_deepspeed, prepare_model_for_kbit_training, random, textwrap, torch, tqdm, warnings, F, Optional, PeftModel, PreTrainedModel, Trainer, is_peft_available, logger, os, torch) + + +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 UnslothBCOConfig(BCOConfig): + """ + + Configuration class for the [`BCOTrainer`]. + + This class includes only the parameters that are specific to BCO training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may + differ from those in [`~transformers.TrainingArguments`]. + + 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: + max_length (`int` or `None`, *optional*, defaults to `1024`): + Maximum length of the sequences (prompt + completion) in the batch. This argument is required if you want + to use the default data collator. + max_prompt_length (`int` or `None`, *optional*, defaults to `512`): + Maximum length of the prompt. This argument is required if you want to use the default data collator. + max_completion_length (`int` or `None`, *optional*, defaults to `None`): + Maximum length of the completion. This argument is required if you want to use the default data collator + and your model is an encoder-decoder. + beta (`float`, *optional*, defaults to `0.1`): + Parameter controlling the deviation from the reference model. Higher β means less deviation from the + reference model. + label_pad_token_id (`int`, *optional*, defaults to `-100`): + Label pad token id. This argument is required if you want to use the default data collator. + padding_value (`int` or `None`, *optional*, defaults to `None`): + Padding value to use. If `None`, the padding value of the tokenizer is used. + truncation_mode (`str`, *optional*, defaults to `"keep_end"`): + Truncation mode to use when the prompt is too long. Possible values are `"keep_end"` or `"keep_start"`. + This argument is required if you want to use the default data collator. + disable_dropout (`bool`, *optional*, defaults to `True`): + Whether to disable dropout in the model and reference model. + generate_during_eval (`bool`, *optional*, defaults to `False`): + If `True`, generates and logs completions from both the model and the reference model to W&B or Comet + during evaluation. + is_encoder_decoder (`bool` or `None`, *optional*, defaults to `None`): + When using the `model_init` argument (callable) to instantiate the model instead of the `model` argument, + you need to specify if the model returned by the callable is an encoder-decoder model. + precompute_ref_log_probs (`bool`, *optional*, defaults to `False`): + Whether to precompute reference model log probabilities for training and evaluation datasets. This is + useful when training without the reference model to reduce the total GPU memory needed. + model_init_kwargs (`dict[str, Any]` or `None`, *optional*, defaults to `None`): + Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model from a + string. + ref_model_init_kwargs (`dict[str, Any]` or `None`, *optional*, defaults to `None`): + Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the reference model + from a string. + dataset_num_proc (`int` or `None`, *optional*, defaults to `None`): + Number of processes to use for processing the dataset. + prompt_sample_size (`int`, *optional*, defaults to `1024`): + Number of prompts that are fed to density ratio classifier. + min_density_ratio (`float`, *optional*, defaults to `0.5`): + Minimum value of the density ratio. The estimated density ratio is clamped to this value. + max_density_ratio (`float`, *optional*, defaults to `10.0`): + Maximum value of the density ratio. The estimated density ratio is clamped to this value. + + """ + 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, + hub_revision = None, + 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, + liger_kernel_config = None, + eval_use_gather_object = False, + average_tokens_across_devices = False, + max_length = 1024, + max_prompt_length = 512, + max_completion_length = None, + beta = 0.1, + label_pad_token_id = -100, + padding_value = None, + truncation_mode = 'keep_end', + disable_dropout = True, + generate_during_eval = False, + is_encoder_decoder = None, + precompute_ref_log_probs = False, + model_init_kwargs = None, + ref_model_init_kwargs = None, + dataset_num_proc = None, + prompt_sample_size = 1024, + min_density_ratio = 0.5, + max_density_ratio = 10.0, + 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, + 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, + max_length = max_length, + max_prompt_length = max_prompt_length, + max_completion_length = max_completion_length, + beta = beta, + label_pad_token_id = label_pad_token_id, + padding_value = padding_value, + truncation_mode = truncation_mode, + disable_dropout = disable_dropout, + generate_during_eval = generate_during_eval, + is_encoder_decoder = is_encoder_decoder, + precompute_ref_log_probs = precompute_ref_log_probs, + model_init_kwargs = model_init_kwargs, + ref_model_init_kwargs = ref_model_init_kwargs, + dataset_num_proc = dataset_num_proc, + prompt_sample_size = prompt_sample_size, + min_density_ratio = min_density_ratio, + max_density_ratio = max_density_ratio,**kwargs) + self.vllm_sampling_params = vllm_sampling_params + self.unsloth_num_chunks = unsloth_num_chunks +pass + +class _UnslothBCOTrainer(Trainer): + r"""""" + + _tag_names = ["trl", "bco"] + + def __init__( + self, + model: Union[PreTrainedModel, nn.Module, str] = None, + ref_model: Optional[Union[PreTrainedModel, nn.Module, str]] = None, + args: BCOConfig = None, + train_dataset: Optional[Dataset] = None, + eval_dataset: Optional[Union[Dataset, dict[str, Dataset]]] = None, + processing_class: Optional[ + Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin] + ] = None, + data_collator: Optional[DataCollator] = None, + model_init: Optional[Callable[[], PreTrainedModel]] = 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, + compute_metrics: Optional[Callable[[EvalLoopOutput], dict]] = None, + model_adapter_name: Optional[str] = None, + ref_adapter_name: Optional[str] = None, + embedding_func: Optional[Callable] = None, + embedding_tokenizer: Optional[PreTrainedTokenizerBase] = None, + ): + if embedding_func is not None and not (is_sklearn_available() and is_joblib_available()): + raise ImportError( + "BCOTrainer with UDM requires the scikit-learn and joblib libraries. Please install it with `pip install scikit-learn joblib`." + ) + + if type(args) is TrainingArguments: + raise ValueError("Please use `BCOConfig` instead `TrainingArguments`.") + + if not isinstance(model, str) and model is not None 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_kwargs to the BCOTrainer. 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 BCOConfig. 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_kwargs to the BCOTrainer. 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 BCOConfig. 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 with `pip install peft` 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 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 args.gradient_checkpointing: + # 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 args.gradient_checkpointing: + # 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." + ) + + if model is not None: + self.is_encoder_decoder = model.config.is_encoder_decoder + elif args.is_encoder_decoder is None: + raise ValueError("When no model is provided, you need to pass the parameter is_encoder_decoder.") + else: + self.is_encoder_decoder = args.is_encoder_decoder + + self.is_peft_model = is_peft_available() and isinstance(model, PeftModel) + self.model_adapter_name = model_adapter_name + self.ref_adapter_name = ref_adapter_name + + 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( + "max_length or a processing_class must be specified when using the default DPODataCollatorWithPadding" + ) + if args.max_length is None: + warnings.warn( + "When using DPODataCollatorWithPadding, you should set `max_length` in the `BCOConfig`. " + "It will be set to `512` by default, but you should do it yourself in the future.", + UserWarning, + ) + max_length = 512 + if args.max_length is not None: + max_length = args.max_length + + if args.max_prompt_length is None: + warnings.warn( + "When using DPODataCollatorWithPadding, you should set `max_prompt_length` in the `BCOConfig`. " + "It will be set to `128` by default, but you should do it yourself in the future.", + UserWarning, + ) + max_prompt_length = 128 + if args.max_prompt_length is not None: + max_prompt_length = args.max_prompt_length + + max_completion_length = None + if args.max_completion_length is None and self.is_encoder_decoder: + warnings.warn( + "When using DPODataCollatorWithPadding with an encoder decoder architecture, you should set `max_completion_length` in the BCOTrainer's init" + " it will be set to `128` by default, but you should do it yourself in the future.", + UserWarning, + ) + max_completion_length = 128 + if args.max_completion_length is not None and self.is_encoder_decoder: + max_completion_length = args.max_completion_length + + if data_collator is None: + data_collator = DPODataCollatorWithPadding( + pad_token_id=processing_class.pad_token_id, + label_pad_token_id=args.label_pad_token_id, + is_encoder_decoder=self.is_encoder_decoder, + ) + + if args.remove_unused_columns: + args.remove_unused_columns = False + # warn users + warnings.warn( + "When using DPODataCollatorWithPadding, you should set `remove_unused_columns=False` in your BCOConfig" + " we have set it for you, but you should do it yourself in the future.", + UserWarning, + ) + + self.use_dpo_data_collator = True + else: + self.use_dpo_data_collator = False + + # 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.max_length = max_length + self.generate_during_eval = args.generate_during_eval + self.label_pad_token_id = args.label_pad_token_id + self.padding_value = args.padding_value if args.padding_value is not None else processing_class.pad_token_id + self.max_prompt_length = max_prompt_length + self.truncation_mode = args.truncation_mode + self.max_completion_length = max_completion_length + self.precompute_ref_log_probs = args.precompute_ref_log_probs + + # 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 + + # metric + self._stored_metrics = defaultdict(lambda: defaultdict(list)) + + # BCO parameter + self.beta = args.beta + self.aux_loss_enabled = getattr(model.config, "output_router_logits", False) + 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, + ) + + # Underlying Distribution Matching argument + self.embedding_func = embedding_func + self.embedding_tokenizer = embedding_tokenizer + + # 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 BCO, the sampled data does not include the + # "input_ids" key. Instead, the available keys are "prompt_input_ids" and "completion_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 + + with PartialState().main_process_first(): + # Apply the chat template if needed + train_dataset = train_dataset.map( + maybe_apply_chat_template, fn_kwargs={"tokenizer": processing_class}, num_proc=args.dataset_num_proc + ) + if eval_dataset is not None: + eval_dataset = eval_dataset.map( + maybe_apply_chat_template, + fn_kwargs={"tokenizer": processing_class}, + num_proc=args.dataset_num_proc, + ) + + # Tokenize and prepare the training datasets + train_dataset = train_dataset.map( + _tokenize, + batched=True, + fn_kwargs={"tokenizer": processing_class, "embedding_tokenizer": self.embedding_tokenizer}, + num_proc=args.dataset_num_proc, + desc="Tokenizing train dataset", + ) + + # Prepare the datasets + fn_kwargs = { + "prefix": "", + "is_encoder_decoder": self.is_encoder_decoder, + "tokenizer": processing_class, + "max_length": self.max_length, + "truncation_mode": self.truncation_mode, + "label_pad_token_id": self.label_pad_token_id, + "max_prompt_length": self.max_prompt_length, + "max_completion_length": self.max_completion_length, + } + train_dataset = train_dataset.map( + _process_tokens, + fn_kwargs=fn_kwargs, + num_proc=args.dataset_num_proc, + desc="Processing tokenized train dataset", + ) + + if eval_dataset is not None: + # Tokenize + eval_dataset = eval_dataset.map( + _tokenize, + fn_kwargs={"tokenizer": processing_class, "embedding_tokenizer": self.embedding_tokenizer}, + batched=True, + num_proc=args.dataset_num_proc, + desc="Tokenizing eval dataset", + ) + + # Process + fn_kwargs = { + "prefix": "", + "is_encoder_decoder": self.is_encoder_decoder, + "tokenizer": processing_class, + "max_length": self.max_length, + "truncation_mode": self.truncation_mode, + "label_pad_token_id": self.label_pad_token_id, + "max_prompt_length": self.max_prompt_length, + "max_completion_length": self.max_completion_length, + } + eval_dataset = eval_dataset.map( + _process_tokens, + fn_kwargs=fn_kwargs, + num_proc=args.dataset_num_proc, + desc="Processing tokenized eval dataset", + ) + + desirable = train_dataset.filter( + lambda x: x["label"], num_proc=args.dataset_num_proc, desc="Filtering desirable examples" + ) + undesirable = train_dataset.filter( + lambda x: not x["label"], num_proc=args.dataset_num_proc, desc="Filtering undesirable examples" + ) + + 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`" + ) + else: + if self.is_deepspeed_enabled: + self.ref_model = prepare_deepspeed(self.ref_model, self.accelerator) + else: + self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True) + + self.running = RunningMoments(accelerator=self.accelerator) + + if self.embedding_func is None or args.resume_from_checkpoint: + return + + chosen_embeddings = self._get_sample_prompt_embeddings(desirable, sample_size=self.args.prompt_sample_size) + rejected_embeddings = self._get_sample_prompt_embeddings(undesirable, sample_size=self.args.prompt_sample_size) + + embeddings = torch.cat((chosen_embeddings, rejected_embeddings), dim=0) + labels = torch.cat( + (torch.ones_like(chosen_embeddings[:, 0]), torch.zeros_like(rejected_embeddings[:, 0])), dim=0 + ) + + self.clf = LogisticRegression(class_weight="balanced").fit( + embeddings.cpu().float().numpy(), labels.cpu().numpy() + ) + chosen_mean = self.clf.score( + chosen_embeddings.cpu().float().numpy(), torch.ones_like(chosen_embeddings[:, 0]).cpu().numpy() + ) + rejected_mean = self.clf.score( + rejected_embeddings.cpu().float().numpy(), torch.zeros_like(rejected_embeddings[:, 0]).cpu().numpy() + ) + logger.info(f"UDM classifier training scores: chosen: {chosen_mean}, rejected: {rejected_mean}") + + @property + def match_underlying_distribution(self): + return self.embedding_func is not None and self.embedding_tokenizer is not None + + def _get_chosen_prob(self, prompt_embeddings: torch.FloatTensor) -> torch.FloatTensor: + """ + Calculates the probability if the given prompt embedding is from desirable dataset. This function calculates + the probability in the process and ensemble across processes. + """ + dtype = prompt_embeddings.dtype + device = prompt_embeddings.device + rank = self.accelerator.process_index + + padded_prompt_embeddings = self.accelerator.pad_across_processes( + prompt_embeddings, pad_index=self.embedding_tokenizer.pad_token_id + ) + sample_size = padded_prompt_embeddings.shape[0] + nonzero = padded_prompt_embeddings.mean(dim=1) != self.embedding_tokenizer.pad_token_id + prompt_embeddings = self.accelerator.gather(padded_prompt_embeddings) + + # cannot predict for all empty values + if prompt_embeddings.shape[0] == 0: + return torch.tensor([], device=device, dtype=dtype) + + prob = self.clf.predict_proba(prompt_embeddings.cpu().float().numpy())[:, 1] + prob = torch.as_tensor(prob, dtype=dtype, device=device) + prob = self.accelerator.reduce(prob, reduction="mean") + + prob = prob[sample_size * rank : sample_size * (rank + 1)] + prob = prob[nonzero] + + return prob + + def _vectorize_prompt(self, input_ids: torch.LongTensor, attention_mask: torch.LongTensor) -> torch.FloatTensor: + """ + Replaces processing_class.pad_token_id to embedding_tokenizer.pad_token_id and applies self.embedding_func + """ + input_ids = torch.where( + input_ids == self.processing_class.pad_token_id, + self.embedding_tokenizer.pad_token_id, + input_ids, + ) + + with torch.no_grad(): + embeddings = self.embedding_func( + input_ids=input_ids, + attention_mask=attention_mask, + ) + + return embeddings + + def _get_prompt_embeddings( + self, batch: dict[str, Union[list, torch.LongTensor]] + ) -> tuple[torch.FloatTensor, torch.FloatTensor]: + """Extract embeddings from frozen embedding model""" + + if not self.match_underlying_distribution: + return None, None + + embeddings = self._vectorize_prompt( + input_ids=batch["embedding_input_ids"], + attention_mask=batch["embedding_attention_mask"], + ) + + chosen_idx = [i for i in range(len(batch["label"])) if batch["label"][i] is True] + rejected_idx = [i for i in range(len(batch["label"])) if batch["label"][i] is False] + + chosen_embeddings = embeddings[chosen_idx, ...] + rejected_embeddings = embeddings[rejected_idx, ...] + + return (chosen_embeddings, rejected_embeddings) + + def _get_sample_prompt_embeddings(self, dataset: Dataset, sample_size: int = 512) -> torch.FloatTensor: + """ + Sample instances from dataset and get prompt embeddings. Used for density ratio classifier training. + """ + n_samples = min(len(dataset), sample_size) + rand_indices = np.random.choice(len(dataset), size=(n_samples,)) + + embedding_dataset = dataset.select(rand_indices) + + dataloader_params = { + "batch_size": self.args.per_device_train_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(embedding_dataset, **dataloader_params)) + + with torch.no_grad(): + all_embeddings = torch.empty(0) + for padded_batch in tqdm(iterable=data_loader, desc="Building sample prompt embeddings"): + embeddings = self._vectorize_prompt( + input_ids=padded_batch["embedding_input_ids"], + attention_mask=padded_batch["embedding_attention_mask"], + ) + embeddings = self.accelerator.gather_for_metrics(embeddings) + all_embeddings = torch.cat((all_embeddings, embeddings.cpu())) + + return all_embeddings + + def _save_optimizer_and_scheduler(self, output_dir): + output_dir = output_dir if output_dir is not None else self.args.output_dir + super()._save_optimizer_and_scheduler(output_dir) + + if self.accelerator.is_main_process: + # When saving optimizer and scheduler to checkpoint, save also the running delta object. + self.running.save_to_json(os.path.join(output_dir, RUNNING_NAME)) + + if self.match_underlying_distribution: + joblib.dump(self.clf, os.path.join(output_dir, CLF_NAME), compress=True) + + def _load_optimizer_and_scheduler(self, checkpoint): + if checkpoint is None: + logger.warning_once(f"Missing Checkpoint {checkpoint}") + return + + super()._load_optimizer_and_scheduler(checkpoint) + + # when loading optimizer and scheduler from checkpoint, also load the running delta object. + running_file = os.path.join(checkpoint, RUNNING_NAME) + if os.path.isfile(running_file): + self.running = RunningMoments.load_from_json(self.accelerator, running_file) + + if self.match_underlying_distribution: + clf_file = os.path.join(checkpoint, CLF_NAME) + if os.path.isfile(clf_file): + self.clf = joblib.load(clf_file) + + @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 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: + dataloader_params = { + "batch_size": self.args.per_device_train_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)) + reference_completion_logps = [] + + for padded_batch in tqdm(iterable=data_loader, desc="Train dataset reference log probs"): + reference_completion_logp = self.compute_reference_log_probs(padded_batch) + + reference_completion_logp = self.accelerator.gather_for_metrics(reference_completion_logp) + reference_completion_logps.append(reference_completion_logp.cpu()) + + self.train_dataset = self.train_dataset.add_column( + name="reference_logps", column=torch.cat(reference_completion_logps).float().numpy() + ) + + 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: + dataloader_params = { + "batch_size": self.args.per_device_eval_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)) + + reference_completion_logps = [] + + for padded_batch in tqdm(iterable=data_loader, desc="Eval dataset reference log probs"): + reference_completion_logp = self.compute_reference_log_probs(padded_batch) + + reference_completion_logp = self.accelerator.gather_for_metrics(reference_completion_logp) + reference_completion_logps.append(reference_completion_logp.cpu()) + + eval_dataset = eval_dataset.add_column( + name="reference_logps", column=torch.cat(reference_completion_logps).float().numpy() + ) + + # Save calculated reference_chosen_logps and reference_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) + + def compute_reference_log_probs(self, padded_batch: dict) -> dict: + """Computes log probabilities of the reference model for a single padded batch of a BCO specific dataset.""" + with torch.no_grad(): + if self.ref_model is None: + with self.null_ref_context(): + if self.is_encoder_decoder: + completion_logits = self.model( + padded_batch["prompt_input_ids"], + attention_mask=padded_batch["prompt_attention_mask"], + decoder_input_ids=padded_batch.get("completion_decoder_input_ids"), + labels=padded_batch["completion_labels"], + ).logits + + else: + completion_logits = self.model( + padded_batch["completion_input_ids"], + attention_mask=padded_batch["completion_attention_mask"], + ).logits + + else: + if self.is_encoder_decoder: + completion_logits = self.ref_model( + padded_batch["prompt_input_ids"], + attention_mask=padded_batch["prompt_attention_mask"], + decoder_input_ids=padded_batch.get("completion_decoder_input_ids"), + labels=padded_batch["completion_labels"], + ).logits + + else: + completion_logits = self.ref_model( + padded_batch["completion_input_ids"], attention_mask=padded_batch["completion_attention_mask"] + ).logits + + completion_logps = self.get_batch_logps( + completion_logits, + padded_batch["completion_labels"], + average_log_prob=False, + is_encoder_decoder=self.is_encoder_decoder, + label_pad_token_id=self.label_pad_token_id, + ) + + return completion_logps + + @staticmethod + def get_batch_logps( + logits: torch.FloatTensor, + labels: torch.LongTensor, + average_log_prob: bool = False, + label_pad_token_id: int = -100, + is_encoder_decoder: bool = False, + ) -> torch.FloatTensor: + """Compute the log probabilities of the given labels under the given logits. + + Args: + logits: Logits of the model (unnormalized). Shape: (batch_size, sequence_length, vocab_size) + labels: + Labels for which to compute the log probabilities. Label tokens with a value of label_pad_token_id are + ignored. Shape: (batch_size, sequence_length) + average_log_prob: + If True, return the average log probability per (non-masked) token. Otherwise, return the sum of the + log probabilities of the (non-masked) tokens. + + Returns: + A tensor of shape (batch_size,) containing the average/sum log probabilities of the given labels under the + given logits. + """ + if logits.shape[:-1] != labels.shape: + raise ValueError("Logits (batch and sequence length dim) and labels must have the same shape.") + + if not is_encoder_decoder: + labels = labels[:, 1:].clone() + logits = logits[:, :-1, :] + else: + # Fixes end-dec RuntimeError + labels = labels.clone() + + loss_mask = labels != label_pad_token_id + + # dummy token; we'll ignore the losses on these tokens later + labels[labels == label_pad_token_id] = 0 + + per_token_logps = selective_log_softmax(logits, labels) + + if average_log_prob: + return (per_token_logps * loss_mask).sum(-1) / loss_mask.sum(-1) + else: + return (per_token_logps * loss_mask).sum(-1) + + def forward( + self, model: nn.Module, batch: dict[str, Union[list, torch.LongTensor]] + ) -> tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: + model_kwargs = ( + { + "labels": batch["completion_labels"], + "decoder_input_ids": batch.get("completion_decoder_input_ids"), + } + if self.is_encoder_decoder + else {} + ) + if self.aux_loss_enabled: + model_kwargs["output_router_logits"] = True + + outputs = model( + batch["completion_input_ids"], + attention_mask=batch["completion_attention_mask"], + **model_kwargs, + ) + completion_logits = outputs.logits + + completion_logps = self.get_batch_logps( + completion_logits, + batch["completion_labels"], + average_log_prob=False, + is_encoder_decoder=self.is_encoder_decoder, + label_pad_token_id=self.label_pad_token_id, + ) + + if completion_logps.shape[0] != len(batch["label"]): + raise ValueError( + "There is a mismatch between the number of examples in this batch and the number of " + "examples for which an output sequence was predicted." + ) + + chosen_idx = [i for i in range(completion_logps.shape[0]) if batch["label"][i] is True] + rejected_idx = [i for i in range(completion_logps.shape[0]) if batch["label"][i] is False] + + chosen_logps = completion_logps[chosen_idx, ...] + rejected_logps = completion_logps[rejected_idx, ...] + + chosen_logits = completion_logits[chosen_idx, ...] + rejected_logits = completion_logits[rejected_idx, ...] + + if self.aux_loss_enabled: + return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, outputs.aux_loss) + else: + return (chosen_logps, rejected_logps, chosen_logits, rejected_logits) + + def _get_udm_weight(self, rejected_embeddings: torch.FloatTensor) -> torch.FloatTensor: + prob_desirable = self._get_chosen_prob(rejected_embeddings) + min_ratio = self.args.min_density_ratio + max_ratio = self.args.max_density_ratio + + weight = (prob_desirable / (1 - prob_desirable + 1e-8)).clamp(min=min_ratio, max=max_ratio) + + return weight + + def bco_loss( + self, + policy_chosen_logps: torch.FloatTensor, + policy_rejected_logps: torch.FloatTensor, + reference_chosen_logps: torch.FloatTensor, + reference_rejected_logps: torch.FloatTensor, + chosen_embeddings: Optional[torch.FloatTensor], + rejected_embeddings: Optional[torch.FloatTensor], + do_train: bool = True, + ) -> tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: + """Compute the BCO loss for a batch of policy and reference model log probabilities. + + Args: + policy_chosen_logps: + Log probabilities of the policy model for the chosen responses. Shape: (num(chosen) in batch_size,) + policy_rejected_logps: + Log probabilities of the policy model for the rejected responses. Shape: (num(rejected) in batch_size,) + reference_chosen_logps: + Log probabilities of the reference model for the chosen responses. Shape: (num(chosen) in batch_size,) + reference_rejected_logps: + Log probabilities of the reference model for the rejected responses. Shape: (num(rejected) in + batch_size,) + chosen_embeddings: embeddings of desirable prompts + rejected_embeddings: embeddings of undesirable prompts + + Returns: + A tuple of four tensors: (losses, chosen_rewards, rejected_rewards, delta). The losses tensor contains the + BCO loss for each example in the batch. The chosen_rewards and rejected_rewards tensors contain the rewards + for the chosen and rejected responses, respectively. The delta value contains the moving average of all + implicit rewards. + """ + + chosen_logratios = policy_chosen_logps - reference_chosen_logps + chosen_rewards = self.beta * chosen_logratios + + rejected_logratios = policy_rejected_logps - reference_rejected_logps + rejected_rewards = self.beta * rejected_logratios + + if do_train: + self.running.update(torch.cat((chosen_rewards, rejected_rewards), 0).detach()) + delta = torch.as_tensor(self.running.mean, device=chosen_rewards.device) + + chosen_losses = -F.logsigmoid(chosen_rewards - delta) + rejected_losses = -F.logsigmoid(-(rejected_rewards - delta)) + + if self.match_underlying_distribution: + chosen_weight = torch.ones_like(chosen_losses) + rejected_weight = self._get_udm_weight(rejected_embeddings) + + losses = torch.cat((chosen_weight * chosen_losses, rejected_weight * rejected_losses), dim=0) + else: + losses = torch.cat((chosen_losses, rejected_losses), dim=0) + + return losses, chosen_rewards, rejected_rewards, delta + + def get_batch_loss_metrics( + self, + model, + batch: dict[str, Union[list, torch.LongTensor]], + do_train: bool = True, + ): + """Compute the BCO loss and other metrics for the given batch of inputs for train or test.""" + metrics = {} + batch = {k: (v.to(self.accelerator.device) if isinstance(v, torch.Tensor) else v) for k, v in batch.items()} + + forward_output = self.forward(model, batch) + ( + policy_chosen_logps, + policy_rejected_logps, + policy_chosen_logits, + policy_rejected_logits, + ) = forward_output[:4] + if self.aux_loss_enabled: + aux_loss = forward_output[4] + + # if reference_logps in batch use them, otherwise use the reference model + if "reference_logps" in batch: + chosen_idx = [i for i in range(batch["reference_logps"].shape[0]) if batch["label"][i] is True] + rejected_idx = [i for i in range(batch["reference_logps"].shape[0]) if batch["label"][i] is False] + + reference_chosen_logps = batch["reference_logps"][chosen_idx, ...] + reference_rejected_logps = batch["reference_logps"][rejected_idx, ...] + else: + with torch.no_grad(): + if self.ref_model is None: + with self.null_ref_context(): + ( + reference_chosen_logps, + reference_rejected_logps, + _, + _, + ) = self.forward(self.model, batch)[:4] + else: + ( + reference_chosen_logps, + reference_rejected_logps, + _, + _, + ) = self.forward(self.ref_model, batch)[:4] + + chosen_embeddings, rejected_embeddings = self._get_prompt_embeddings(batch) + + losses, chosen_rewards, rejected_rewards, delta = self.bco_loss( + policy_chosen_logps, + policy_rejected_logps, + reference_chosen_logps, + reference_rejected_logps, + chosen_embeddings, + rejected_embeddings, + do_train=do_train, + ) + metrics["delta"] = self.accelerator.gather_for_metrics(delta).mean().item() + + num_chosen = torch.Tensor([len(chosen_rewards)]).to(self.accelerator.device) + num_rejected = torch.Tensor([len(rejected_rewards)]).to(self.accelerator.device) + + all_num_chosen = self.accelerator.gather_for_metrics(num_chosen).sum().item() + all_num_rejected = self.accelerator.gather_for_metrics(num_rejected).sum().item() + + if all_num_chosen > 0: + metrics["rewards/chosen_sum"] = ( + self.accelerator.gather_for_metrics(chosen_rewards.nansum()).nansum().item() + ) + metrics["logps/chosen_sum"] = ( + self.accelerator.gather_for_metrics(policy_chosen_logps.nansum()).nansum().item() + ) + metrics["logits/chosen_sum"] = ( + self.accelerator.gather_for_metrics(policy_chosen_logits.nansum()).nansum().item() + ) + metrics["count/chosen"] = all_num_chosen + + if all_num_rejected > 0: + metrics["rewards/rejected_sum"] = ( + self.accelerator.gather_for_metrics(rejected_rewards.nansum()).nansum().item() + ) + metrics["logps/rejected_sum"] = ( + self.accelerator.gather_for_metrics(policy_rejected_logps.nansum()).nansum().item() + ) + metrics["logits/rejected_sum"] = ( + self.accelerator.gather_for_metrics(policy_rejected_logits.nansum()).nansum().item() + ) + metrics["count/rejected"] = all_num_rejected + + loss = losses.nanmean() + if self.aux_loss_enabled: + loss += self.aux_loss_coef * aux_loss + + return loss, 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]]]: + compute_loss_context_manager = ( + autocast(self.accelerator.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) + + # 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 + if self.accelerator.is_main_process: + self.store_metrics(metrics, train_eval="train") + + if return_outputs: + return (loss, metrics) + return loss + + 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 _get_train_sampler(self, dataset: Optional[Dataset] = None) -> Optional[torch.utils.data.Sampler]: + if dataset is None: + dataset = self.train_dataset + if dataset is None or not has_length(dataset): + return None + return SequentialSampler(dataset) + + 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. + generate_context_manager = ( + autocast(self.accelerator.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.processing_class.pad_token_id, + ) + + # if reference_output in batch use that otherwise use the reference model + if "reference_output" in batch: + reference_output = batch["reference_output"] + else: + if self.ref_model is None: + with self.null_ref_context(): + reference_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.processing_class.pad_token_id, + ) + else: + reference_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.processing_class.pad_token_id, + ) + + policy_output = pad_to_length(policy_output, self.max_length, self.processing_class.pad_token_id) + policy_output_decoded = self.processing_class.batch_decode(policy_output, skip_special_tokens=True) + + reference_output = pad_to_length(reference_output, self.max_length, self.processing_class.pad_token_id) + reference_output_decoded = self.processing_class.batch_decode(reference_output, skip_special_tokens=True) + + return policy_output_decoded, reference_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 = [] + + prediction_context_manager = ( + autocast(self.accelerator.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, do_train=False) + + # force log the metrics + if self.accelerator.is_main_process: + 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 = {} + if "logits/chosen_sum" in metrics: + logits_dict["eval_logits/chosen"] = metrics["logits/chosen_sum"] + if "logits/rejected_sum" in metrics: + logits_dict["eval_logits/rejected"] = metrics["logits/rejected_sum"] + 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 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) + + target_indicies = [i for i in range(len(random_batch["label"])) if random_batch["label"][i] is False] + target_batch = { + "prompt_input_ids": random_batch["prompt_input_ids"][target_indicies], + "prompt_attention_mask": random_batch["prompt_attention_mask"][target_indicies], + "prompt": itemgetter(*target_indicies)(random_batch["prompt"]), + } + policy_output_decoded, ref_output_decoded = self.generate_from_model_and_ref(self.model, target_batch) + + table = pd.DataFrame( + columns=["Prompt", "Policy", "Ref Model"], + data=[ + [prompt, pol[len(prompt) :], ref[len(prompt) :]] + for prompt, pol, ref in zip(target_batch["prompt"], policy_output_decoded, ref_output_decoded) + ], + ) + if "wandb" in self.args.report_to: + 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" + # train metrics should have no prefix, eval should have 'eval_' + prefix = "eval_" if train_eval == "eval" else "" + # accumulate average metrics from sums and lengths + for split in ["chosen", "rejected"]: + if f"count/{split}" in self._stored_metrics[train_eval]: + count_sum = torch.Tensor(self._stored_metrics[train_eval][f"count/{split}"]).sum().item() + for metric in ["rewards", "logps", "logits"]: + logs[f"{prefix}{metric}/{split}"] = ( + torch.Tensor(self._stored_metrics[train_eval][f"{metric}/{split}_sum"]).sum().item() + / count_sum + ) + # delete obsolete metric + del self._stored_metrics[train_eval][f"{metric}/{split}_sum"] + del self._stored_metrics[train_eval][f"count/{split}"] + # calculate reward margin + if f"{prefix}rewards/chosen" in logs and f"{prefix}rewards/rejected" in logs: + logs[f"{prefix}rewards/margins"] = logs[f"{prefix}rewards/chosen"] - logs[f"{prefix}rewards/rejected"] + # Add averaged stored metrics to logs + for key, metrics in self._stored_metrics[train_eval].items(): + logs[f"{prefix}{key}"] = torch.Tensor(metrics).mean().item() + del self._stored_metrics[train_eval] + return super().log(logs, start_time) + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) + + 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 + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + citation = textwrap.dedent("""\ + @article{jung2024binary, + title = {{Binary Classifier Optimization for Large Language Model Alignment}}, + author = {Seungjae Jung and Gunsoo Han and Daniel Wontae Nam and Kyoung{-}Woon On}, + year = 2024, + eprint = {arXiv:2404.04656} + }""") + + 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="BCO", + trainer_citation=citation, + paper_title="Binary Classifier Optimization for Large Language Model Alignment", + paper_id="2404.04656", + ) + + model_card.save(os.path.join(self.args.output_dir, "README.md")) +class UnslothBCOTrainer(_UnslothBCOTrainer): + """ + + Initialize BCOTrainer from [BCO](https://huggingface.co/papers/2404.04656) paper. + + 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 (`BCOConfig`): + The arguments to use for training. + 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. + data_collator (`transformers.DataCollator`, *optional*, defaults to `None`): + 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. + model_init (`Callable[[], transformers.PreTrainedModel]`): + The model initializer to use for training. If None is specified, the default model initializer will be + used. + 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. + 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. + model_adapter_name (`str`, defaults to `None`): + Name of the train target PEFT adapter, when using LoRA with multiple adapters. + ref_adapter_name (`str`, defaults to `None`): + Name of the reference PEFT adapter, when using LoRA with multiple adapters. + + """ + def __init__( + self, + model = None, + ref_model = None, + args = None, + train_dataset = None, + eval_dataset = None, + processing_class = None, + data_collator = None, + model_init = None, + callbacks = None, + preprocess_logits_for_metrics = None, + peft_config = None, + compute_metrics = None, + model_adapter_name = None, + ref_adapter_name = None, + embedding_func = None, + embedding_tokenizer = None, + **kwargs + ): + if args is None: args = UnslothBCOConfig() + 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 + 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) + 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 + 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('bco_trainer', other_metrics) + + super().__init__( + model = model, + ref_model = ref_model, + args = args, + train_dataset = train_dataset, + eval_dataset = eval_dataset, + processing_class = processing_class, + data_collator = data_collator, + model_init = model_init, + callbacks = callbacks, + preprocess_logits_for_metrics = preprocess_logits_for_metrics, + peft_config = peft_config, + compute_metrics = compute_metrics, + model_adapter_name = model_adapter_name, + ref_adapter_name = ref_adapter_name, + embedding_func = embedding_func, + embedding_tokenizer = embedding_tokenizer,**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 diff --git a/unsloth_compiled_cache/UnslothCPOTrainer.py b/unsloth_compiled_cache/UnslothCPOTrainer.py new file mode 100644 index 0000000..265c943 --- /dev/null +++ b/unsloth_compiled_cache/UnslothCPOTrainer.py @@ -0,0 +1,1596 @@ +""" +2025.7.1 +2025.7.1 +4.53.1 +0.19.1 +__UNSLOTH_VERSIONING__ +""" +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from trl.trainer.cpo_trainer import (Any, AutoModelForCausalLM, BaseImageProcessor, CPOConfig, CPOTrainer, Callable, DPODataCollatorWithPadding, DataCollator, DataLoader, Dataset, EvalLoopOutput, F, FeatureExtractionMixin, Literal, Optional, PartialState, Path, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, Trainer, TrainerCallback, Union, add_bos_token_if_needed, add_eos_token_if_needed, autocast, defaultdict, disable_dropout_in_model, generate_model_card, get_comet_experiment_url, inspect, is_comet_available, is_peft_available, is_torch_fx_proxy, is_wandb_available, log_table_to_comet_experiment, maybe_apply_chat_template, maybe_extract_prompt, nn, np, nullcontext, os, pad_to_length, pd, peft_module_casting_to_bf16, prepare_model_for_kbit_training, random, textwrap, torch, warnings, F, Optional, PeftModel, PreTrainedModel, Trainer, is_peft_available, os, torch) + + +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 UnslothCPOConfig(CPOConfig): + """ + + Configuration class for the [`CPOTrainer`]. + + This class includes only the parameters that are specific to CPO training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may + differ from those in [`~transformers.TrainingArguments`]. + + 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: + max_length (`int` or `None`, *optional*, defaults to `1024`): + Maximum length of the sequences (prompt + completion) in the batch. This argument is required if you want + to use the default data collator. + max_prompt_length (`int` or `None`, *optional*, defaults to `512`): + Maximum length of the prompt. This argument is required if you want to use the default data collator. + max_completion_length (`int` or `None`, *optional*, defaults to `None`): + Maximum length of the completion. This argument is required if you want to use the default data collator + and your model is an encoder-decoder. + 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). + label_smoothing (`float`, *optional*, defaults to `0.0`): + Label smoothing factor. This argument is required if you want to use the default data collator. + 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. + - `"simpo"`: SimPO loss from the [SimPO](https://huggingface.co/papers/2405.14734) paper. + + disable_dropout (`bool`, *optional*, defaults to `True`): + Whether to disable dropout in the model. + cpo_alpha (`float`, *optional*, defaults to `1.0`): + Weight of the BC regularizer in CPO training. + simpo_gamma (`float`, *optional*, defaults to `0.5`): + Target reward margin for the SimPO loss, used only when the `loss_type="simpo"`. + label_pad_token_id (`int`, *optional*, defaults to `-100`): + Label pad token id. This argument is required if you want to use the default data collator. + padding_value (`int` or `None`, *optional*, defaults to `None`): + Padding value to use. If `None`, the padding value of the tokenizer is used. + truncation_mode (`str`,*optional*, defaults to `"keep_end"`): + Truncation mode to use when the prompt is too long. Possible values are `"keep_end"` or `"keep_start"`. + This argument is required if you want to use the default data collator. + generate_during_eval (`bool`, *optional*, defaults to `False`): + If `True`, generates and logs completions from the model to W&B or Comet during evaluation. + is_encoder_decoder (`bool` or `None`, *optional*, defaults to `None`): + When using the `model_init` argument (callable) to instantiate the model instead of the `model` argument, + you need to specify if the model returned by the callable is an encoder-decoder model. + model_init_kwargs (`dict[str, Any]` or `None`, *optional*, defaults to `None`): + Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model from a + string. + dataset_num_proc (`int` or `None`, *optional*, defaults to `None`): + Number of processes to use for processing the dataset. + + """ + 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, + hub_revision = None, + 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, + liger_kernel_config = None, + eval_use_gather_object = False, + average_tokens_across_devices = False, + max_length = 1024, + max_prompt_length = 512, + max_completion_length = None, + beta = 0.1, + label_smoothing = 0.0, + loss_type = 'sigmoid', + disable_dropout = True, + cpo_alpha = 1.0, + simpo_gamma = 0.5, + label_pad_token_id = -100, + padding_value = None, + truncation_mode = 'keep_end', + generate_during_eval = False, + is_encoder_decoder = None, + model_init_kwargs = None, + dataset_num_proc = None, + 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, + 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, + max_length = max_length, + max_prompt_length = max_prompt_length, + max_completion_length = max_completion_length, + beta = beta, + label_smoothing = label_smoothing, + loss_type = loss_type, + disable_dropout = disable_dropout, + cpo_alpha = cpo_alpha, + simpo_gamma = simpo_gamma, + label_pad_token_id = label_pad_token_id, + padding_value = padding_value, + truncation_mode = truncation_mode, + generate_during_eval = generate_during_eval, + is_encoder_decoder = is_encoder_decoder, + model_init_kwargs = model_init_kwargs, + dataset_num_proc = dataset_num_proc,**kwargs) + self.vllm_sampling_params = vllm_sampling_params + self.unsloth_num_chunks = unsloth_num_chunks +pass + +class _UnslothCPOTrainer(Trainer): + r"""""" + + _tag_names = ["trl", "cpo"] + + def __init__( + self, + model: Optional[Union[PreTrainedModel, nn.Module, str]] = None, + args: Optional[CPOConfig] = 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, + 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, + compute_metrics: Optional[Callable[[EvalLoopOutput], dict]] = None, + ): + if args.model_init_kwargs is None: + model_init_kwargs = {} + elif not isinstance(model, str): + raise ValueError("You passed model_kwargs to the CPOTrainer. 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 CPOConfig. Expected a string with either `torch.dtype` or 'auto', but got {torch_dtype}." + ) + model_init_kwargs["torch_dtype"] = torch_dtype + + if isinstance(model, str): + model = AutoModelForCausalLM.from_pretrained(model, **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 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 args.gradient_checkpointing: + # 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 args.gradient_checkpointing: + # 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." + ) + + if model is not None: + self.is_encoder_decoder = model.config.is_encoder_decoder + elif args.is_encoder_decoder is None: + raise ValueError("When no model is provided, you need to pass the parameter is_encoder_decoder.") + else: + self.is_encoder_decoder = args.is_encoder_decoder + + if self.is_encoder_decoder: + self.decoder_start_token_id = model.config.decoder_start_token_id + self.pad_token_id = model.config.pad_token_id + + if processing_class is None: + raise ValueError("processing_class must be specified to tokenize a CPO dataset.") + if args.max_length is None: + warnings.warn( + "`max_length` is not set in the CPOConfig's init" + " it will default to `512` by default, but you should do it yourself in the future.", + UserWarning, + ) + max_length = 512 + else: + max_length = args.max_length + if args.max_prompt_length is None: + warnings.warn( + "`max_prompt_length` is not set in the CPOConfig's init" + " it will default to `128` by default, but you should do it yourself in the future.", + UserWarning, + ) + max_prompt_length = 128 + else: + max_prompt_length = args.max_prompt_length + + if not max_prompt_length < max_length: + raise ValueError( + f"max_prompt_length ({max_prompt_length}) should be strictly less than max_length ({max_length})." + ) + + if args.max_completion_length is None and self.is_encoder_decoder: + warnings.warn( + "When using an encoder decoder architecture, you should set `max_completion_length` in the CPOConfig's init" + " it will default to `128` by default, but you should do it yourself in the future.", + UserWarning, + ) + max_completion_length = 128 + else: + max_completion_length = args.max_completion_length + + if data_collator is None: + data_collator = DPODataCollatorWithPadding( + pad_token_id=processing_class.pad_token_id, + label_pad_token_id=args.label_pad_token_id, + is_encoder_decoder=self.is_encoder_decoder, + ) + + if args.remove_unused_columns: + args.remove_unused_columns = False + # warn users + warnings.warn( + "When using DPODataCollatorWithPadding, you should set `remove_unused_columns=False` in your TrainingArguments" + " we have set it for you, but you should do it yourself in the future.", + UserWarning, + ) + + self.use_dpo_data_collator = True + else: + self.use_dpo_data_collator = False + + # Disable dropout in the model + if args.disable_dropout: + disable_dropout_in_model(model) + + self.max_length = max_length + self.generate_during_eval = args.generate_during_eval + self.label_pad_token_id = args.label_pad_token_id + self.padding_value = args.padding_value if args.padding_value is not None else processing_class.pad_token_id + self.max_prompt_length = max_prompt_length + self.truncation_mode = args.truncation_mode + self.max_completion_length = max_completion_length + self.processing_class = processing_class + + if args.loss_type in ["hinge", "ipo"] 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 CPOTrainer. Please use KTOTrainer.") + + self.beta = args.beta + self.label_smoothing = args.label_smoothing + self.loss_type = args.loss_type + self.cpo_alpha = args.cpo_alpha + self.aux_loss_enabled = getattr(model.config, "output_router_logits", False) + 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, + ) + + if args.loss_type == "simpo": + self.simpo_gamma = args.simpo_gamma + + self._stored_metrics = defaultdict(lambda: defaultdict(list)) + + # 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 CPO, 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 + + # Compute that only on the main process for faster data processing. + # see: https://github.com/huggingface/trl/pull/1255 + with PartialState().main_process_first(): + # Extract the prompt if needed, and apply the chat template if needed + train_dataset = train_dataset.map(maybe_extract_prompt, num_proc=args.dataset_num_proc) + train_dataset = train_dataset.map( + maybe_apply_chat_template, fn_kwargs={"tokenizer": processing_class}, num_proc=args.dataset_num_proc + ) + if eval_dataset is not None: + eval_dataset = eval_dataset.map(maybe_extract_prompt, num_proc=args.dataset_num_proc) + eval_dataset = eval_dataset.map( + maybe_apply_chat_template, + fn_kwargs={"tokenizer": processing_class}, + num_proc=args.dataset_num_proc, + ) + + # tokenize the dataset + train_dataset = train_dataset.map(self.tokenize_row, num_proc=args.dataset_num_proc) + if eval_dataset is not None: + eval_dataset = eval_dataset.map(self.tokenize_row, num_proc=args.dataset_num_proc) + + 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`." + ) + + def build_tokenized_answer(self, prompt, answer): + """ + Llama tokenizer does satisfy `enc(a + b) = enc(a) + enc(b)`. It does ensure `enc(a + b) = enc(a) + enc(a + + b)[len(enc(a)):]`. Reference: + https://github.com/EleutherAI/lm-evaluation-harness/pull/531#issuecomment-1595586257 + """ + + full_tokenized = self.processing_class(prompt + answer, add_special_tokens=False) + prompt_input_ids = self.processing_class(prompt, add_special_tokens=False)["input_ids"] + + answer_input_ids = full_tokenized["input_ids"][len(prompt_input_ids) :] + answer_attention_mask = full_tokenized["attention_mask"][len(prompt_input_ids) :] + + # Concat tokens to form `enc(a) + enc(a + b)[len(enc(a)):]` + full_concat_input_ids = np.concatenate([prompt_input_ids, answer_input_ids]) + + # Prepare input tokens for token by token comparison + full_input_ids = np.array(full_tokenized["input_ids"]) + + if len(full_input_ids) != len(full_concat_input_ids): + raise ValueError("Prompt input ids and answer input ids should have the same length.") + + # On some tokenizers, like Llama-2 tokenizer, there are occasions where tokens + # can be merged together when tokenizing prompt+answer. This could result + # on the last token from the prompt being different when tokenized on its own + # vs when done as prompt+answer. + response_token_ids_start_idx = len(prompt_input_ids) + + # If tokenized prompt is different than both prompt+answer, then it means the + # last token has changed due to merging. + if prompt_input_ids != full_tokenized["input_ids"][:response_token_ids_start_idx]: + response_token_ids_start_idx -= 1 + + prompt_input_ids = full_tokenized["input_ids"][:response_token_ids_start_idx] + prompt_attention_mask = full_tokenized["attention_mask"][:response_token_ids_start_idx] + + if len(prompt_input_ids) != len(prompt_attention_mask): + raise ValueError("Prompt input ids and attention mask should have the same length.") + + answer_input_ids = full_tokenized["input_ids"][response_token_ids_start_idx:] + answer_attention_mask = full_tokenized["attention_mask"][response_token_ids_start_idx:] + + return dict( + prompt_input_ids=prompt_input_ids, + prompt_attention_mask=prompt_attention_mask, + input_ids=answer_input_ids, + attention_mask=answer_attention_mask, + ) + + def tokenize_row(self, feature, model: Optional[Union[PreTrainedModel, nn.Module]] = None) -> dict: + """Tokenize a single row from a CPO specific dataset. + + At this stage, we don't convert to PyTorch tensors yet; we just handle the truncation in case the prompt + + chosen or prompt + rejected responses is/are too long. First we truncate the prompt; if we're still too long, + we truncate the chosen/rejected. + + We also create the labels for the chosen/rejected responses, which are of length equal to the sum of the length + of the prompt and the chosen/rejected response, with label_pad_token_id for the prompt tokens. + """ + batch = {} + prompt = feature["prompt"] + chosen = feature["chosen"] + rejected = feature["rejected"] + + if not self.is_encoder_decoder: + # Check issues below for more details + # 1. https://github.com/huggingface/trl/issues/907 + # 2. https://github.com/EleutherAI/lm-evaluation-harness/pull/531#issuecomment-1595586257 + # 3. https://github.com/LianjiaTech/BELLE/issues/337 + + if not isinstance(prompt, str): + raise ValueError(f"prompt should be an str but got {type(prompt)}") + prompt_tokens = self.processing_class(prompt, add_special_tokens=False) + prompt_tokens = {f"prompt_{k}": v for k, v in prompt_tokens.items()} + + if not isinstance(chosen, str): + raise ValueError(f"chosen should be an str but got {type(chosen)}") + chosen_tokens = self.build_tokenized_answer(prompt, chosen) + + if not isinstance(rejected, str): + raise ValueError(f"rejected should be an str but got {type(rejected)}") + rejected_tokens = self.build_tokenized_answer(prompt, rejected) + + # Last prompt token might get merged by tokenizer and + # it should not be included for generation if that happens + prompt_len_input_ids = len(prompt_tokens["prompt_input_ids"]) + + chosen_prompt_len_input_ids = len(chosen_tokens["prompt_input_ids"]) + rejected_prompt_len_input_ids = len(rejected_tokens["prompt_input_ids"]) + prompt_len_input_ids = min(chosen_prompt_len_input_ids, rejected_prompt_len_input_ids) + + for k, v in prompt_tokens.items(): + prompt_tokens[k] = v[:prompt_len_input_ids] + + # Make sure prompts only have one different token at most an + # and length only differs by 1 at most + num_diff_tokens = sum( + [a != b for a, b in zip(chosen_tokens["prompt_input_ids"], rejected_tokens["prompt_input_ids"])] + ) + num_diff_len = abs(chosen_prompt_len_input_ids - rejected_prompt_len_input_ids) + if num_diff_tokens > 1 or num_diff_len > 1: + raise ValueError( + "Chosen and rejected prompt_input_ids might only differ on the " + "last token due to tokenizer merge ops." + ) + + # add BOS token to head of prompt. Avoid adding if it's already there + prompt_tokens, chosen_tokens, rejected_tokens = add_bos_token_if_needed( + self.processing_class.bos_token_id, + prompt_len_input_ids, + prompt_tokens, + chosen_prompt_len_input_ids, + chosen_tokens, + rejected_prompt_len_input_ids, + rejected_tokens, + ) + + # add EOS token to end of answer. Avoid adding if it's already there + chosen_tokens, rejected_tokens = add_eos_token_if_needed( + self.processing_class.eos_token_id, chosen_tokens, rejected_tokens + ) + + longer_response_length = max(len(chosen_tokens["input_ids"]), len(rejected_tokens["input_ids"])) + + # if combined sequence is too long, truncate the prompt + for answer_tokens in [chosen_tokens, rejected_tokens, prompt_tokens]: + if len(answer_tokens["prompt_input_ids"]) + longer_response_length > self.max_length: + if self.truncation_mode == "keep_start": + for k in ["prompt_input_ids", "prompt_attention_mask"]: + answer_tokens[k] = answer_tokens[k][: self.max_prompt_length] + elif self.truncation_mode == "keep_end": + for k in ["prompt_input_ids", "prompt_attention_mask"]: + answer_tokens[k] = answer_tokens[k][-self.max_prompt_length :] + else: + raise ValueError(f"Unknown truncation mode: {self.truncation_mode}") + + # if that's still too long, truncate the response + for answer_tokens in [chosen_tokens, rejected_tokens]: + if len(answer_tokens["prompt_input_ids"]) + longer_response_length > self.max_length: + for k in ["input_ids", "attention_mask"]: + answer_tokens[k] = answer_tokens[k][: self.max_length - self.max_prompt_length] + + # Create labels + chosen_sequence_tokens = { + k: chosen_tokens[f"prompt_{k}"] + chosen_tokens[k] for k in ["input_ids", "attention_mask"] + } + rejected_sequence_tokens = { + k: rejected_tokens[f"prompt_{k}"] + rejected_tokens[k] for k in ["input_ids", "attention_mask"] + } + chosen_sequence_tokens["labels"] = chosen_sequence_tokens["input_ids"][:] + chosen_sequence_tokens["labels"][: len(chosen_tokens["prompt_input_ids"])] = [ + self.label_pad_token_id + ] * len(chosen_tokens["prompt_input_ids"]) + rejected_sequence_tokens["labels"] = rejected_sequence_tokens["input_ids"][:] + rejected_sequence_tokens["labels"][: len(rejected_tokens["prompt_input_ids"])] = [ + self.label_pad_token_id + ] * len(rejected_tokens["prompt_input_ids"]) + + for k, toks in { + "chosen_": chosen_sequence_tokens, + "rejected_": rejected_sequence_tokens, + "": prompt_tokens, + }.items(): + for type_key, tokens in toks.items(): + if type_key == "token_type_ids": + continue + batch[f"{k}{type_key}"] = tokens + + else: + chosen_tokens = self.processing_class( + chosen, truncation=True, max_length=self.max_completion_length, add_special_tokens=True + ) + rejected_tokens = self.processing_class( + rejected, truncation=True, max_length=self.max_completion_length, add_special_tokens=True + ) + prompt_tokens = self.processing_class( + prompt, truncation=True, max_length=self.max_prompt_length, add_special_tokens=True + ) + + batch["chosen_labels"] = chosen_tokens["input_ids"] + batch["rejected_labels"] = rejected_tokens["input_ids"] + batch["prompt_input_ids"] = prompt_tokens["input_ids"] + batch["prompt_attention_mask"] = prompt_tokens["attention_mask"] + + if model is not None and hasattr(model, "prepare_decoder_input_ids_from_labels"): + batch["rejected_decoder_input_ids"] = model.prepare_decoder_input_ids_from_labels( + labels=torch.tensor(batch["rejected_labels"]) + ) + batch["chosen_decoder_input_ids"] = model.prepare_decoder_input_ids_from_labels( + labels=torch.tensor(batch["chosen_labels"]) + ) + + return batch + + @staticmethod + def concatenated_inputs( + batch: dict[str, Union[list, torch.LongTensor]], + is_encoder_decoder: bool = False, + label_pad_token_id: int = -100, + padding_value: int = 0, + device: Optional[torch.device] = None, + ) -> dict[str, torch.LongTensor]: + """Concatenate the chosen and rejected inputs into a single tensor. + + Args: + batch: + A batch of data. Must contain the keys 'chosen_input_ids' and 'rejected_input_ids', which are tensors + of shape (batch_size, sequence_length). + is_encoder_decoder: + Whether the model is an encoder-decoder model. + label_pad_token_id: + The label pad token id. + padding_value: + The padding value to use for the concatenated inputs_ids. + device: + The device for the concatenated inputs. + + Returns: + A dictionary containing the concatenated inputs under the key 'concatenated_input_ids'. + """ + concatenated_batch = {} + + if is_encoder_decoder: + max_length = max(batch["chosen_labels"].shape[1], batch["rejected_labels"].shape[1]) + else: + max_length = max(batch["chosen_input_ids"].shape[1], batch["rejected_input_ids"].shape[1]) + + for k in batch: + if k.startswith("chosen") and isinstance(batch[k], torch.Tensor): + if "labels" in k or is_encoder_decoder: + pad_value = label_pad_token_id + elif k.endswith("_input_ids"): + pad_value = padding_value + elif k.endswith("_attention_mask"): + pad_value = 0 + concatenated_key = k.replace("chosen", "concatenated") + concatenated_batch[concatenated_key] = pad_to_length(batch[k], max_length, pad_value=pad_value) + for k in batch: + if k.startswith("rejected") and isinstance(batch[k], torch.Tensor): + if "labels" in k or is_encoder_decoder: + pad_value = label_pad_token_id + elif k.endswith("_input_ids"): + pad_value = padding_value + elif k.endswith("_attention_mask"): + pad_value = 0 + concatenated_key = k.replace("rejected", "concatenated") + concatenated_batch[concatenated_key] = torch.cat( + ( + concatenated_batch[concatenated_key], + pad_to_length(batch[k], max_length, pad_value=pad_value), + ), + dim=0, + ).to(device=device) + + if is_encoder_decoder: + concatenated_batch["concatenated_input_ids"] = batch["prompt_input_ids"].repeat(2, 1).to(device=device) + concatenated_batch["concatenated_attention_mask"] = ( + batch["prompt_attention_mask"].repeat(2, 1).to(device=device) + ) + + return concatenated_batch + + def cpo_loss( + self, + policy_chosen_logps: torch.FloatTensor, + policy_rejected_logps: torch.FloatTensor, + ) -> tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: + """Compute the CPO loss for a batch of policy and reference model log probabilities. + + Args: + policy_chosen_logps: + Log probabilities of the policy model for the chosen responses. Shape: (batch_size,) + policy_rejected_logps: + Log probabilities of the policy model for the rejected responses. Shape: (batch_size,) + + Returns: + A tuple of three tensors: (losses, chosen_rewards, rejected_rewards). The losses tensor contains the CPO + loss for each example in the batch. The chosen_rewards and rejected_rewards tensors contain the rewards for + the chosen and rejected responses, respectively. + """ + logits = (policy_chosen_logps - policy_rejected_logps).to(self.accelerator.device) + + # The beta is a temperature parameter for the CPO 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 CPO loss. + + if self.loss_type == "simpo": + gamma_logratios = self.simpo_gamma / self.beta + logits = logits - gamma_logratios + # This reduces to Equation 3 from the CPO paper when label_smoothing -> 0. + losses = ( + -F.logsigmoid(self.beta * logits) * (1 - self.label_smoothing) + - F.logsigmoid(-self.beta * logits) * self.label_smoothing + ) + elif self.loss_type == "sigmoid": + # This reduces to Equation 3 from the CPO paper when label_smoothing -> 0. + losses = ( + -F.logsigmoid(self.beta * logits) * (1 - self.label_smoothing) + - F.logsigmoid(-self.beta * logits) * 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 + else: + raise ValueError( + f"Unknown loss type: {self.loss_type}. Should be one of ['sigmoid', 'hinge', 'ipo', 'simpo']" + ) + + chosen_rewards = self.beta * (policy_chosen_logps.to(self.accelerator.device)).detach() + rejected_rewards = self.beta * (policy_rejected_logps.to(self.accelerator.device)).detach() + + return losses, chosen_rewards, rejected_rewards + + @staticmethod + def get_batch_logps( + logits: torch.FloatTensor, + labels: torch.LongTensor, + average_log_prob: bool = False, + label_pad_token_id: int = -100, + is_encoder_decoder: bool = False, + ) -> torch.FloatTensor: + """Compute the log probabilities of the given labels under the given logits. + + Args: + logits: Logits of the model (unnormalized). Shape: (batch_size, sequence_length, vocab_size) + labels: + Labels for which to compute the log probabilities. Label tokens with a value of label_pad_token_id are + ignored. Shape: (batch_size, sequence_length) + average_log_prob: + If True, return the average log probability per (non-masked) token. Otherwise, return the sum of the + log probabilities of the (non-masked) tokens. + label_pad_token_id: The label pad token id. + is_encoder_decoder: Whether the model is an encoder-decoder model. + + Returns: + A tensor of shape (batch_size,) containing the average/sum log probabilities of the given labels under the + given logits. + """ + if logits.shape[:-1] != labels.shape: + raise ValueError("Logits (batch and sequence length dim) and labels must have the same shape.") + + if not is_encoder_decoder: + labels = labels[:, 1:].clone() + logits = logits[:, :-1, :] + loss_mask = labels != label_pad_token_id + + # dummy token; we'll ignore the losses on these tokens later + labels[labels == label_pad_token_id] = 0 + + per_token_logps = selective_log_softmax(logits, labels) + + if average_log_prob: + return (per_token_logps * loss_mask).sum(-1) / loss_mask.sum(-1) + else: + return (per_token_logps * loss_mask).sum(-1) + + def concatenated_forward( + self, model: nn.Module, batch: dict[str, Union[list, torch.LongTensor]] + ) -> tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: + """Run 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. + """ + concatenated_batch = self.concatenated_inputs( + batch, + is_encoder_decoder=self.is_encoder_decoder, + label_pad_token_id=self.label_pad_token_id, + padding_value=self.padding_value, + device=self.accelerator.device, + ) + len_chosen = batch["chosen_labels"].shape[0] + + model_kwargs = ( + { + "decoder_input_ids": self._shift_right(concatenated_batch["concatenated_labels"]), + } + if self.is_encoder_decoder + else {} + ) + + if self.aux_loss_enabled: + model_kwargs["output_router_logits"] = True + + outputs = model( + concatenated_batch["concatenated_input_ids"], + attention_mask=concatenated_batch["concatenated_attention_mask"], + use_cache=False, + **model_kwargs, + ) + all_logits = outputs.logits + + def cross_entropy_loss(logits, labels): + if not self.is_encoder_decoder: + # Shift so that tokens < n predict n + logits = logits[..., :-1, :].contiguous() + labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = nn.CrossEntropyLoss() + logits = logits.view(-1, logits.shape[-1]) + labels = labels.view(-1) + # Enable model parallelism + labels = labels.to(logits.device) + loss = loss_fct(logits, labels) + return loss + + labels = concatenated_batch["concatenated_labels"].clone() + + if self.cpo_alpha == 0: + nll_loss = torch.tensor(0.0).to(self.accelerator.device) + else: + nll_loss = cross_entropy_loss(all_logits[:len_chosen], labels[:len_chosen]) + + all_logps = self.get_batch_logps( + all_logits, + concatenated_batch["concatenated_labels"], + average_log_prob=self.loss_type in ["ipo", "simpo"], + is_encoder_decoder=self.is_encoder_decoder, + label_pad_token_id=self.label_pad_token_id, + ) + + chosen_logps = all_logps[:len_chosen] + rejected_logps = all_logps[len_chosen:] + + chosen_logits = all_logits[:len_chosen] + rejected_logits = all_logits[len_chosen:] + + if self.aux_loss_enabled: + return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, nll_loss, outputs.aux_loss) + + return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, nll_loss) + + def get_batch_loss_metrics( + self, + model, + batch: dict[str, Union[list, torch.LongTensor]], + train_eval: Literal["train", "eval"] = "train", + ): + """Compute the CPO loss and other metrics for the given batch of inputs for train or test.""" + metrics = {} + + forward_output = self.concatenated_forward(model, batch) + ( + policy_chosen_logps, + policy_rejected_logps, + policy_chosen_logits, + policy_rejected_logits, + policy_nll_loss, + ) = forward_output[:5] + if self.aux_loss_enabled: + aux_loss = forward_output[5] + + losses, chosen_rewards, rejected_rewards = self.cpo_loss( + policy_chosen_logps, + policy_rejected_logps, + ) + + loss = losses.mean() + self.cpo_alpha * policy_nll_loss + reward_accuracies = (chosen_rewards > rejected_rewards).float() + + 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/rejected"] = ( + self.accelerator.gather_for_metrics(policy_rejected_logps).detach().mean().item() + ) + metrics[f"{prefix}logps/chosen"] = ( + self.accelerator.gather_for_metrics(policy_chosen_logps).detach().mean().item() + ) + metrics[f"{prefix}logits/rejected"] = ( + self.accelerator.gather_for_metrics(policy_rejected_logits.detach().mean()).mean().item() + ) + metrics[f"{prefix}logits/chosen"] = ( + self.accelerator.gather_for_metrics(policy_chosen_logits.detach().mean()).mean().item() + ) + metrics[f"{prefix}nll_loss"] = self.accelerator.gather_for_metrics(policy_nll_loss).detach().mean().item() + + if self.aux_loss_enabled: + loss += self.aux_loss_coef * aux_loss + + return loss, 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]]]: + compute_loss_context_manager = ( + autocast(self.accelerator.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") + + # force log the metrics + self.store_metrics(metrics, train_eval="train") + + if return_outputs: + return (loss, metrics) + return loss + + def generate_from_model(self, model, batch: dict[str, torch.LongTensor]) -> 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. + generate_context_manager = ( + autocast(self.accelerator.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.processing_class.pad_token_id, + ) + + policy_output = pad_to_length(policy_output, self.max_length, self.processing_class.pad_token_id) + policy_output_decoded = self.processing_class.batch_decode(policy_output, skip_special_tokens=True) + + return policy_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 = [] + + prediction_context_manager = ( + autocast(self.accelerator.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 = self.generate_from_model(self.model, random_batch) + + table = pd.DataFrame( + columns=["Prompt", "Policy"], + data=[ + [prompt, pol[len(prompt) :]] for prompt, pol in zip(random_batch["prompt"], policy_output_decoded) + ], + ) + if "wandb" in self.args.report_to: + 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] + return super().log(logs, start_time) + + def _shift_right(self, input_ids): + if self.decoder_start_token_id is None: + raise ValueError( + "model.config.decoder_start_token_id has to be defined. It is usually set to the pad_token_id." + ) + + # shift inputs to the right + if is_torch_fx_proxy(input_ids): + # Item assignment is not supported natively for proxies. + shifted_input_ids = torch.full(input_ids.shape[:-1] + (1,), self.decoder_start_token_id) + shifted_input_ids = torch.cat([shifted_input_ids, input_ids[..., :-1]], dim=-1) + else: + shifted_input_ids = input_ids.new_zeros(input_ids.shape) + shifted_input_ids[..., 1:] = input_ids[..., :-1].clone() + shifted_input_ids[..., 0] = self.decoder_start_token_id + + if self.pad_token_id is None: + raise ValueError("model.config.pad_token_id has to be defined.") + # replace possible -100 values in labels by `pad_token_id` + shifted_input_ids.masked_fill_(shifted_input_ids == -100, self.pad_token_id) + + return shifted_input_ids + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) + + 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 + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + citation = textwrap.dedent("""\ + @inproceedings{xu2024contrastive, + title = {{Contrastive Preference Optimization: Pushing the Boundaries of LLM Performance in Machine Translation}}, + author = {Haoran Xu and Amr Sharaf and Yunmo Chen and Weiting Tan and Lingfeng Shen and Benjamin Van Durme and Kenton Murray and Young Jin Kim}, + 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=51iwkioZpn} + }""") + + 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="CPO", + trainer_citation=citation, + paper_title="Contrastive Preference Optimization: Pushing the Boundaries of LLM Performance in Machine Translation", + paper_id="2401.08417", + ) + model_card.save(os.path.join(self.args.output_dir, "README.md")) +class UnslothCPOTrainer(_UnslothCPOTrainer): + """ + + Initialize CPOTrainer. + + Args: + model (`transformers.PreTrainedModel`): + The model to train, preferably an `AutoModelForSequenceClassification`. + args (`CPOConfig`): + The CPO 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 (`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. + 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. + 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. + + """ + def __init__( + self, + model = None, + args = None, + data_collator = None, + train_dataset = None, + eval_dataset = None, + processing_class = None, + model_init = None, + callbacks = None, + preprocess_logits_for_metrics = None, + peft_config = None, + compute_metrics = None, + **kwargs + ): + if args is None: args = UnslothCPOConfig() + 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 + 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) + 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 + 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('cpo_trainer', other_metrics) + + 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, + callbacks = callbacks, + preprocess_logits_for_metrics = preprocess_logits_for_metrics, + peft_config = peft_config, + compute_metrics = compute_metrics,**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 diff --git a/unsloth_compiled_cache/UnslothDDPOTrainer.py b/unsloth_compiled_cache/UnslothDDPOTrainer.py new file mode 100644 index 0000000..9f8e9db --- /dev/null +++ b/unsloth_compiled_cache/UnslothDDPOTrainer.py @@ -0,0 +1,889 @@ +""" +2025.7.1 +2025.7.1 +4.53.1 +0.19.1 +__UNSLOTH_VERSIONING__ +""" +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from trl.trainer.ddpo_trainer import (Accelerator, Any, Callable, DDPOConfig, DDPOStableDiffusionPipeline, DDPOTrainer, Optional, Path, PerPromptStatTracker, ProjectConfiguration, PyTorchModelHubMixin, Union, defaultdict, futures, generate_model_card, get_comet_experiment_url, is_wandb_available, logger, os, set_seed, textwrap, torch, warn) + + +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 UnslothDDPOConfig(DDPOConfig): + """ + + Configuration class for the [`DDPOTrainer`]. + + 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: + exp_name (`str`, *optional*, defaults to `os.path.basename(sys.argv[0])[: -len(".py")]`): + Name of this experiment (by default is the file name without the extension name). + run_name (`str`, *optional*, defaults to `""`): + Name of this run. + seed (`int`, *optional*, defaults to `0`): + Random seed. + log_with (`Literal["wandb", "tensorboard"]]` or `None`, *optional*, defaults to `None`): + Log with either 'wandb' or 'tensorboard', check + https://huggingface.co/docs/accelerate/usage_guides/tracking for more details. + tracker_kwargs (`Dict`, *optional*, defaults to `{}`): + Keyword arguments for the tracker (e.g. wandb_project). + accelerator_kwargs (`Dict`, *optional*, defaults to `{}`): + Keyword arguments for the accelerator. + project_kwargs (`Dict`, *optional*, defaults to `{}`): + Keyword arguments for the accelerator project config (e.g. `logging_dir`). + tracker_project_name (`str`, *optional*, defaults to `"trl"`): + Name of project to use for tracking. + logdir (`str`, *optional*, defaults to `"logs"`): + Top-level logging directory for checkpoint saving. + num_epochs (`int`, *optional*, defaults to `100`): + Number of epochs to train. + save_freq (`int`, *optional*, defaults to `1`): + Number of epochs between saving model checkpoints. + num_checkpoint_limit (`int`, *optional*, defaults to `5`): + Number of checkpoints to keep before overwriting old ones. + mixed_precision (`str`, *optional*, defaults to `"fp16"`): + Mixed precision training. + allow_tf32 (`bool`, *optional*, defaults to `True`): + Allow `tf32` on Ampere GPUs. + resume_from (`str`, *optional*, defaults to `""`): + Resume training from a checkpoint. + sample_num_steps (`int`, *optional*, defaults to `50`): + Number of sampler inference steps. + sample_eta (`float`, *optional*, defaults to `1.0`): + Eta parameter for the DDIM sampler. + sample_guidance_scale (`float`, *optional*, defaults to `5.0`): + Classifier-free guidance weight. + sample_batch_size (`int`, *optional*, defaults to `1`): + Batch size (per GPU) to use for sampling. + sample_num_batches_per_epoch (`int`, *optional*, defaults to `2`): + Number of batches to sample per epoch. + train_batch_size (`int`, *optional*, defaults to `1`): + Batch size (per GPU) to use for training. + train_use_8bit_adam (`bool`, *optional*, defaults to `False`): + Use 8bit Adam optimizer from bitsandbytes. + train_learning_rate (`float`, *optional*, defaults to `3e-4`): + Learning rate. + train_adam_beta1 (`float`, *optional*, defaults to `0.9`): + Adam beta1. + train_adam_beta2 (`float`, *optional*, defaults to `0.999`): + Adam beta2. + train_adam_weight_decay (`float`, *optional*, defaults to `1e-4`): + Adam weight decay. + train_adam_epsilon (`float`, *optional*, defaults to `1e-8`): + Adam epsilon. + train_gradient_accumulation_steps (`int`, *optional*, defaults to `1`): + Number of gradient accumulation steps. + train_max_grad_norm (`float`, *optional*, defaults to `1.0`): + Maximum gradient norm for gradient clipping. + train_num_inner_epochs (`int`, *optional*, defaults to `1`): + Number of inner epochs per outer epoch. + train_cfg (`bool`, *optional*, defaults to `True`): + Whether to use classifier-free guidance during training. + train_adv_clip_max (`float`, *optional*, defaults to `5.0`): + Clip advantages to the range. + train_clip_range (`float`, *optional*, defaults to `1e-4`): + PPO clip range. + train_timestep_fraction (`float`, *optional*, defaults to `1.0`): + Fraction of timesteps to train on. + per_prompt_stat_tracking (`bool`, *optional*, defaults to `False`): + Whether to track statistics for each prompt separately. + per_prompt_stat_tracking_buffer_size (`int`, *optional*, defaults to `16`): + Number of reward values to store in the buffer for each prompt. + per_prompt_stat_tracking_min_count (`int`, *optional*, defaults to `16`): + Minimum number of reward values to store in the buffer. + async_reward_computation (`bool`, *optional*, defaults to `False`): + Whether to compute rewards asynchronously. + max_workers (`int`, *optional*, defaults to `2`): + Maximum number of workers to use for async reward computation. + negative_prompts (`str`, *optional*, defaults to `""`): + Comma-separated list of prompts to use as negative examples. + push_to_hub (`bool`, *optional*, defaults to `False`): + Whether to push the final model checkpoint to the Hub. + + """ + 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, + exp_name = '001测试unsloth', + run_name = '', + seed = 3407, + log_with = None, + tracker_project_name = 'trl', + logdir = 'logs', + num_epochs = 100, + save_freq = 1, + num_checkpoint_limit = 5, + mixed_precision = 'fp16', + allow_tf32 = True, + resume_from = '', + sample_num_steps = 50, + sample_eta = 1.0, + sample_guidance_scale = 5.0, + sample_batch_size = 1, + sample_num_batches_per_epoch = 2, + train_batch_size = 1, + train_use_8bit_adam = False, + train_learning_rate = 5e-05, + train_adam_beta1 = 0.9, + train_adam_beta2 = 0.999, + train_adam_weight_decay = 0.01, + train_adam_epsilon = 1e-08, + train_gradient_accumulation_steps = 2, + train_max_grad_norm = 1.0, + train_num_inner_epochs = 1, + train_cfg = True, + train_adv_clip_max = 5.0, + train_clip_range = 0.0001, + train_timestep_fraction = 1.0, + per_prompt_stat_tracking = False, + per_prompt_stat_tracking_buffer_size = 16, + per_prompt_stat_tracking_min_count = 16, + async_reward_computation = False, + max_workers = 2, + negative_prompts = '', + push_to_hub = False, + vllm_sampling_params = None, + unsloth_num_chunks = -1, + **kwargs, + ): + + super().__init__( + exp_name = exp_name, + run_name = run_name, + seed = seed, + log_with = log_with, + tracker_project_name = tracker_project_name, + logdir = logdir, + num_epochs = num_epochs, + save_freq = save_freq, + num_checkpoint_limit = num_checkpoint_limit, + mixed_precision = mixed_precision, + allow_tf32 = allow_tf32, + resume_from = resume_from, + sample_num_steps = sample_num_steps, + sample_eta = sample_eta, + sample_guidance_scale = sample_guidance_scale, + sample_batch_size = sample_batch_size, + sample_num_batches_per_epoch = sample_num_batches_per_epoch, + train_batch_size = train_batch_size, + train_use_8bit_adam = train_use_8bit_adam, + train_learning_rate = train_learning_rate, + train_adam_beta1 = train_adam_beta1, + train_adam_beta2 = train_adam_beta2, + train_adam_weight_decay = train_adam_weight_decay, + train_adam_epsilon = train_adam_epsilon, + train_gradient_accumulation_steps = train_gradient_accumulation_steps, + train_max_grad_norm = train_max_grad_norm, + train_num_inner_epochs = train_num_inner_epochs, + train_cfg = train_cfg, + train_adv_clip_max = train_adv_clip_max, + train_clip_range = train_clip_range, + train_timestep_fraction = train_timestep_fraction, + per_prompt_stat_tracking = per_prompt_stat_tracking, + per_prompt_stat_tracking_buffer_size = per_prompt_stat_tracking_buffer_size, + per_prompt_stat_tracking_min_count = per_prompt_stat_tracking_min_count, + async_reward_computation = async_reward_computation, + max_workers = max_workers, + negative_prompts = negative_prompts, + push_to_hub = push_to_hub,**kwargs) + self.vllm_sampling_params = vllm_sampling_params + self.unsloth_num_chunks = unsloth_num_chunks +pass + +class _UnslothDDPOTrainer(PyTorchModelHubMixin): + """""" + + _tag_names = ["trl", "ddpo"] + + def __init__( + self, + config: DDPOConfig, + reward_function: Callable[[torch.Tensor, tuple[str], tuple[Any]], torch.Tensor], + prompt_function: Callable[[], tuple[str, Any]], + sd_pipeline: DDPOStableDiffusionPipeline, + image_samples_hook: Optional[Callable[[Any, Any, Any], Any]] = None, + ): + if image_samples_hook is None: + warn("No image_samples_hook provided; no images will be logged") + + self.prompt_fn = prompt_function + self.reward_fn = reward_function + self.config = config + self.image_samples_callback = image_samples_hook + + accelerator_project_config = ProjectConfiguration(**self.config.project_kwargs) + + if self.config.resume_from: + self.config.resume_from = os.path.normpath(os.path.expanduser(self.config.resume_from)) + if "checkpoint_" not in os.path.basename(self.config.resume_from): + # get the most recent checkpoint in this directory + checkpoints = list( + filter( + lambda x: "checkpoint_" in x, + os.listdir(self.config.resume_from), + ) + ) + if len(checkpoints) == 0: + raise ValueError(f"No checkpoints found in {self.config.resume_from}") + checkpoint_numbers = sorted([int(x.split("_")[-1]) for x in checkpoints]) + self.config.resume_from = os.path.join( + self.config.resume_from, + f"checkpoint_{checkpoint_numbers[-1]}", + ) + + accelerator_project_config.iteration = checkpoint_numbers[-1] + 1 + + # number of timesteps within each trajectory to train on + self.num_train_timesteps = int(self.config.sample_num_steps * self.config.train_timestep_fraction) + + self.accelerator = Accelerator( + log_with=self.config.log_with, + mixed_precision=self.config.mixed_precision, + project_config=accelerator_project_config, + # we always accumulate gradients across timesteps; we want config.train.gradient_accumulation_steps to be the + # number of *samples* we accumulate across, so we need to multiply by the number of training timesteps to get + # the total number of optimizer steps to accumulate across. + gradient_accumulation_steps=self.config.train_gradient_accumulation_steps * self.num_train_timesteps, + **self.config.accelerator_kwargs, + ) + + is_okay, message = self._config_check() + if not is_okay: + raise ValueError(message) + + is_using_tensorboard = config.log_with is not None and config.log_with == "tensorboard" + + if self.accelerator.is_main_process: + self.accelerator.init_trackers( + self.config.tracker_project_name, + config=dict(ddpo_trainer_config=config.to_dict()) if not is_using_tensorboard else config.to_dict(), + init_kwargs=self.config.tracker_kwargs, + ) + + logger.info(f"\n{config}") + + set_seed(self.config.seed, device_specific=True) + + self.sd_pipeline = sd_pipeline + + self.sd_pipeline.set_progress_bar_config( + position=1, + disable=not self.accelerator.is_local_main_process, + leave=False, + desc="Timestep", + dynamic_ncols=True, + ) + + # For mixed precision training we cast all non-trainable weights [vae, non-lora text_encoder and non-lora unet] to half-precision + # as these weights are only used for inference, keeping weights in full precision is not required. + if self.accelerator.mixed_precision == "fp16": + inference_dtype = torch.float16 + elif self.accelerator.mixed_precision == "bf16": + inference_dtype = torch.bfloat16 + else: + inference_dtype = torch.float32 + + self.sd_pipeline.vae.to(self.accelerator.device, dtype=inference_dtype) + self.sd_pipeline.text_encoder.to(self.accelerator.device, dtype=inference_dtype) + self.sd_pipeline.unet.to(self.accelerator.device, dtype=inference_dtype) + + trainable_layers = self.sd_pipeline.get_trainable_layers() + + self.accelerator.register_save_state_pre_hook(self._save_model_hook) + self.accelerator.register_load_state_pre_hook(self._load_model_hook) + + # Enable TF32 for faster training on Ampere GPUs, + # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices + if self.config.allow_tf32: + torch.backends.cuda.matmul.allow_tf32 = True + + self.optimizer = self._setup_optimizer( + trainable_layers.parameters() if not isinstance(trainable_layers, list) else trainable_layers + ) + + self.neg_prompt_embed = self.sd_pipeline.text_encoder( + self.sd_pipeline.tokenizer( + [""] if self.config.negative_prompts is None else self.config.negative_prompts, + return_tensors="pt", + padding="max_length", + truncation=True, + max_length=self.sd_pipeline.tokenizer.model_max_length, + ).input_ids.to(self.accelerator.device) + )[0] + + if config.per_prompt_stat_tracking: + self.stat_tracker = PerPromptStatTracker( + config.per_prompt_stat_tracking_buffer_size, + config.per_prompt_stat_tracking_min_count, + ) + + # NOTE: for some reason, autocast is necessary for non-lora training but for lora training it isn't necessary and it uses + # more memory + self.autocast = self.sd_pipeline.autocast or self.accelerator.autocast + + if hasattr(self.sd_pipeline, "use_lora") and self.sd_pipeline.use_lora: + unet, self.optimizer = self.accelerator.prepare(trainable_layers, self.optimizer) + self.trainable_layers = list(filter(lambda p: p.requires_grad, unet.parameters())) + else: + self.trainable_layers, self.optimizer = self.accelerator.prepare(trainable_layers, self.optimizer) + + if self.config.async_reward_computation: + self.executor = futures.ThreadPoolExecutor(max_workers=config.max_workers) + + if config.resume_from: + logger.info(f"Resuming from {config.resume_from}") + self.accelerator.load_state(config.resume_from) + self.first_epoch = int(config.resume_from.split("_")[-1]) + 1 + else: + self.first_epoch = 0 + + def compute_rewards(self, prompt_image_pairs, is_async=False): + if not is_async: + rewards = [] + for images, prompts, prompt_metadata in prompt_image_pairs: + reward, reward_metadata = self.reward_fn(images, prompts, prompt_metadata) + rewards.append( + ( + torch.as_tensor(reward, device=self.accelerator.device), + reward_metadata, + ) + ) + else: + rewards = self.executor.map(lambda x: self.reward_fn(*x), prompt_image_pairs) + rewards = [ + (torch.as_tensor(reward.result(), device=self.accelerator.device), reward_metadata.result()) + for reward, reward_metadata in rewards + ] + + return zip(*rewards) + + def step(self, epoch: int, global_step: int): + """ + Perform a single step of training. + + Args: + epoch (int): The current epoch. + global_step (int): The current global step. + + Side Effects: + - Model weights are updated + - Logs the statistics to the accelerator trackers. + - If `self.image_samples_callback` is not None, it will be called with the prompt_image_pairs, global_step, + and the accelerator tracker. + + Returns: + global_step (int): The updated global step. + + """ + samples, prompt_image_data = self._generate_samples( + iterations=self.config.sample_num_batches_per_epoch, + batch_size=self.config.sample_batch_size, + ) + + # collate samples into dict where each entry has shape (num_batches_per_epoch * sample.batch_size, ...) + samples = {k: torch.cat([s[k] for s in samples]) for k in samples[0].keys()} + rewards, rewards_metadata = self.compute_rewards( + prompt_image_data, is_async=self.config.async_reward_computation + ) + + for i, image_data in enumerate(prompt_image_data): + image_data.extend([rewards[i], rewards_metadata[i]]) + + if self.image_samples_callback is not None: + self.image_samples_callback(prompt_image_data, global_step, self.accelerator.trackers[0]) + + rewards = torch.cat(rewards) + rewards = self.accelerator.gather(rewards).cpu().numpy() + + self.accelerator.log( + { + "reward": rewards, + "epoch": epoch, + "reward_mean": rewards.mean(), + "reward_std": rewards.std(), + }, + step=global_step, + ) + + if self.config.per_prompt_stat_tracking: + # gather the prompts across processes + prompt_ids = self.accelerator.gather(samples["prompt_ids"]).cpu().numpy() + prompts = self.sd_pipeline.tokenizer.batch_decode(prompt_ids, skip_special_tokens=True) + advantages = self.stat_tracker.update(prompts, rewards) + else: + advantages = (rewards - rewards.mean()) / (rewards.std() + 1e-8) + + # ungather advantages; keep the entries corresponding to the samples on this process + samples["advantages"] = ( + torch.as_tensor(advantages) + .reshape(self.accelerator.num_processes, -1)[self.accelerator.process_index] + .to(self.accelerator.device) + ) + + del samples["prompt_ids"] + + total_batch_size, num_timesteps = samples["timesteps"].shape + + for inner_epoch in range(self.config.train_num_inner_epochs): + # shuffle samples along batch dimension + perm = torch.randperm(total_batch_size, device=self.accelerator.device) + samples = {k: v[perm] for k, v in samples.items()} + + # shuffle along time dimension independently for each sample + # still trying to understand the code below + perms = torch.stack( + [torch.randperm(num_timesteps, device=self.accelerator.device) for _ in range(total_batch_size)] + ) + + for key in ["timesteps", "latents", "next_latents", "log_probs"]: + samples[key] = samples[key][ + torch.arange(total_batch_size, device=self.accelerator.device)[:, None], + perms, + ] + + original_keys = samples.keys() + original_values = samples.values() + # rebatch them as user defined train_batch_size is different from sample_batch_size + reshaped_values = [v.reshape(-1, self.config.train_batch_size, *v.shape[1:]) for v in original_values] + + # Transpose the list of original values + transposed_values = zip(*reshaped_values) + # Create new dictionaries for each row of transposed values + samples_batched = [dict(zip(original_keys, row_values)) for row_values in transposed_values] + + self.sd_pipeline.unet.train() + global_step = self._train_batched_samples(inner_epoch, epoch, global_step, samples_batched) + # ensure optimization step at the end of the inner epoch + if not self.accelerator.sync_gradients: + raise ValueError( + "Optimization step should have been performed by this point. Please check calculated gradient accumulation settings." + ) + + if epoch != 0 and epoch % self.config.save_freq == 0 and self.accelerator.is_main_process: + self.accelerator.save_state() + + return global_step + + def calculate_loss(self, latents, timesteps, next_latents, log_probs, advantages, embeds): + """ + Calculate the loss for a batch of an unpacked sample + + Args: + latents (torch.Tensor): + The latents sampled from the diffusion model, shape: [batch_size, num_channels_latents, height, width] + timesteps (torch.Tensor): + The timesteps sampled from the diffusion model, shape: [batch_size] + next_latents (torch.Tensor): + The next latents sampled from the diffusion model, shape: [batch_size, num_channels_latents, height, + width] + log_probs (torch.Tensor): + The log probabilities of the latents, shape: [batch_size] + advantages (torch.Tensor): + The advantages of the latents, shape: [batch_size] + embeds (torch.Tensor): + The embeddings of the prompts, shape: [2*batch_size or batch_size, ...] Note: the "or" is because if + train_cfg is True, the expectation is that negative prompts are concatenated to the embeds + + Returns: + loss (torch.Tensor), approx_kl (torch.Tensor), clipfrac (torch.Tensor) (all of these are of shape (1,)) + """ + with self.autocast(): + if self.config.train_cfg: + noise_pred = self.sd_pipeline.unet( + torch.cat([latents] * 2), + torch.cat([timesteps] * 2), + embeds, + ).sample + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + self.config.sample_guidance_scale * ( + noise_pred_text - noise_pred_uncond + ) + else: + noise_pred = self.sd_pipeline.unet( + latents, + timesteps, + embeds, + ).sample + # compute the log prob of next_latents given latents under the current model + + scheduler_step_output = self.sd_pipeline.scheduler_step( + noise_pred, + timesteps, + latents, + eta=self.config.sample_eta, + prev_sample=next_latents, + ) + + log_prob = scheduler_step_output.log_probs + + advantages = torch.clamp( + advantages, + -self.config.train_adv_clip_max, + self.config.train_adv_clip_max, + ) + + ratio = torch.exp(log_prob - log_probs) + + loss = self.loss(advantages, self.config.train_clip_range, ratio) + + approx_kl = 0.5 * torch.mean((log_prob - log_probs) ** 2) + + clipfrac = torch.mean((torch.abs(ratio - 1.0) > self.config.train_clip_range).float()) + + return loss, approx_kl, clipfrac + + def loss( + self, + advantages: torch.Tensor, + clip_range: float, + ratio: torch.Tensor, + ): + unclipped_loss = -advantages * ratio + clipped_loss = -advantages * torch.clamp( + ratio, + 1.0 - clip_range, + 1.0 + clip_range, + ) + return torch.mean(torch.maximum(unclipped_loss, clipped_loss)) + + def _setup_optimizer(self, trainable_layers_parameters): + if self.config.train_use_8bit_adam: + import bitsandbytes + + optimizer_cls = bitsandbytes.optim.AdamW8bit + else: + optimizer_cls = torch.optim.AdamW + + return optimizer_cls( + trainable_layers_parameters, + lr=self.config.train_learning_rate, + betas=(self.config.train_adam_beta1, self.config.train_adam_beta2), + weight_decay=self.config.train_adam_weight_decay, + eps=self.config.train_adam_epsilon, + ) + + def _save_model_hook(self, models, weights, output_dir): + self.sd_pipeline.save_checkpoint(models, weights, output_dir) + weights.pop() # ensures that accelerate doesn't try to handle saving of the model + + def _load_model_hook(self, models, input_dir): + self.sd_pipeline.load_checkpoint(models, input_dir) + models.pop() # ensures that accelerate doesn't try to handle loading of the model + + def _generate_samples(self, iterations, batch_size): + """ + Generate samples from the model + + Args: + iterations (int): Number of iterations to generate samples for + batch_size (int): Batch size to use for sampling + + Returns: + samples (list[dict[str, torch.Tensor]]), prompt_image_pairs (list[list[Any]]) + """ + samples = [] + prompt_image_pairs = [] + self.sd_pipeline.unet.eval() + + sample_neg_prompt_embeds = self.neg_prompt_embed.repeat(batch_size, 1, 1) + + for _ in range(iterations): + prompts, prompt_metadata = zip(*[self.prompt_fn() for _ in range(batch_size)]) + + prompt_ids = self.sd_pipeline.tokenizer( + prompts, + return_tensors="pt", + padding="max_length", + truncation=True, + max_length=self.sd_pipeline.tokenizer.model_max_length, + ).input_ids.to(self.accelerator.device) + prompt_embeds = self.sd_pipeline.text_encoder(prompt_ids)[0] + + with self.autocast(): + sd_output = self.sd_pipeline( + prompt_embeds=prompt_embeds, + negative_prompt_embeds=sample_neg_prompt_embeds, + num_inference_steps=self.config.sample_num_steps, + guidance_scale=self.config.sample_guidance_scale, + eta=self.config.sample_eta, + output_type="pt", + ) + + images = sd_output.images + latents = sd_output.latents + log_probs = sd_output.log_probs + + latents = torch.stack(latents, dim=1) # (batch_size, num_steps + 1, ...) + log_probs = torch.stack(log_probs, dim=1) # (batch_size, num_steps, 1) + timesteps = self.sd_pipeline.scheduler.timesteps.repeat(batch_size, 1) # (batch_size, num_steps) + + samples.append( + { + "prompt_ids": prompt_ids, + "prompt_embeds": prompt_embeds, + "timesteps": timesteps, + "latents": latents[:, :-1], # each entry is the latent before timestep t + "next_latents": latents[:, 1:], # each entry is the latent after timestep t + "log_probs": log_probs, + "negative_prompt_embeds": sample_neg_prompt_embeds, + } + ) + prompt_image_pairs.append([images, prompts, prompt_metadata]) + + return samples, prompt_image_pairs + + def _train_batched_samples(self, inner_epoch, epoch, global_step, batched_samples): + """ + Train on a batch of samples. Main training segment + + Args: + inner_epoch (int): The current inner epoch + epoch (int): The current epoch + global_step (int): The current global step + batched_samples (list[dict[str, torch.Tensor]]): The batched samples to train on + + Side Effects: + - Model weights are updated + - Logs the statistics to the accelerator trackers. + + Returns: + global_step (int): The updated global step + """ + info = defaultdict(list) + for _i, sample in enumerate(batched_samples): + if self.config.train_cfg: + # concat negative prompts to sample prompts to avoid two forward passes + embeds = torch.cat([sample["negative_prompt_embeds"], sample["prompt_embeds"]]) + else: + embeds = sample["prompt_embeds"] + + for j in range(self.num_train_timesteps): + with self.accelerator.accumulate(self.sd_pipeline.unet): + loss, approx_kl, clipfrac = self.calculate_loss( + sample["latents"][:, j], + sample["timesteps"][:, j], + sample["next_latents"][:, j], + sample["log_probs"][:, j], + sample["advantages"], + embeds, + ) + info["approx_kl"].append(approx_kl) + info["clipfrac"].append(clipfrac) + info["loss"].append(loss) + + self.accelerator.backward(loss) + if self.accelerator.sync_gradients: + self.accelerator.clip_grad_norm_( + self.trainable_layers.parameters() + if not isinstance(self.trainable_layers, list) + else self.trainable_layers, + self.config.train_max_grad_norm, + ) + self.optimizer.step() + self.optimizer.zero_grad() + + # Checks if the accelerator has performed an optimization step behind the scenes + if self.accelerator.sync_gradients: + # log training-related stuff + info = {k: torch.mean(torch.stack(v)) for k, v in info.items()} + info = self.accelerator.reduce(info, reduction="mean") + info.update({"epoch": epoch, "inner_epoch": inner_epoch}) + self.accelerator.log(info, step=global_step) + global_step += 1 + info = defaultdict(list) + return global_step + + def _config_check(self) -> tuple[bool, str]: + samples_per_epoch = ( + self.config.sample_batch_size * self.accelerator.num_processes * self.config.sample_num_batches_per_epoch + ) + total_train_batch_size = ( + self.config.train_batch_size + * self.accelerator.num_processes + * self.config.train_gradient_accumulation_steps + ) + + if not self.config.sample_batch_size >= self.config.train_batch_size: + return ( + False, + f"Sample batch size ({self.config.sample_batch_size}) must be greater than or equal to the train batch size ({self.config.train_batch_size})", + ) + if not self.config.sample_batch_size % self.config.train_batch_size == 0: + return ( + False, + f"Sample batch size ({self.config.sample_batch_size}) must be divisible by the train batch size ({self.config.train_batch_size})", + ) + if not samples_per_epoch % total_train_batch_size == 0: + return ( + False, + f"Number of samples per epoch ({samples_per_epoch}) must be divisible by the total train batch size ({total_train_batch_size})", + ) + return True, "" + + def train(self, epochs: Optional[int] = None): + """ + Train the model for a given number of epochs + """ + global_step = 0 + if epochs is None: + epochs = self.config.num_epochs + for epoch in range(self.first_epoch, epochs): + global_step = self.step(epoch, global_step) + + def _save_pretrained(self, save_directory): + self.sd_pipeline.save_pretrained(save_directory) + self.create_model_card() + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) + + 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 + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + citation = textwrap.dedent("""\ + @inproceedings{black2024training, + title = {{Training Diffusion Models with Reinforcement Learning}}, + author = {Kevin Black and Michael Janner and Yilun Du and Ilya Kostrikov and Sergey Levine}, + year = 2024, + booktitle = {The Twelfth International Conference on Learning Representations, {ICLR} 2024, Vienna, Austria, May 7-11, 2024}, + publisher = {OpenReview.net}, + url = {https://openreview.net/forum?id=YCWjhGrJFD}, + }""") + + 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="DDPO", + trainer_citation=citation, + paper_title="Training Diffusion Models with Reinforcement Learning", + paper_id="2305.13301", + ) + + model_card.save(os.path.join(self.args.output_dir, "README.md")) +class UnslothDDPOTrainer(_UnslothDDPOTrainer): + """ + + The DDPOTrainer uses Deep Diffusion Policy Optimization to optimise diffusion models. Note, this trainer is heavily + inspired by the work here: https://github.com/kvablack/ddpo-pytorch As of now only Stable Diffusion based pipelines + are supported + + Attributes: + **config** (`DDPOConfig`) -- Configuration object for DDPOTrainer. Check the documentation of `PPOConfig` for more: + details. + **reward_function** (Callable[[torch.Tensor, tuple[str], tuple[Any]], torch.Tensor]) -- Reward function to be used: + **prompt_function** (Callable[[], tuple[str, Any]]) -- Function to generate prompts to guide model + **sd_pipeline** (`DDPOStableDiffusionPipeline`) -- Stable Diffusion pipeline to be used for training. + **image_samples_hook** (Optional[Callable[[Any, Any, Any], Any]]) -- Hook to be called to log images + + """ + def __init__( + self, + config, + reward_function, + prompt_function, + sd_pipeline, + image_samples_hook = None, + **kwargs + ): + if args is None: args = UnslothDDPOConfig() + other_metrics = [] + + from unsloth_zoo.logging_utils import PatchRLStatistics + PatchRLStatistics('ddpo_trainer', other_metrics) + + super().__init__( + config = config, + reward_function = reward_function, + prompt_function = prompt_function, + sd_pipeline = sd_pipeline, + image_samples_hook = image_samples_hook,**kwargs) + +pass diff --git a/unsloth_compiled_cache/UnslothDPOTrainer.py b/unsloth_compiled_cache/UnslothDPOTrainer.py new file mode 100644 index 0000000..ddc35b5 --- /dev/null +++ b/unsloth_compiled_cache/UnslothDPOTrainer.py @@ -0,0 +1,2443 @@ +""" +2025.7.1 +2025.7.1 +4.53.1 +0.19.1 +__UNSLOTH_VERSIONING__ +""" +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from trl.trainer.dpo_trainer import (Any, AutoModelForCausalLM, AutoTokenizer, BaseImageProcessor, Callable, DPOConfig, DPOTrainer, DataCollator, DataCollatorForPreference, DataLoader, Dataset, EvalLoopOutput, F, FDivergenceConstants, FDivergenceType, FeatureExtractionMixin, IterableDataset, Literal, MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES, Optional, PartialState, Path, PeftConfig, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, RunningMoments, SyncRefModelCallback, Trainer, TrainerCallback, Union, autocast, 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, get_peft_model, inspect, is_comet_available, is_liger_kernel_available, is_peft_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, shift_tokens_right, textwrap, torch, tqdm, warnings, F, Optional, PeftModel, PreTrainedModel, Trainer, is_peft_available, os, torch) + + +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`]. + + This class includes only the parameters that are specific to DPO training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may + differ from those in [`~transformers.TrainingArguments`]. + + 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 + + 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. + + use_liger_loss (`bool`, *optional*, defaults to `False`): + Whether to use Liger loss. + base_model_attribute_name (`str`, *optional*, defaults to `"model"`): + Name of the attribute in the model that contains the base model. This is used to get the base model from + the model when the model does not have a `get_decoder` method in the case when `use_liger_loss` is `True`. + 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, + hub_revision = None, + 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, + liger_kernel_config = None, + 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', + use_liger_loss = False, + base_model_attribute_name = 'model', + 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, + 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, + 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, + use_liger_loss = use_liger_loss, + base_model_attribute_name = base_model_attribute_name, + 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): + """""" + + _tag_names = ["trl", "dpo"] + + def __init__( + self, + model: Union[str, nn.Module, PreTrainedModel], + ref_model: Optional[Union[PreTrainedModel, nn.Module, str]] = None, + args: Optional[DPOConfig] = None, + data_collator: Optional[DataCollator] = None, # type: ignore + train_dataset: Optional[Union[Dataset, IterableDataset]] = None, + eval_dataset: Optional[Union[Dataset, IterableDataset, dict[str, Union[Dataset, IterableDataset]]]] = None, + processing_class: Optional[ + Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin] + ] = None, + compute_metrics: Optional[Callable[[EvalLoopOutput], dict]] = None, + callbacks: Optional[list[TrainerCallback]] = None, + optimizers: tuple[Optional[torch.optim.Optimizer], Optional[torch.optim.lr_scheduler.LambdaLR]] = (None, None), + optimizer_cls_and_kwargs: Optional[tuple[type[torch.optim.Optimizer], dict[str, Any]]] = None, + preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None, + peft_config: Optional["PeftConfig"] = None, + ): + # Args + model_id = model if isinstance(model, str) else model.config._name_or_path + if args is None: + model_name = model_id.split("/")[-1] + args = DPOConfig(f"{model_name}-DPO") + + # Handle the tokenizer + if processing_class is None: + processing_class = AutoTokenizer.from_pretrained(model_id) + + 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." + ) + + # Model + 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 not None and not isinstance(model, str): + warnings.warn( + "You passed model_init_kwargs to the `DPOConfig`, but your model is already instantiated. " + "The `model_init_kwargs` will be ignored." + ) + if isinstance(model, str): + model = self._create_model_from_path(model, args) + + if args.ref_model_init_kwargs is not None and not isinstance(ref_model, str): + warnings.warn( + "You passed ref_model_init_kwargs to the `DPOConfig`, but your ref_model is already instantiated. " + "The `ref_model_init_kwargs` will be ignored." + ) + if isinstance(ref_model, str): + ref_model = self._create_model_from_path(ref_model, args, is_ref=True) + + # PEFT configuration and model wrapping + model = self._prepare_peft_model(model, ref_model, peft_config, args) + + 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) + + # 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) + + # Liger kernel + if args.use_liger_loss: + if not is_liger_kernel_available(): + raise ImportError( + "You set `use_liger_loss=True` but the liger kernel is not available. " + "Please install liger-kernel first: `pip install liger-kernel`" + ) + if args.loss_type != "sigmoid": + raise ValueError( + "You set `use_liger_loss=True` but the loss type is not `sigmoid`. " + "Please set `loss_type='sigmoid'` to use the liger kernel." + ) + self.dpo_loss_fn = LigerFusedLinearDPOLoss( + ignore_index=args.label_pad_token_id, + beta=args.beta, + use_ref_model=not args.reference_free, + average_log_prob=False, + ) + # 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 + + # Data collator + if data_collator is None: + data_collator = DataCollatorForPreference(pad_token_id=self.padding_value) + + 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." + ) + if args.per_device_train_batch_size == 1: + warnings.warn( + "You are using a per_device_train_batch_size of 1 with padding-free training. Using a batch size " + "of 1 anihilate the benefits of padding-free training. Please consider increasing the batch size " + "to at least 2." + ) + 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 + + # 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, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + optimizer_cls_and_kwargs=optimizer_cls_and_kwargs, + 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 _create_model_from_path(self, model_path: str, args: DPOConfig, is_ref: bool = False) -> PreTrainedModel: + """Creates a model from a path or model identifier.""" + if not is_ref: + model_init_kwargs = args.model_init_kwargs or {} + else: + model_init_kwargs = args.ref_model_init_kwargs or {} + + # Handle torch dtype + torch_dtype = model_init_kwargs.get("torch_dtype") + if isinstance(torch_dtype, torch.dtype) or torch_dtype == "auto" or torch_dtype is None: + pass # torch_dtype is already a torch.dtype or "auto" or None + elif isinstance(torch_dtype, str): # it's a str, but not "auto" + torch_dtype = getattr(torch, torch_dtype) + model_init_kwargs["torch_dtype"] = torch_dtype + else: + raise ValueError( + "Invalid `torch_dtype` passed to `DPOConfig`. Expected either 'auto' or a string representing " + f"a `torch.dtype` (e.g., 'float32'), but got {torch_dtype}." + ) + # Disable caching if gradient checkpointing is enabled (not supported) + # if args.gradient_checkpointing: + # model_init_kwargs["use_cache"] = False + + # Create model + model = AutoModelForCausalLM.from_pretrained(model_path, **model_init_kwargs) + return model + + def _prepare_peft_model( + self, model: PreTrainedModel, ref_model: PreTrainedModel, peft_config: Any, args: DPOConfig + ) -> PreTrainedModel: + """Prepares a model for PEFT training.""" + # 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) + + else: + model = self._prepare_gradient_checkpointing(model, args) + + # get peft model with the given config + model = get_peft_model(model, peft_config) + 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 + + else: + model = self._prepare_gradient_checkpointing(model, args) + + return model + + def _prepare_gradient_checkpointing(self, model: PreTrainedModel, args: DPOConfig): + """Prepare the gradienting checkpointing for the model.""" + # 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. + if args.gradient_checkpointing: + # 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) + + return model + + 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 = {} + if isinstance(dataset, Dataset): # IterableDataset does not support num_proc nor writer_batch_size + map_kwargs["num_proc"] = args.dataset_num_proc + map_kwargs["writer_batch_size"] = 10 + + 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.""" + compte_ref_context_manager = ( + autocast(self.accelerator.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 _compute_loss_liger(self, model: nn.Module, batch: dict[str, Union[list, torch.LongTensor]]): + unwrapped_model = self.accelerator.unwrap_model(model) + concatenated_batch = self.concatenated_inputs(batch, padding_value=self.padding_value) + + model_kwargs = {} + 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_attention_mask = concatenated_batch["prompt_attention_mask"] + completion_attention_mask = concatenated_batch["completion_attention_mask"] + + if self.is_encoder_decoder: + # 1. Get encoder outputs + encoder_outputs = unwrapped_model.get_encoder()( + concatenated_batch["prompt_input_ids"], + attention_mask=concatenated_batch["prompt_attention_mask"], + return_dict=True, + ) + # 2. Prepare decoder inputs + decoder_input_ids = shift_tokens_right( + concatenated_batch["completion_input_ids"], + unwrapped_model.config.decoder_start_token_id, + ) + # 3. Get decoder outputs + decoder_outputs = unwrapped_model.get_decoder()( + input_ids=decoder_input_ids, + attention_mask=concatenated_batch["completion_attention_mask"], + encoder_hidden_states=encoder_outputs.last_hidden_state, + encoder_attention_mask=concatenated_batch["prompt_attention_mask"], + use_cache=False, + ) + hidden_states = decoder_outputs.last_hidden_state + + ref_hidden_states = None + if not self.reference_free and self.ref_model is not None: + unwrapped_ref_model = self.accelerator.unwrap_model(self.ref_model) + ref_encoder_outputs = unwrapped_ref_model.get_encoder()( + concatenated_batch["prompt_input_ids"], + attention_mask=concatenated_batch["prompt_attention_mask"], + return_dict=True, + ) + ref_decoder_outputs = unwrapped_ref_model.get_decoder()( + input_ids=decoder_input_ids, + attention_mask=concatenated_batch["completion_attention_mask"], + encoder_hidden_states=ref_encoder_outputs.last_hidden_state, + encoder_attention_mask=concatenated_batch["prompt_attention_mask"], + use_cache=False, + ) + ref_hidden_states = ref_decoder_outputs.last_hidden_state + elif not self.reference_free: + with self.null_ref_context(): + ref_encoder_outputs = unwrapped_model.get_encoder()( + concatenated_batch["prompt_input_ids"], + attention_mask=concatenated_batch["prompt_attention_mask"], + return_dict=True, + ) + ref_decoder_outputs = unwrapped_model.get_decoder()( + input_ids=decoder_input_ids, + attention_mask=concatenated_batch["completion_attention_mask"], + encoder_hidden_states=ref_encoder_outputs.last_hidden_state, + encoder_attention_mask=concatenated_batch["prompt_attention_mask"], + use_cache=False, + ) + ref_hidden_states = ref_decoder_outputs.last_hidden_state + + labels = concatenated_batch["completion_input_ids"] + loss_mask = completion_attention_mask.bool() + else: + # For decoder-only models + input_ids = torch.cat( + (concatenated_batch["prompt_input_ids"], concatenated_batch["completion_input_ids"]), dim=1 + ) + attention_mask = torch.cat( + (concatenated_batch["prompt_attention_mask"], concatenated_batch["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) + + # Add logits_to_keep optimization + if self.use_logits_to_keep: + first_compute_index = loss_mask.nonzero(as_tuple=True)[1].min() + logits_to_keep = (loss_mask.shape[1] - first_compute_index).item() + 1 + model_kwargs["logits_to_keep"] = logits_to_keep + + model_kwargs["output_hidden_states"] = True + + # Add padding-free training support + if self.padding_free: + 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 + + # Get the base model outputs (before LM head) + if hasattr(unwrapped_model, "get_decoder"): + base_model = unwrapped_model.get_decoder() + else: + base_model = getattr(unwrapped_model, self.args.base_model_attribute_name, unwrapped_model) + + outputs = base_model( + input_ids, + use_cache=False, + **model_kwargs, + ) + hidden_states = outputs.last_hidden_state[:, :-1] + + # Get reference hidden states if needed + ref_hidden_states = None + if not self.reference_free and self.ref_model is not None: + unwrapped_ref_model = self.accelerator.unwrap_model(self.ref_model) + if hasattr(unwrapped_ref_model, "get_decoder"): + ref_base_model = unwrapped_ref_model.get_decoder() + else: + ref_base_model = getattr( + unwrapped_ref_model, self.args.base_model_attribute_name, unwrapped_ref_model + ) + + ref_outputs = ref_base_model( + input_ids, + use_cache=False, + **model_kwargs, + ) + ref_hidden_states = ref_outputs.last_hidden_state[:, :-1] + elif not self.reference_free: + if hasattr(unwrapped_model, "get_decoder"): + ref_base_model = unwrapped_model.get_decoder() + else: + ref_base_model = getattr(unwrapped_model, self.args.base_model_attribute_name, unwrapped_model) + with self.null_ref_context(): + ref_outputs = ref_base_model( + input_ids, + use_cache=False, + **model_kwargs, + ) + ref_hidden_states = ref_outputs.last_hidden_state[:, :-1] + + masked_input_ids = torch.where(loss_mask != 0, input_ids, self.label_pad_token_id) + labels = masked_input_ids[:, 1:] # Shift right for casual LM + + # Get the LM head + lm_head = unwrapped_model.get_output_embeddings() + + # Get reference model weights if needed + ref_weight = None + ref_bias = None + if not self.reference_free: + if self.ref_model is not None: + unwrapped_ref_model = self.accelerator.unwrap_model(self.ref_model) + ref_lm_head = unwrapped_ref_model.get_output_embeddings() + else: + with self.null_ref_context(): + ref_lm_head = unwrapped_model.get_output_embeddings() + ref_weight = ref_lm_head.weight + ref_bias = ref_lm_head.bias if hasattr(ref_lm_head, "bias") else None + + # Compute loss using Liger kernel + loss_output = self.dpo_loss_fn( + lm_head.weight, + hidden_states, + labels, + bias=lm_head.bias if hasattr(lm_head, "bias") else None, + ref_input=ref_hidden_states if not self.reference_free else None, + ref_weight=ref_weight if not self.reference_free else None, + ref_bias=ref_bias if not self.reference_free else None, + ) + ( + loss, + (chosen_logps, rejected_logps, chosen_logits_mean, rejected_logits_mean, nll_loss, *aux_outputs), + ) = loss_output + + output = { + "loss": loss, + "chosen_logps": chosen_logps, + "rejected_logps": rejected_logps, + "mean_chosen_logits": chosen_logits_mean, + "mean_rejected_logits": rejected_logits_mean, + "nll_loss": nll_loss, + "chosen_rewards": aux_outputs[0], + "rejected_rewards": aux_outputs[1], + } + if self.aux_loss_enabled: + output["aux_loss"] = outputs.aux_loss + + return output + + 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 + + model_kwargs["output_hidden_states"] = True + + 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[:, 1:].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, :-1] if not self.is_encoder_decoder else logits[:num_examples] + chosen_labels = labels[:num_examples, :-1] if not self.is_encoder_decoder else 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 = {} + + if self.args.use_liger_loss: + model_output = self._compute_loss_liger(model, batch) + losses = model_output["loss"] + chosen_rewards = model_output["chosen_rewards"] + rejected_rewards = model_output["rejected_rewards"] + else: + 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]]]: + compute_loss_context_manager = ( + autocast(self.accelerator.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. + generate_context_manager = ( + autocast(self.accelerator.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 = [] + + prediction_context_manager = ( + autocast(self.accelerator.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] + return super().log(logs, start_time) + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) + + 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 + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + 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): + """ + + Trainer for Direct Preference Optimization (DPO) method. + + This class is a wrapper around the [`transformers.Trainer`] class and inherits all of its attributes and methods. + + Args: + model (`Union[str, PreTrainedModel]`): + Model to be trained. Can be either: + + - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or a + path to a *directory* containing model weights saved using + [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded + using [`~transformers.AutoModelForCausalLM.from_pretrained`] with the keyword arguments in + `args.model_init_kwargs`. + - A [`~transformers.PreTrainedModel`] object. Only causal language models are supported. + 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`], *optional*, defaults to `None`): + Configuration for this trainer. If `None`, a default configuration is used. + data_collator (`DataCollator`, *optional*): + Function to use to form a batch from a list of elements of the processed `train_dataset` or `eval_dataset`. + Will default to [`DataCollatorForPreference`]. + train_dataset ([`~datasets.Dataset`] or [`~datasets.IterableDataset`]): + Dataset to use for training. DPO supports [preference](#preference) type and. The format of the samples can + be either: + + - [Standard](dataset_formats#standard): Each sample contains plain text. + - [Conversational](dataset_formats#conversational): Each sample contains structured messages (e.g., role + and content). + eval_dataset ([`~datasets.Dataset`], [`~datasets.IterableDataset`] or `dict[str, Union[Dataset, IterableDataset]]`): + Dataset to use for evaluation. It must meet the same requirements as `train_dataset`. + processing_class ([`~transformers.PreTrainedTokenizerBase`], *optional*, defaults to `None`): + Processing class used to process the data. If `None`, the processing class is loaded from the model's name + with [`~transformers.AutoTokenizer.from_pretrained`]. + compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): + The function that will be used to compute metrics at evaluation. Must take a [`EvalPrediction`] and return + a dictionary string to metric values. *Note* When passing TrainingArgs with `batch_eval_metrics` set to + `True`, your compute_metrics function must take a boolean `compute_result` argument. This will be triggered + after the last eval batch to signal that the function needs to calculate and return the global summary + statistics rather than accumulating the batch-level statistics. + callbacks (list of [`~transformers.TrainerCallback`], *optional*, defaults to `None`): + List of callbacks to customize the training loop. Will add those to the list of default callbacks detailed + in [here](https://huggingface.co/docs/transformers/main_classes/callback). + + If you want to remove one of the default callbacks used, use the [`~transformers.Trainer.remove_callback`] + method. + optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`, *optional*, defaults to `(None, None)`): + A tuple containing the optimizer and the scheduler to use. Will default to an instance of [`AdamW`] on your + model and a scheduler given by [`get_linear_schedule_with_warmup`] controlled by `args`. + optimizer_cls_and_kwargs (`Tuple[Type[torch.optim.Optimizer], Dict[str, Any]]`, *optional*, defaults to `None`): + A tuple containing the optimizer class and keyword arguments to use. Overrides `optim` and `optim_args` in + `args`. Incompatible with the `optimizers` argument. + preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`, *optional*, defaults to `None`): + A function that preprocess the logits right before caching them at each evaluation step. Must take two + tensors, the logits and the labels, and return the logits once processed as desired. The modifications made + by this function will be reflected in the predictions received by `compute_metrics`. + + Note that the labels (second parameter) will be `None` if the dataset does not have them. + peft_config ([`~peft.PeftConfig`], *optional*, defaults to `None`): + PEFT configuration used to wrap the model. If `None`, the model is not wrapped. + + """ + def __init__( + self, + model, + ref_model = None, + args = None, + data_collator = None, + train_dataset = None, + eval_dataset = None, + processing_class = None, + compute_metrics = None, + callbacks = None, + optimizer_cls_and_kwargs = None, + preprocess_logits_for_metrics = None, + peft_config = None, + **kwargs + ): + if args is None: args = UnslothDPOConfig() + 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 + 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) + 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 + 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, + compute_metrics = compute_metrics, + callbacks = callbacks, + optimizer_cls_and_kwargs = optimizer_cls_and_kwargs, + 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 diff --git a/unsloth_compiled_cache/UnslothGKDTrainer.py b/unsloth_compiled_cache/UnslothGKDTrainer.py new file mode 100644 index 0000000..5b57477 --- /dev/null +++ b/unsloth_compiled_cache/UnslothGKDTrainer.py @@ -0,0 +1,860 @@ +""" +2025.7.1 +2025.7.1 +4.53.1 +0.19.1 +__UNSLOTH_VERSIONING__ +""" +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from trl.trainer.gkd_trainer import (Any, AutoModelForCausalLM, BaseImageProcessor, Callable, DataCollator, DataCollatorForChatML, Dataset, EvalPrediction, F, FeatureExtractionMixin, GKDConfig, GKDTrainer, GenerationConfig, Optional, PeftConfig, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, SFTTrainer, TrainerCallback, Union, disable_dropout_in_model, empty_cache, generate_model_card, get_comet_experiment_url, is_wandb_available, nn, os, prepare_deepspeed, random, textwrap, torch, 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 +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 UnslothGKDConfig(GKDConfig): + """ + + Configuration class for [`GKDTrainer`]. + + This class includes only the parameters that are specific to GKD training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] and [`SFTConfig`] documentation. + + Args: + temperature (`float`, *optional*, defaults to `0.9`): + Temperature for sampling. The higher the temperature, the more random the completions. + lmbda (`float`, *optional*, defaults to `0.5`): + Lambda parameter that controls the student data fraction (i.e., the proportion of on-policy + student-generated outputs). + beta (`float`, *optional*, defaults to `0.5`): + Interpolation coefficient between `0.0` and `1.0` of the Generalized Jensen-Shannon Divergence loss. When + beta is `0.0`, the loss is the KL divergence. When beta is `1.0`, the loss is the Inverse KL Divergence. + max_new_tokens (`int`, *optional*, defaults to `128`): + Maximum number of tokens to generate per completion. + teacher_model_name_or_path (`str` or `None`, *optional*, defaults to `None`): + Model name or path of the teacher model. If `None`, the teacher model will be the same as the model being + trained. + teacher_model_init_kwargs (`dict[str, Any]]` or `None`, *optional*, defaults to `None`): + Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the teacher model + from a string. + disable_dropout (`bool`, *optional*, defaults to `True`): + Whether to disable dropout in the model. + seq_kd (`bool`, *optional*, defaults to `False`): + Seq_kd parameter that controls whether to perform Sequence-Level KD (can be viewed as supervised FT on + teacher-generated output). + + """ + 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, + hub_revision = None, + 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, + liger_kernel_config = None, + eval_use_gather_object = False, + average_tokens_across_devices = True, + model_init_kwargs = None, + chat_template_path = None, + dataset_text_field = 'text', + dataset_kwargs = None, + dataset_num_proc = None, + eos_token = None, + pad_token = None, + max_length = 1024, + packing = False, + packing_strategy = 'ffd', + padding_free = False, + pad_to_multiple_of = None, + eval_packing = None, + completion_only_loss = None, + assistant_only_loss = False, + activation_offloading = False, + max_seq_length = None, + temperature = 0.9, + lmbda = 0.5, + beta = 0.5, + max_new_tokens = 128, + teacher_model_name_or_path = None, + teacher_model_init_kwargs = None, + disable_dropout = True, + seq_kd = 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() + if temperature <= 0: + raise MathError('Unsloth: Please set a positive non-zero temperature since your results will be wrong.') + elif temperature >= 10: + raise MathError('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, + 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, + 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, + model_init_kwargs = model_init_kwargs, + chat_template_path = chat_template_path, + dataset_text_field = dataset_text_field, + dataset_kwargs = dataset_kwargs, + dataset_num_proc = dataset_num_proc, + eos_token = eos_token, + pad_token = pad_token, + max_length = max_length, + packing = packing, + packing_strategy = packing_strategy, + padding_free = padding_free, + pad_to_multiple_of = pad_to_multiple_of, + eval_packing = eval_packing, + completion_only_loss = completion_only_loss, + assistant_only_loss = assistant_only_loss, + activation_offloading = activation_offloading, + max_seq_length = max_seq_length, + temperature = temperature, + lmbda = lmbda, + beta = beta, + max_new_tokens = max_new_tokens, + teacher_model_name_or_path = teacher_model_name_or_path, + teacher_model_init_kwargs = teacher_model_init_kwargs, + disable_dropout = disable_dropout, + seq_kd = seq_kd,**kwargs) + self.vllm_sampling_params = vllm_sampling_params + self.unsloth_num_chunks = unsloth_num_chunks +pass + +class _UnslothGKDTrainer(SFTTrainer): + _tag_names = ["trl", "gkd"] + + def __init__( + self, + model: Optional[Union[PreTrainedModel, nn.Module, str]] = None, + teacher_model: Union[PreTrainedModel, nn.Module, str] = None, + args: Optional[GKDConfig] = None, + data_collator: Optional[DataCollator] = None, # type: ignore + train_dataset: Optional[Dataset] = None, + eval_dataset: Optional[Union[Dataset, dict[str, Dataset]]] = None, + processing_class: Optional[ + Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin] + ] = 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, + peft_config: Optional["PeftConfig"] = None, + formatting_func: Optional[Callable] = None, + ): + # add remove_unused_columns=False to the dataclass args + args.remove_unused_columns = False + data_collator = DataCollatorForChatML(tokenizer=processing_class, max_length=args.max_length) + + super().__init__( + model, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + peft_config=peft_config, + formatting_func=formatting_func, + ) + + if args.teacher_model_init_kwargs is None: + teacher_model_init_kwargs = {} + elif not isinstance(teacher_model, str): + raise ValueError( + "You passed teacher_model_init_kwargs to the GKDConfig, but your teacher_model is already instantiated." + ) + else: + teacher_model_init_kwargs = args.teacher_model_init_kwargs + teacher_model_init_kwargs["torch_dtype"] = ( + teacher_model_init_kwargs["torch_dtype"] + if teacher_model_init_kwargs["torch_dtype"] in ["auto", None] + else getattr(torch, teacher_model_init_kwargs["torch_dtype"]) + ) + + if isinstance(teacher_model, str): + teacher_model = AutoModelForCausalLM.from_pretrained(teacher_model, **teacher_model_init_kwargs) + + # Disable dropout in the model + if args.disable_dropout: + disable_dropout_in_model(self.model) + + if self.is_deepspeed_enabled: + self.teacher_model = prepare_deepspeed(teacher_model, self.accelerator) + else: + self.teacher_model = self.accelerator.prepare_model(teacher_model, evaluation_mode=True) + + self.lmbda = args.lmbda + self.beta = args.beta + self.temperature = args.temperature + self.seq_kd = args.seq_kd + + self.generation_config = GenerationConfig( + max_new_tokens=args.max_new_tokens, + temperature=args.temperature, + do_sample=True, + top_k=0, + use_cache=False if args.gradient_checkpointing else True, + pad_token_id=self.processing_class.pad_token_id, + ) + # Set custom EOS tokens if they are specified by the model's generation + # config. This is important for models with the Llama 3 chat template, + # which use special tokens <|eot_id|> and <|eom_id|> to mark the end of + # turns or messages. + if ( + hasattr(self.model.generation_config, "eos_token_id") + and self.model.generation_config.eos_token_id is not None + ): + self.generation_config.eos_token_id = self.model.generation_config.eos_token_id + + @staticmethod + def generalized_jsd_loss( + student_logits, teacher_logits, labels=None, beta=0.5, temperature=1.0, reduction="batchmean" + ): + """ + Compute the generalized Jensen-Shannon Divergence loss for knowledge distillation using F.kl_div. See Eq. (1) + of https://huggingface.co/papers/2306.13649 for the definition. + + Args: + student_logits: + Tensor of shape (batch_size, sequence_length, vocab_size) + teacher_logits: + Tensor of shape (batch_size, sequence_length, vocab_size) + labels: + Tensor of shape (batch_size, sequence_length) with -100 for padding tokens to ignore when computing + loss + beta: + Interpolation coefficient between 0 and 1 (default: 0.5) + temperature: + Softmax temperature (default: 1.0) + reduction: + Specifies the reduction to apply to the output (default: 'batchmean') + + Returns: + loss: Scalar tensor with the generalized JSD loss + """ + + # Apply temperature scaling + student_logits = student_logits / temperature + teacher_logits = teacher_logits / temperature + + # Compute log probabilities for student and probabilities for teacher + student_log_probs = F.log_softmax(student_logits, dim=-1) + teacher_log_probs = F.log_softmax(teacher_logits, dim=-1) + + if beta == 0: + jsd = F.kl_div(student_log_probs, teacher_log_probs, reduction="none", log_target=True) + elif beta == 1: + jsd = F.kl_div(teacher_log_probs, student_log_probs, reduction="none", log_target=True) + else: + # Compute the log of the mixture distribution + # log(a + b) = log(exp(log(a)) + exp(log(b))) -> for mixture + beta = torch.tensor(beta, dtype=student_log_probs.dtype) + mixture_log_probs = torch.logsumexp( + torch.stack([student_log_probs + torch.log(1 - beta), teacher_log_probs + torch.log(beta)]), + dim=0, + ) + + # Compute KL divergences using F.kl_div + # PyTorch differs from the standard mathematical definition, so the order of the probability distributions is swapped compared to that defined in the paper. + kl_teacher = F.kl_div(mixture_log_probs, teacher_log_probs, reduction="none", log_target=True) + kl_student = F.kl_div(mixture_log_probs, student_log_probs, reduction="none", log_target=True) + + # Compute the Generalized Jensen-Shannon Divergence + jsd = beta * kl_teacher + (1 - beta) * kl_student + + # Masking + if labels is not None: + mask = labels != -100 + jsd = jsd[mask] + + # Apply reduction + if reduction == "batchmean": + return jsd.sum() / mask.sum() if labels is not None else jsd.sum() / (jsd.size(0) * jsd.size(1)) + elif reduction == "sum": + return jsd.sum() + elif reduction == "mean": + return jsd.mean() + else: + return jsd + + def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): + # compute student output + outputs_student = model( + input_ids=inputs["input_ids"], + attention_mask=inputs["attention_mask"], + ) + + # compute teacher output in eval mode + self.teacher_model.eval() + with torch.no_grad(): + outputs_teacher = self.teacher_model( + input_ids=inputs["input_ids"], + attention_mask=inputs["attention_mask"], + ) + + # slice the logits for the generated tokens using the inputs["prompts"] lengths + prompt_lengths = inputs["prompts"].shape[1] + shifted_student_logits = outputs_student.logits[:, prompt_lengths - 1 : -1, :] + shifted_teacher_logits = outputs_teacher.logits[:, prompt_lengths - 1 : -1, :] + shifted_labels = inputs["labels"][:, prompt_lengths:] + + # compute loss + loss = self.generalized_jsd_loss( + student_logits=shifted_student_logits, + teacher_logits=shifted_teacher_logits, + labels=shifted_labels, + beta=self.beta, + ) + + # empty cache + empty_cache() + + # Return loss + return (loss, outputs_student) if return_outputs else loss + + @staticmethod + def generate_on_policy_outputs(model, inputs, generation_config, pad_token_id=None): + # Generate output with respect to the prompt only + generated_outputs = model.generate( + input_ids=inputs["prompts"], + attention_mask=inputs.get("prompt_attention_mask", None), + generation_config=generation_config, + return_dict_in_generate=True, + ) + + # Get the generated token IDs + generated_tokens = generated_outputs.sequences + # Calculate new attention mask + new_attention_mask = torch.ones_like(generated_tokens) + new_labels = generated_tokens.clone() + + # If there's pad_token_id, set attention mask to 0 for padding tokens + if pad_token_id is not None: + new_labels[new_labels == pad_token_id] = -100 + new_attention_mask[generated_tokens == pad_token_id] = 0 + + return generated_tokens, new_attention_mask, new_labels + + def training_step( + self, model: nn.Module, inputs: dict[str, Union[torch.Tensor, Any]], num_items_in_batch: Optional[int] = None + ) -> torch.Tensor: + """ + Perform a training step for the Generalized Knowledge Distillation (GKD) model. + + This method implements the on-policy learning approach described in the GKD paper. With probability + `self.lmbda`, it generates new responses using the student model, which are then used for training instead of + the original inputs. + """ + if self.seq_kd: + with unwrap_model_for_generation(self.teacher_model, self.accelerator) as unwrapped_model: + new_input_ids, new_attention_mask, new_labels = self.generate_on_policy_outputs( + unwrapped_model, inputs, self.generation_config, self.processing_class.pad_token_id + ) + inputs["input_ids"] = new_input_ids + inputs["attention_mask"] = new_attention_mask + inputs["labels"] = new_labels + if random.random() <= self.lmbda: + with unwrap_model_for_generation(model, self.accelerator) as unwrapped_model: + new_input_ids, new_attention_mask, new_labels = self.generate_on_policy_outputs( + unwrapped_model, inputs, self.generation_config, self.processing_class.pad_token_id + ) + inputs["input_ids"] = new_input_ids + inputs["attention_mask"] = new_attention_mask + inputs["labels"] = new_labels + + loss = super().training_step(model, inputs, num_items_in_batch) + return loss + + 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 + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + citation = textwrap.dedent("""\ + @inproceedings{agarwal2024on-policy, + title = {{On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes}}, + author = {Rishabh Agarwal and Nino Vieillard and Yongchao Zhou and Piotr Stanczyk and Sabela Ramos Garea and Matthieu Geist and Olivier Bachem}, + year = 2024, + booktitle = {The Twelfth International Conference on Learning Representations, {ICLR} 2024, Vienna, Austria, May 7-11, 2024}, + publisher = {OpenReview.net}, + url = {https://openreview.net/forum?id=3zKtaqxLhW}, + }""") + + 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="GKD", + trainer_citation=citation, + paper_title="On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes", + paper_id="2306.13649", + ) + + model_card.save(os.path.join(self.args.output_dir, "README.md")) +class UnslothGKDTrainer(_UnslothGKDTrainer): + """ + + """ + def __init__( + self, + model = None, + teacher_model = None, + args = None, + data_collator = None, + train_dataset = None, + eval_dataset = None, + processing_class = None, + compute_metrics = None, + callbacks = None, + preprocess_logits_for_metrics = None, + peft_config = None, + formatting_func = None, + **kwargs + ): + if args is None: args = UnslothGKDConfig() + 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 + 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) + 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 + 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('gkd_trainer', other_metrics) + + super().__init__( + model = model, + teacher_model = teacher_model, + args = args, + data_collator = data_collator, + train_dataset = train_dataset, + eval_dataset = eval_dataset, + processing_class = processing_class, + compute_metrics = compute_metrics, + callbacks = callbacks, + preprocess_logits_for_metrics = preprocess_logits_for_metrics, + peft_config = peft_config, + formatting_func = formatting_func,**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 diff --git a/unsloth_compiled_cache/UnslothGRPOTrainer.py b/unsloth_compiled_cache/UnslothGRPOTrainer.py new file mode 100644 index 0000000..40d3463 --- /dev/null +++ b/unsloth_compiled_cache/UnslothGRPOTrainer.py @@ -0,0 +1,2497 @@ +""" +2025.7.1 +2025.7.1 +4.53.1 +0.19.1 +__UNSLOTH_VERSIONING__ +""" +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from trl.trainer.grpo_trainer import (Any, AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer, DataLoader, Dataset, FSDP, GRPOConfig, GRPOTrainer, GenerationConfig, IterableDataset, Optional, Path, PeftConfig, PreTrainedModel, PreTrainedTokenizerBase, RepeatSampler, RewardFunc, Sampler, SyncRefModelCallback, Trainer, TrainerCallback, Union, VLLMClient, _ForwardRedirection, apply_chat_template, broadcast_object_list, datasets, defaultdict, deque, disable_dropout_in_model, gather, gather_object, generate_model_card, get_comet_experiment_url, identity, is_conversational, is_datasets_available, is_liger_kernel_available, is_peft_available, is_peft_model, is_rich_available, is_vllm_available, is_wandb_available, maybe_apply_chat_template, nanmax, nanmin, nanstd, nn, nullcontext, os, pad, partial, prepare_deepspeed, prepare_fsdp, print_prompt_completions_sample, profiling_context, profiling_decorator, re, seed_worker, set_seed, shuffle_tensor_dict, split_tensor_dict, textwrap, torch, transformers, unwrap_model_for_generation, version, warnings, Any, FSDP, Union, apply_chat_template, broadcast_object_list, gather, gather_object, is_conversational, maybe_apply_chat_template, nanstd, nullcontext, os, pad, profiling_context, re, torch, unwrap_model_for_generation, os, re, torch, transformers, re, Any, Union, os, profiling_decorator, re, shuffle_tensor_dict, split_tensor_dict, torch, Optional, PreTrainedModel, Trainer, is_peft_available, os, re, torch, FSDP, nn, os, re, GRPOTrainer, Trainer, gather, os, re, torch) + + +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 + +def grpo_compute_loss( + ref_logits, + new_logits, + old_logits, + input_ids, + mask, + beta, + advantages, + **kwargs +): + # All Unsloth Zoo code licensed under LGPLv3 + # Set defaults for optional arguments + loss_type = kwargs.get("loss_type", "grpo") + epsilon_low = kwargs.get("epsilon_low", 0.2) + epsilon_high = kwargs.get("epsilon_high", 0.2) + max_completion_length = kwargs.get("max_completion_length", 8192) + delta = kwargs.get("delta", None) + temperature = kwargs.get("temperature", 1.0) + logit_scale_multiply = kwargs.get("logit_scale_multiply", 0.0) + logit_scale_divide = kwargs.get("logit_scale_divide", 0.0) + logit_softcapping = kwargs.get("logit_softcapping", 0.0) + + input_ids = input_ids.unsqueeze(-1) + + # Optional logit softcapping and logit dividing + if logit_scale_multiply != 0: new_logits = new_logits * logit_scale_multiply + if logit_scale_divide != 0: new_logits = new_logits / logit_scale_divide + if logit_softcapping != 0: new_logits = new_logits * torch.tanh(new_logits / logit_softcapping) + + new_logits = new_logits.to(torch.float32) + # See https://huggingface.co/blog/the_n_implementation_details_of_rlhf_with_ppo#policy-training-implementation-details + if temperature != 1.0: new_logits = new_logits / temperature + new_x = torch.gather(new_logits, dim = -1, index = input_ids).squeeze(-1) + new = new_x - torch.logsumexp(new_logits, dim = -1) + + # x_i - logsumexp(x_i) + with torch.no_grad(): + if beta != 0.0: + assert ref_logits is not None, "ref_logits should not be None when beta != 0.0" + + # Optional logit softcapping and logit dividing + if logit_scale_multiply != 0: ref_logits = ref_logits * logit_scale_multiply + if logit_scale_divide != 0: ref_logits = ref_logits / logit_scale_divide + if logit_softcapping != 0: ref_logits = ref_logits * torch.tanh(ref_logits / logit_softcapping) + + ref_logits = ref_logits.to(torch.float32) + # See https://huggingface.co/blog/the_n_implementation_details_of_rlhf_with_ppo#policy-training-implementation-details + if temperature != 1.0: ref_logits = ref_logits / temperature + ref_x = torch.gather(ref_logits, dim = -1, index = input_ids).squeeze(-1) + ref = ref_x - torch.logsumexp(ref_logits, dim = -1) + pass + + if old_logits is not None: + # Optional logit softcapping and logit dividing + if logit_scale_multiply != 0: old_logits = old_logits * logit_scale_multiply + if logit_scale_divide != 0: old_logits = old_logits / logit_scale_divide + if logit_softcapping != 0: old_logits = old_logits * torch.tanh(old_logits / logit_softcapping) + + old_logits = old_logits.to(torch.float32) + # See https://huggingface.co/blog/the_n_implementation_details_of_rlhf_with_ppo#policy-training-implementation-details + if temperature != 1.0: old_logits = old_logits / temperature + old_x = torch.gather(old_logits, dim = -1, index = input_ids).squeeze(-1) + old = old_x - torch.logsumexp(old_logits, dim = -1) + pass + pass + + # Reverse KL + # Note that this is a low variance low bias estimator for the KL divergence as used in GRPO paper + if beta != 0.0: + kl_i = torch.exp(ref - new) - (ref - new) - 1.0 + + else: + kl_i = 0.0 # set it to 0 to not effect the downstream computation + # Full correct reverse KL divergence?? Missing term maybe? + # kl_i = torch.exp(new) * kl_i + + # Below is forward KL (normal KL) + # kl_i = torch.exp(old) * (old - new) + if old_logits is not None: + coef_1 = torch.exp(new - old) + else: + coef_1 = torch.exp(new - new.detach()) + coef_2 = torch.clamp(coef_1, 1 - epsilon_low, 1 + epsilon_high) + + if delta is not None: + loss_1 = torch.clamp(coef_1, max=delta) * advantages.unsqueeze(1) + else: + loss_1 = coef_1 * advantages.unsqueeze(1) + pass + + # Must detach - otherwise gradients are not propagated correctly! + # exp(x - x) == 1 + # loss_i = torch.exp(new - new.detach()) * advantages.unsqueeze(1) + + loss_2 = coef_2 * advantages.unsqueeze(1) + loss_i = -torch.min(loss_1, loss_2) + if beta != 0.0: + loss_i = loss_i + beta * kl_i + + mask = mask.to(torch.float32) + n_mask_per_reward = mask.sum(1) + + # https://github.com/huggingface/trl/blob/main/trl/trainer/grpo_trainer.py#L1363-L1370 + if loss_type == "grpo": + loss = ((loss_i * mask).sum(-1) / mask.sum(-1).clamp(min=1.0)).mean() + elif loss_type == "bnpo": + loss = (loss_i * mask).sum() / mask.sum().clamp(min=1.0) + elif loss_type == "dr_grpo": + loss = (loss_i * mask).sum() / (loss_i.size(0) * max_completion_length) + else: + raise ValueError(f"Unknown loss type: {loss_type}") + + # loss = (loss_i * mask).sum() / mask.sum() + + # Get metrics as well which are folded + with torch.inference_mode(): + completion_length = n_mask_per_reward.mean() + mean_kl_per_reward = (kl_i * mask).sum(1) / n_mask_per_reward + mean_kl = mean_kl_per_reward.mean() + pass + + return loss, completion_length, mean_kl + +class UnslothEfficientGRPO(torch.autograd.Function): + # All Unsloth Zoo code licensed under LGPLv3 + @staticmethod + def forward(ctx, _new_hidden_states, _old_hidden_states, _ref_hidden_states, lm_head, _input_ids, _mask, _advantages, beta, scaler = None, n_chunks = 1, extra_kwargs=None): + if extra_kwargs is None: + extra_kwargs = {} + def compute_loss(new_hidden_states, old_hidden_states, ref_hidden_states, input_ids, mask, advantages, scaling): + new_logits = torch.matmul(new_hidden_states, lm_head.t()) + new_logits = new_logits[:, :-1, :] # exclude the last logit: it corresponds to the next token pred + with torch.no_grad(): + if beta != 0.0: + ref_logits = torch.matmul(ref_hidden_states, lm_head.t()) + ref_logits = ref_logits[:, :-1, :] # exclude the last logit: it corresponds to the next token pred + else: + ref_logits = None + if old_hidden_states is not None: + old_logits = torch.matmul(old_hidden_states, lm_head.t()) + old_logits = old_logits[:, :-1, :] # exclude the last logit: it corresponds to the next token pred + else: + old_logits = None + # if old_hidden_states is not None: + # old_logits = torch.matmul(old_hidden_states, lm_head.t()) #last logit already excluded + # old_logits = old_logits[:, :-1, :] # exclude the last logit: it corresponds to the next token pred + # else: + # old_logits = None + # unsloth_zoo/rl_replacements.py + loss, completion_length, mean_kl = grpo_compute_loss( + ref_logits, + new_logits, + old_logits, + input_ids, + mask, + beta, + advantages, + **extra_kwargs, + ) + + # Scale loss if needed for mixed precision training + scaled_loss = loss * scaling + # Must add .loss.detach otherwise autograd uses 2x VRAM + return scaled_loss, (loss.detach(), completion_length, mean_kl,) + pass + + device =_new_hidden_states.device + grad_inputs = torch.empty_like(_new_hidden_states) + accumulated_loss = torch.zeros(1, device = device) + accumulated_completion_length = torch.zeros(1, device = device) + accumulated_mean_kl = torch.zeros(1, device = device) + + def accumulate_chunk(new_hidden_states_j, old_hidden_states_j, ref_hidden_states_j, input_ids_j, mask_j, advantages_j, scaling): + (chunk_grad_input,), (chunk_loss, (unscaled_loss, chunk_completion_length, chunk_mean_kl,)) = torch.func.grad_and_value( + compute_loss, + argnums = (0,), + has_aux = True, + )(new_hidden_states_j, old_hidden_states_j, ref_hidden_states_j, input_ids_j, mask_j, advantages_j, scaling) + accumulated_loss .add_(unscaled_loss) + accumulated_completion_length.add_(chunk_completion_length) + accumulated_mean_kl .add_(chunk_mean_kl) + return chunk_grad_input + pass + + accumulate_chunk = torch.compile( + accumulate_chunk, + fullgraph = True, + options = torch_compile_options, + ) + + grad_inputs_chunks = torch.chunk(grad_inputs, chunks = n_chunks, dim = 0) + new_hidden_states = torch.chunk(_new_hidden_states, chunks = n_chunks, dim = 0) + if _old_hidden_states is not None: + old_hidden_states = torch.chunk(_old_hidden_states, chunks = n_chunks, dim = 0) + else: + old_hidden_states = [None] * n_chunks + ref_hidden_states = torch.chunk(_ref_hidden_states, chunks = n_chunks, dim = 0) + input_ids = torch.chunk(_input_ids, chunks = n_chunks, dim = 0) + mask = torch.chunk(_mask, chunks = n_chunks, dim = 0) + advantages = torch.chunk(_advantages, chunks = n_chunks, dim = 0) + + # Get mixed precision scaling if seen + scaling = scaler.get_scale() if scaler is not None else 1.0 + + # Force torch.compile to use dynamic shapes for seqlen dim + mark_dynamic = lambda x: torch._dynamo.mark_dynamic(x, 1) + + for (grad_inputs_j, new_hidden_states_j, old_hidden_states_j, ref_hidden_states_j, input_ids_j, mask_j, advantages_j,) in \ + zip(grad_inputs_chunks, new_hidden_states, old_hidden_states, ref_hidden_states, input_ids, mask, advantages): + + mark_dynamic(new_hidden_states_j) + mark_dynamic(ref_hidden_states_j) + if old_hidden_states_j is not None: + mark_dynamic(old_hidden_states_j) + mark_dynamic(input_ids_j) + mark_dynamic(mask_j) + + + grad_inputs_j.copy_(accumulate_chunk(new_hidden_states_j, old_hidden_states_j,ref_hidden_states_j, input_ids_j, mask_j, advantages_j, scaling)) + pass + + grad_inputs .div_(n_chunks) + accumulated_loss .div_(n_chunks) + accumulated_completion_length.div_(n_chunks) + accumulated_mean_kl .div_(n_chunks) + ctx.save_for_backward(grad_inputs) + return ( + accumulated_loss, + accumulated_completion_length, + accumulated_mean_kl, + ) + pass + + @staticmethod + def backward(ctx, grad_output, dcompletion_length, dmean_kl): + (grad_input,) = ctx.saved_tensors + return (grad_input, None, None, None, None, None, None, None, None, None, None) + pass + +def grpo_accumulated_loss( + trainer, + input_ids, + attention_mask, + logits_to_keep, + completion_mask, + advantages, + old_hidden_states, + n_chunks = -1, + **kwargs, +): + # All Unsloth Zoo code licensed under LGPLv3 + bsz, qlen = input_ids.shape + + # Find closest multiple + factors = [i for i in range(1, bsz + 1) if bsz % i == 0] + if n_chunks == -1: n_chunks = bsz + n_chunks = factors[min(np.searchsorted(factors, n_chunks), len(factors)-1)] + + if not hasattr(trainer, '_autocast_dtype'): + trainer._autocast_dtype = torch.float16 if os.environ.get('ACCELERATE_MIXED_PRECISION', 'fp16') == 'fp16' else torch.bfloat16 + if os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1': trainer._autocast_dtype = torch.float16 + pass + os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" + + completion_input_ids = input_ids[:, -logits_to_keep:] + lm_head = trainer.model.get_output_embeddings().weight + + with torch.amp.autocast(device_type = "cuda", dtype = trainer._autocast_dtype): + with torch.inference_mode(), trainer.accelerator.unwrap_model(trainer.model, keep_fp32_wrapper = False).disable_adapter(): + ref_hidden_states = trainer.model( + input_ids = input_ids, + attention_mask = attention_mask, + logits_to_keep = logits_to_keep + 1, + ).logits + pass + new_hidden_states = trainer.model( + input_ids = input_ids, + attention_mask = attention_mask, + logits_to_keep = logits_to_keep + 1, + ).logits + + loss, completion_length, mean_kl = UnslothEfficientGRPO.apply( + new_hidden_states, + old_hidden_states, + ref_hidden_states, + lm_head, + completion_input_ids, + completion_mask, + advantages, + trainer.beta, + trainer.accelerator.scaler, + n_chunks, + kwargs # pass kwargs as a dict + ) + pass + # Must force not returning hidden states but logits otherwise gibberish + os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "0" + return loss, completion_length, mean_kl + + # Old non efficient code path + new_logits = torch.matmul(new_hidden_states, lm_head.t()) + new_logits = new_logits[:, :-1, :] # exclude the last logit: it corresponds to the next token pred + old_logits = torch.matmul(old_hidden_states, lm_head.t()) + old_logits = old_logits[:, :-1, :] # exclude the last logit: it corresponds to the next token pred + loss, completion_length, mean_kl = grpo_compute_loss( + old_logits, + new_logits, + completion_input_ids, + completion_mask, + trainer.beta, + advantages, + ) + return loss, completion_length, mean_kl + pass + +@torch.compile(dynamic = True, fullgraph = True, options = torch_compile_options) +def grpo_compute_loss_slow( + ref_logits, + new_logits, + old_logits, + input_ids, + mask, + beta, + advantages, + **kwargs +): + # All Unsloth Zoo code licensed under LGPLv3 + # Set defaults for optional arguments + loss_type = kwargs.get("loss_type", "grpo") + epsilon_low = kwargs.get("epsilon_low", 0.2) + epsilon_high = kwargs.get("epsilon_high", 0.2) + max_completion_length = kwargs.get("max_completion_length", 8192) + delta = kwargs.get("delta", None) + temperature = kwargs.get("temperature", 1.0) + logit_scale_multiply = kwargs.get("logit_scale_multiply", 0.0) + logit_scale_divide = kwargs.get("logit_scale_divide", 0.0) + logit_softcapping = kwargs.get("logit_softcapping", 0.0) + + input_ids = input_ids.unsqueeze(-1) + + # Optional logit softcapping and logit dividing + if logit_scale_multiply != 0: new_logits = new_logits * logit_scale_multiply + if logit_scale_divide != 0: new_logits = new_logits / logit_scale_divide + if logit_softcapping != 0: new_logits = new_logits * torch.tanh(new_logits / logit_softcapping) + + new_logits = new_logits.to(torch.float32) + # See https://huggingface.co/blog/the_n_implementation_details_of_rlhf_with_ppo#policy-training-implementation-details + if temperature != 1.0: new_logits = new_logits / temperature + new_x = torch.gather(new_logits, dim = -1, index = input_ids).squeeze(-1) + new = new_x - torch.logsumexp(new_logits, dim = -1) + + # x_i - logsumexp(x_i) + with torch.no_grad(): + if beta != 0.0: + assert ref_logits is not None, "ref_logits should not be None when beta != 0.0" + + # Optional logit softcapping and logit dividing + if logit_scale_multiply != 0: ref_logits = ref_logits * logit_scale_multiply + if logit_scale_divide != 0: ref_logits = ref_logits / logit_scale_divide + if logit_softcapping != 0: ref_logits = ref_logits * torch.tanh(ref_logits / logit_softcapping) + + ref_logits = ref_logits.to(torch.float32) + # See https://huggingface.co/blog/the_n_implementation_details_of_rlhf_with_ppo#policy-training-implementation-details + if temperature != 1.0: ref_logits = ref_logits / temperature + ref_x = torch.gather(ref_logits, dim = -1, index = input_ids).squeeze(-1) + ref = ref_x - torch.logsumexp(ref_logits, dim = -1) + pass + + if old_logits is not None: + # Optional logit softcapping and logit dividing + if logit_scale_multiply != 0: old_logits = old_logits * logit_scale_multiply + if logit_scale_divide != 0: old_logits = old_logits / logit_scale_divide + if logit_softcapping != 0: old_logits = old_logits * torch.tanh(old_logits / logit_softcapping) + + old_logits = old_logits.to(torch.float32) + # See https://huggingface.co/blog/the_n_implementation_details_of_rlhf_with_ppo#policy-training-implementation-details + if temperature != 1.0: old_logits = old_logits / temperature + old_x = torch.gather(old_logits, dim = -1, index = input_ids).squeeze(-1) + old = old_x - torch.logsumexp(old_logits, dim = -1) + pass + pass + + # Reverse KL + # Note that this is a low variance low bias estimator for the KL divergence as used in GRPO paper + if beta != 0.0: + kl_i = torch.exp(ref - new) - (ref - new) - 1.0 + + else: + kl_i = 0.0 # set it to 0 to not effect the downstream computation + # Full correct reverse KL divergence?? Missing term maybe? + # kl_i = torch.exp(new) * kl_i + + # Below is forward KL (normal KL) + # kl_i = torch.exp(old) * (old - new) + if old_logits is not None: + coef_1 = torch.exp(new - old) + else: + coef_1 = torch.exp(new - new.detach()) + coef_2 = torch.clamp(coef_1, 1 - epsilon_low, 1 + epsilon_high) + + if delta is not None: + loss_1 = torch.clamp(coef_1, max=delta) * advantages.unsqueeze(1) + else: + loss_1 = coef_1 * advantages.unsqueeze(1) + pass + + # Must detach - otherwise gradients are not propagated correctly! + # exp(x - x) == 1 + # loss_i = torch.exp(new - new.detach()) * advantages.unsqueeze(1) + + loss_2 = coef_2 * advantages.unsqueeze(1) + loss_i = -torch.min(loss_1, loss_2) + if beta != 0.0: + loss_i = loss_i + beta * kl_i + + mask = mask.to(torch.float32) + n_mask_per_reward = mask.sum(1) + + # https://github.com/huggingface/trl/blob/main/trl/trainer/grpo_trainer.py#L1363-L1370 + if loss_type == "grpo": + loss = ((loss_i * mask).sum(-1) / mask.sum(-1).clamp(min=1.0)).mean() + elif loss_type == "bnpo": + loss = (loss_i * mask).sum() / mask.sum().clamp(min=1.0) + elif loss_type == "dr_grpo": + loss = (loss_i * mask).sum() / (loss_i.size(0) * max_completion_length) + else: + raise ValueError(f"Unknown loss type: {loss_type}") + + # loss = (loss_i * mask).sum() / mask.sum() + + # Get metrics as well which are folded + with torch.inference_mode(): + completion_length = n_mask_per_reward.mean() + mean_kl_per_reward = (kl_i * mask).sum(1) / n_mask_per_reward + mean_kl = mean_kl_per_reward.mean() + pass + + return loss, completion_length, mean_kl + +def vLLMSamplingParams(**kwargs): + from vllm import SamplingParams + sampling_params = SamplingParams(**kwargs) + sampling_params._set_kwargs = kwargs + return sampling_params +@dataclass +class UnslothGRPOConfig(GRPOConfig): + """ + + Configuration class for the [`GRPOTrainer`]. + + This class includes only the parameters that are specific to GRPO training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may + differ from those in [`~transformers.TrainingArguments`]. + + 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 (`str`, `dict[str, Any]` or `None`, *optional*, defaults to `None`): + Keyword arguments for [`~transformers.AutoModelForCausalLM.from_pretrained`], used when the `model` + argument of the [`GRPOTrainer`] is provided as a string. + disable_dropout (`bool`, *optional*, defaults to `False`): + Whether to disable dropout in the model. This is useful for training with a reference model, as it prevents + the model from generating different logprobs for the same input. + + > Parameters that control the data preprocessing + + remove_unused_columns (`bool`, *optional*, defaults to `False`): + Whether to only keep the column `"prompt"` in the dataset. If you use a custom reward function that + requires any column other than `"prompts"` and `"completions"`, you should keep this to `False`. + max_prompt_length (`int` or `None`, *optional*, defaults to `512`): + Maximum length of the prompt. If the prompt is longer than this value, it will be truncated left. + num_generations (`int` or `None`, *optional*, defaults to `8`): + Number of generations per prompt to sample. The effective batch size (num_processes * per_device_batch_size + * gradient_accumulation_steps) must be evenly divisible by this value. + max_completion_length (`int` or `None`, *optional*, defaults to `256`): + Maximum length of the generated completion. + ds3_gather_for_generation (`bool`, *optional*, defaults to `True`): + This setting applies to DeepSpeed ZeRO-3. If enabled, the policy model weights are gathered for generation, + improving generation speed. However, disabling this option allows training models that exceed the VRAM + capacity of a single GPU, albeit at the cost of slower generation. Disabling this option is not compatible + with vLLM generation. + shuffle_dataset (`bool`, *optional*, defaults to `True`): + Whether to shuffle the training dataset. + + > Parameters that control generation + + generation_batch_size: (`int` or `None`, *optional*, defaults to `None`): + Batch size to use for generation. If `None`, it defaults to the effective training batch size: + `per_device_train_batch_size * num_processes * gradient_accumulation_steps`. + steps_per_generations: (`int` or `None`, *optional*, defaults to `None`): + Number of optimization steps per generation. If `None`, it defaults to gradient_accumulation_steps. + temperature (`float`, defaults to `1.0`): + Temperature for sampling. The higher the temperature, the more random the completions. + top_p (`float`, *optional*, defaults to `1.0`): + Float that controls the cumulative probability of the top tokens to consider. Must be in (0, 1]. Set to + `1.0` to consider all tokens. + top_k (`int` or `None`, *optional*, defaults to `None`): + Number of highest probability vocabulary tokens to keep for top-k-filtering. If `None`, top-k-filtering is + disabled and all tokens are considered. + min_p (`float` or `None`, *optional*, defaults to `None`): + Minimum token probability, which will be scaled by the probability of the most likely token. It must be a + value between `0.0` and `1.0`. Typical values are in the `0.01-0.2` range. + repetition_penalty (`float`, *optional*, defaults to `1.0`): + Float that penalizes new tokens based on whether they appear in the prompt and the generated text so far. + Values > `1.0` encourage the model to use new tokens, while values < `1.0` encourage the model to repeat + tokens. + cache_implementation (`str` or `None`, *optional*, defaults to `None`): + Implementation of the cache method for faster generation when use_vllm is set to False. + generation_kwargs (`dict[str, Any]` or `None`, *optional*, defaults to `None`): + Additional keyword arguments to pass to `GenerationConfig` (if using transformers) or `SamplingParams` (if + using vLLM) when sampling completions. This can be used to further customize the generation behavior, such + as setting `supress_tokens`, `num_beams`, etc. If it contains keys that conflict with the other generation + parameters (like `min_p`, `top_p`, etc.), they will override them. + + > Parameters that control generation acceleration powered by vLLM + + use_vllm (`bool`, *optional*, defaults to `False`): + Whether to use vLLM for generating completions. If set to `True`, the trainer will use vLLM for generation + instead of the default model.generate(). Requires `vllm` to be installed. + vllm_mode (`str`, *optional*, defaults to `"server"`): + Mode to use for vLLM integration when `use_vllm` is set to `True`. Must be one of `"server"` or + `"colocate"`. + + - `"server"`: The trainer will send generation requests to a separate vLLM server. Make sure a TRL vLLM + server is running (start with `trl vllm-serve`). + - `"colocate"`: vLLM will run in the same process and share the training GPUs. This avoids the need for a + separate server but may cause resource contention with training. + vllm_guided_decoding_regex (`str` or `None`, *optional*, defaults to `None`): + Regex for vLLM guided decoding. If `None` (default), guided decoding is disabled. + + > Parameters that control the vLLM server (only used when `vllm_mode` is `"server"`) + vllm_server_base_url (`str` or `None`, *optional*, defaults to `None`): + Base URL for the vLLM server (e.g., `"http://localhost:8000"`). If provided, `vllm_server_host` and + `vllm_server_port` are ignored. + vllm_server_host (`str`, *optional*, defaults to `"0.0.0.0"`): + Host of the vLLM server to connect to. Ignored if `vllm_server_base_url` is provided. + vllm_server_port (`int`, *optional*, defaults to `8000`): + Port of the vLLM server to connect to. Ignored if `vllm_server_base_url` is provided. + vllm_server_timeout (`float`, *optional*, defaults to `240.0`): + Total timeout duration in seconds to wait for the vLLM server to be up. If the server is not up after the + timeout, a `ConnectionError` is raised. + + > Parameters that control colocated vLLM execution (only used when `vllm_mode` is `"colocate"`) + + vllm_gpu_memory_utilization (`float`, *optional*, defaults to `0.3`): + Control the GPU memory utilization for vLLM. This setting only applies when `vllm_mode` is set to + `"colocate"`. If you are using `vllm_mode="server"`, this parameter must be passed separately when + launching the vLLM server via the `--vllm_gpu_memory_utilization` flag. + vllm_tensor_parallel_size (`int`, *optional*, defaults to `1`): + Control the tensor parallel size for vLLM. This setting only applies when `vllm_mode` is set to + `"colocate"`. If you are using `vllm_mode="server"`, this parameter must be passed separately when + launching the vLLM server via the `--vllm_tensor_parallel_size` flag. + + > Parameters that control the training + + beta (`float`, *optional*, defaults to `0.0`): + KL coefficient. If `0.0` (default), the reference model is not loaded, reducing memory usage and improving + training speed. + num_iterations (`int`, *optional*, defaults to `1`): + Number of iterations per batch (denoted as μ in the algorithm). + epsilon (`float`, *optional*, defaults to `0.2`): + Epsilon value for clipping. + delta: (`float` or `None`, *optional*, defaults to `None`): + Enables the upper clipping bound in two-sided GRPO loss when set to a float. If `None` (default), standard + GRPO clipping is used. Recommended to be greater than `1 + ε` when enabled. This method is introduced in + the [INTELLECT-2 tech report](https://huggingface.co/papers/2505.07291). + epsilon_high (`float` or `None`, *optional*, defaults to `None`): + Upper-bound epsilon value for clipping. If not specified, it defaults to the same value as the lower-bound + specified in argument `epsilon`. Paper [DAPO](https://huggingface.co/papers/2503.14476) recommends `0.28`. + reward_weights (`list[float]` or `None`, *optional*, defaults to `None`): + Weights for each reward function. Must match the number of reward functions. If `None`, all rewards are + weighted equally with weight `1.0`. + scale_rewards (`bool`, *optional*, defaults to `True`): + Whether to scale the rewards by dividing them by their standard deviation. If `True` (default), the rewards + are normalized by the standard deviation, ensuring they have unit variance. If `False`, no scaling is + applied. The [Dr. GRPO paper](https://huggingface.co/papers/2503.20783) recommends not scaling the rewards, + as scaling by the standard deviation introduces a question-level difficulty bias. + loss_type (`str`, *optional*, defaults to `"bnpo"`): + Specifies the loss formulation to use. Supported values are: + + - `"grpo"`: Aggregates token-level losses by normalizing over sequence length. Not recommended due to + length bias—this approach tends to prefer shorter completions with positive advantages and longer ones + with negative advantages. + - `"bnpo"`: Aggregates token-level losses by normalizing number of active token in the local batch. + Note that normalization is performed over the local batch only, so results may slightly vary depending + on the local batch size, despite a constant effective batch size. When using + `per_device_train_batch_size==1`, the loss is equivalent to the GRPO loss. + - `"dr_grpo"`: Aggregates token-level losses by normalizing with a global constant. This method was + introduced in the [Dr. GRPO paper](https://huggingface.co/papers/2503.20783) to eliminate length bias. + The value of the constant corresponds to `max_completion_length`. + mask_truncated_completions (`bool`, *optional*, defaults to `False`): + When enabled, truncated completions are excluded from the loss calculation, preventing them from being + incorrectly penalized and introducing noise during training. According to the + [DAPO](https://huggingface.co/papers/2503.14476) paper, this is a good practice for training stability. + 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 originates 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`. + use_liger_loss (`bool`, *optional*, defaults to `False`): + Whether to use the Liger GRPO loss. + + > Parameters that control the logging + + log_completions (`bool`, *optional*, defaults to `False`): + Whether to log a sample of (prompt, completion) pairs every `logging_steps` steps. If `rich` is installed, + it prints the sample. If `wandb` logging is enabled, it logs it to `wandb`. + num_completions_to_print (`int` or `None`, *optional*, defaults to `None`): + Number of completions to print with `rich`. If `None`, all completions are logged. + wandb_log_unique_prompts (`bool`, *optional*, defaults to `False`): + Whether to log unique prompts in wandb. If `True`, only unique prompts are logged. If `False`, all prompts + are logged. + + """ + 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 = False, + 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, + hub_revision = None, + 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, + liger_kernel_config = None, + eval_use_gather_object = False, + average_tokens_across_devices = False, + model_init_kwargs = None, + disable_dropout = False, + max_prompt_length = 512, + num_generations = 8, + max_completion_length = 256, + ds3_gather_for_generation = True, + shuffle_dataset = True, + generation_batch_size = None, + steps_per_generation = None, + temperature = 1.0, + top_p = 1.0, + top_k = None, + min_p = None, + generation_kwargs = {}, + repetition_penalty = 1.0, + cache_implementation = None, + use_vllm = False, + vllm_server_base_url = None, + vllm_mode = 'colocate', + vllm_guided_decoding_regex = None, + vllm_server_host = '0.0.0.0', + vllm_server_port = 8000, + vllm_server_timeout = 240.0, + vllm_gpu_memory_utilization = 0.3, + vllm_tensor_parallel_size = 1, + beta = 0.001, + num_iterations = 1, + epsilon = 0.2, + delta = None, + epsilon_high = None, + reward_weights = None, + scale_rewards = True, + loss_type = 'bnpo', + mask_truncated_completions = False, + sync_ref_model = False, + ref_model_mixup_alpha = 0.6, + ref_model_sync_steps = 512, + use_liger_loss = False, + log_completions = False, + num_completions_to_print = None, + wandb_log_unique_prompts = 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 loss_type.lower() == 'dr_grpo': + loss_type = 'dr_grpo' + elif loss_type.lower() == 'dapo': + loss_type = 'dapo' + if loss_type.lower() == 'dr_grpo': + if scale_rewards == None: + scale_rewards = True + elif scale_rewards == True: + print('Unsloth: The Dr GRPO paper recommends setting `scale_rewards` to False! Will override. Set it to `None` to force False.') + scale_rewards = False + elif loss_type.lower() == 'dapo': + print('Unsloth: The DAPO paper recommends `mask_truncated_completions = True`') + print('Unsloth: The DAPO paper recommends `epsilon_high = 0.28`') + print('Unsloth: The DAPO paper recommends setting `beta = 0.0` to remove the KL term') + mask_truncated_completions = True + epsilon_high = 0.28 + beta = 0.0 + loss_type = 'bnpo' + + if (per_device_train_batch_size // num_generations) * num_generations != per_device_train_batch_size: + print('Unsloth: We now expect `per_device_train_batch_size` to be a multiple of `num_generations`.\nWe will change the batch size of ' + str(per_device_train_batch_size) + ' to the `num_generations` of ' + str(num_generations)) + per_device_train_batch_size = num_generations + + if temperature <= 0: + raise MathError('Unsloth: Please set a positive non-zero temperature since your results will be wrong.') + elif temperature >= 10: + raise MathError('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, + 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, + 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, + model_init_kwargs = model_init_kwargs, + disable_dropout = disable_dropout, + max_prompt_length = max_prompt_length, + num_generations = num_generations, + max_completion_length = max_completion_length, + ds3_gather_for_generation = ds3_gather_for_generation, + shuffle_dataset = shuffle_dataset, + generation_batch_size = generation_batch_size, + steps_per_generation = steps_per_generation, + temperature = temperature, + top_p = top_p, + top_k = top_k, + min_p = min_p, + generation_kwargs = generation_kwargs, + repetition_penalty = repetition_penalty, + cache_implementation = cache_implementation, + use_vllm = use_vllm, + vllm_server_base_url = vllm_server_base_url, + vllm_mode = vllm_mode, + vllm_guided_decoding_regex = vllm_guided_decoding_regex, + vllm_server_host = vllm_server_host, + vllm_server_port = vllm_server_port, + vllm_server_timeout = vllm_server_timeout, + vllm_gpu_memory_utilization = vllm_gpu_memory_utilization, + vllm_tensor_parallel_size = vllm_tensor_parallel_size, + beta = beta, + num_iterations = num_iterations, + epsilon = epsilon, + delta = delta, + epsilon_high = epsilon_high, + reward_weights = reward_weights, + scale_rewards = scale_rewards, + loss_type = loss_type, + mask_truncated_completions = mask_truncated_completions, + sync_ref_model = sync_ref_model, + ref_model_mixup_alpha = ref_model_mixup_alpha, + ref_model_sync_steps = ref_model_sync_steps, + use_liger_loss = use_liger_loss, + log_completions = log_completions, + num_completions_to_print = num_completions_to_print, + wandb_log_unique_prompts = wandb_log_unique_prompts,**kwargs) + self.vllm_sampling_params = vllm_sampling_params + self.unsloth_num_chunks = unsloth_num_chunks +pass + +class _UnslothGRPOTrainer(Trainer): + """""" + + _tag_names = ["trl", "grpo"] + + def __init__( + self, + model: Union[str, PreTrainedModel], + reward_funcs: Union[RewardFunc, list[RewardFunc]], + args: Optional[GRPOConfig] = None, + train_dataset: Optional[Union[Dataset, IterableDataset]] = None, + eval_dataset: Optional[Union[Dataset, IterableDataset, dict[str, Union[Dataset, IterableDataset]]]] = None, + processing_class: Optional[PreTrainedTokenizerBase] = None, + reward_processing_classes: Optional[Union[PreTrainedTokenizerBase, list[PreTrainedTokenizerBase]]] = None, + callbacks: Optional[list[TrainerCallback]] = None, + optimizers: tuple[Optional[torch.optim.Optimizer], Optional[torch.optim.lr_scheduler.LambdaLR]] = (None, None), + peft_config: Optional["PeftConfig"] = None, + ): + + if hasattr(model, 'vllm_engine') and hasattr(args, 'use_vllm'): + if (getattr(args, 'use_vllm', False) == False): + args.use_vllm = True + args.vllm_mode='colocate' + # Args + if args is None: + model_name = model if isinstance(model, str) else model.config._name_or_path + model_name = model_name.split("/")[-1] + args = GRPOConfig(f"{model_name}-GRPO") + + # Models + # Trained model + model_init_kwargs = args.model_init_kwargs or {} + if isinstance(model, str): + model_id = model + torch_dtype = model_init_kwargs.get("torch_dtype") + if isinstance(torch_dtype, torch.dtype) or torch_dtype == "auto" or torch_dtype is None: + pass # torch_dtype is already a torch.dtype or "auto" or None + elif isinstance(torch_dtype, str): # it's a str, but not "auto" + torch_dtype = getattr(torch, torch_dtype) + model_init_kwargs["torch_dtype"] = torch_dtype + else: + raise ValueError( + "Invalid `torch_dtype` passed to `GRPOConfig`. Expected either 'auto' or a string representing " + f"a `torch.dtype` (e.g., 'float32'), but got {torch_dtype}." + ) + # Disable caching if gradient checkpointing is enabled [not supported] + model_init_kwargs["use_cache"] = ( + False if args.gradient_checkpointing else model_init_kwargs.get("use_cache") + ) + model = AutoModelForCausalLM.from_pretrained(model, **model_init_kwargs) + else: + model_id = model.config._name_or_path + if args.model_init_kwargs is not None: + raise ValueError( + "You passed `model_init_kwargs` to the `GRPOConfig`, but your model is already instantiated. " + "This argument can only be used when the `model` argument is a string." + ) + + if False: + if not is_peft_available(): + raise ImportError("PEFT is required to use `peft_config`. Run `pip install peft`.") + model = model + + # Enable gradient checkpointing if requested + if args.gradient_checkpointing: + model = self._enable_gradient_checkpointing(model, args) + + # Processing class + if processing_class is None: + processing_class = AutoTokenizer.from_pretrained(model.config._name_or_path, padding_side="left") + if processing_class.pad_token is None: + processing_class.pad_token = processing_class.eos_token + + # Reward functions + if not isinstance(reward_funcs, list): + reward_funcs = [reward_funcs] + self.reward_func_names = [] + for i, reward_func in enumerate(reward_funcs): + if isinstance(reward_func, str): + reward_funcs[i] = AutoModelForSequenceClassification.from_pretrained( + reward_func, num_labels=1, **model_init_kwargs + ) + if isinstance(reward_funcs[i], nn.Module): # Use Module over PretrainedModel for compat w/ compiled models + self.reward_func_names.append(reward_funcs[i].config._name_or_path.split("/")[-1]) + else: + self.reward_func_names.append(reward_funcs[i].__name__) + self.reward_funcs = reward_funcs + + # Reward weights + if args.reward_weights is not None: + if len(args.reward_weights) != len(reward_funcs): + raise ValueError( + f"Number of reward weights ({len(args.reward_weights)}) must match number of reward " + f"functions ({len(reward_funcs)})" + ) + self.reward_weights = torch.tensor(args.reward_weights, dtype=torch.float32) + else: + self.reward_weights = torch.ones(len(reward_funcs), dtype=torch.float32) + + # Reward processing class + if reward_processing_classes is None: + reward_processing_classes = [None] * len(reward_funcs) + elif not isinstance(reward_processing_classes, list): + reward_processing_classes = [reward_processing_classes] + else: + if len(reward_processing_classes) != len(reward_funcs): + raise ValueError("The number of reward processing classes must match the number of reward functions.") + + for i, (reward_processing_class, reward_func) in enumerate(zip(reward_processing_classes, reward_funcs)): + if isinstance(reward_func, PreTrainedModel): + if reward_processing_class is None: + reward_processing_class = AutoTokenizer.from_pretrained(reward_func.config._name_or_path) + if reward_processing_class.pad_token_id is None: + reward_processing_class.pad_token = reward_processing_class.eos_token + # The reward model computes the reward for the latest non-padded token in the input sequence. + # So it's important to set the pad token ID to the padding token ID of the processing class. + reward_func.config.pad_token_id = reward_processing_class.pad_token_id + reward_processing_classes[i] = reward_processing_class + self.reward_processing_classes = reward_processing_classes + + # Training arguments + self.max_prompt_length = args.max_prompt_length + self.max_completion_length = args.max_completion_length # = |o_i| in the GRPO paper + self.num_generations = args.num_generations # = G in the GRPO paper + self.temperature = args.temperature + self.top_p = args.top_p + self.top_k = args.top_k + self.min_p = args.min_p + self.repetition_penalty = args.repetition_penalty + self.use_vllm = args.use_vllm + self.vllm_mode = args.vllm_mode + self.vllm_gpu_memory_utilization = args.vllm_gpu_memory_utilization # only applies to colocation mode + self.vllm_tensor_parallel_size = args.vllm_tensor_parallel_size # only applies to colocation mode + self.use_liger_loss = args.use_liger_loss + self.loss_type = args.loss_type + self.scale_rewards = args.scale_rewards + self.mask_truncated_completions = args.mask_truncated_completions + + # Datasets + self.shuffle_dataset = args.shuffle_dataset + + if ( + isinstance(train_dataset, IterableDataset) + or isinstance(eval_dataset, IterableDataset) + or ( + isinstance(eval_dataset, dict) and any(isinstance(ds, IterableDataset) for ds in eval_dataset.values()) + ) + ): + # See https://github.com/huggingface/trl/issues/3213 + raise NotImplementedError( + "Iterable datasets are not yet supported in GRPOTrainer. Please use a standard dataset instead." + ) + + # Multi-step + self.num_iterations = args.num_iterations # = 𝜇 in the GRPO paper + self.epsilon_low = args.epsilon + self.epsilon_high = args.epsilon_high if args.epsilon_high is not None else args.epsilon + # Tracks the number of iterations [forward + backward passes], including those within a grad accum cycle + self._step = 0 + # Buffer the batch to reuse generated outputs across multiple updates. For more details, see + # `_get_train_sampler` and `_prepare_inputs`. + self._buffered_inputs = None + + # 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 GRPO, the sampled data does not include the + # "input_ids" key. Instead, the available keys is "prompt". 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 + + super().__init__( + model=model, + args=args, + data_collator=identity, # No data collation is needed in GRPO + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + callbacks=callbacks, + optimizers=optimizers, + ) + + # Reference model + self.beta = args.beta + if self.beta == 0.0: + # If beta is 0.0, the reference model is not needed + self.ref_model = None + elif is_peft_model(model): + # If PEFT is used, the reference model is not needed since the adapter can be disabled + # to revert to the initial model. + self.ref_model = None + else: + # For deepspeed, fsdp or non-distributed models, create a reference model from scratch + self.ref_model = AutoModelForCausalLM.from_pretrained(model_id, **model_init_kwargs) + + # Disable dropout in the models + if args.disable_dropout: + disable_dropout_in_model(model) + if self.ref_model is not None: + disable_dropout_in_model(self.ref_model) + + # Liger loss + if self.use_liger_loss: + if not is_liger_kernel_available(): + raise ImportError( + "Liger is required to use `liger_loss` as the GRPO loss. Run `pip install liger-kernel`." + ) + # redirect the model.module forward to the model forward to ensure pre-forward hooks are called + self._forward_redirection = _ForwardRedirection() + + self.liger_grpo_loss = LigerFusedLinearGRPOLoss( + beta=self.beta, + epsilon_low=self.epsilon_low, + epsilon_high=self.epsilon_high, + temperature=self.temperature, + use_ref_model=self.beta != 0.0, + loss_type=self.loss_type, + max_completion_length=self.max_completion_length, + ) + + # Initialize the metrics + self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)} + self._total_train_tokens = 0 + self.log_completions = args.log_completions + self.wandb_log_unique_prompts = args.wandb_log_unique_prompts + self.num_completions_to_print = args.num_completions_to_print + # maxlen is set to the total number of forward passes per step. This value of `maxlen` ensures we log only the + # final optimization step. + maxlen = self.accelerator.num_processes * args.per_device_train_batch_size * args.steps_per_generation + self._textual_logs = { + "prompt": deque(maxlen=maxlen), + "completion": deque(maxlen=maxlen), + "rewards": defaultdict(lambda: deque(maxlen=maxlen)), + "advantages": deque(maxlen=maxlen), + } + + # Ensure each process receives a unique seed to prevent duplicate completions when generating with + # transformers if num_generations exceeds per_device_train_batch_size. We could skip it if we use vLLM, but + # it's safer to set it in all cases. + set_seed(args.seed, device_specific=True) + + if self.use_vllm: + if not is_vllm_available(): + raise ImportError( + "vLLM is not available and `use_vllm` is set to True. Please install vLLM with " + "`pip install vllm` to use it." + ) + + if self.vllm_mode == "server" and self.accelerator.is_main_process: + if args.vllm_server_base_url is not None: + base_url = args.vllm_server_base_url + else: + base_url = f"http://{args.vllm_server_host}:{args.vllm_server_port}" + self.vllm_client = VLLMClient(base_url=base_url, connection_timeout=args.vllm_server_timeout) + self.vllm_client.init_communicator() + + elif self.vllm_mode == "colocate": + if not self.accelerator.num_processes % self.vllm_tensor_parallel_size == 0: + raise ValueError( + f"vllm_tensor_parallel_size ({self.vllm_tensor_parallel_size}) must divide world size " + f"({self.accelerator.num_processes}) evenly." + ) + + if self.vllm_tensor_parallel_size > 1: + self.tp_group, _ = torch.distributed.new_subgroups_by_enumeration( + [ + list(range(i * self.vllm_tensor_parallel_size, (i + 1) * self.vllm_tensor_parallel_size)) + for i in range(self.accelerator.num_processes // self.vllm_tensor_parallel_size) + ] + ) + + self.llm = model.vllm_engine + self.guided_decoding_regex = args.vllm_guided_decoding_regex + + self._last_loaded_step = -1 + self.accelerator.wait_for_everyone() + else: + generation_kwargs = { + "max_new_tokens": self.max_completion_length, + "do_sample": True, + "pad_token_id": processing_class.pad_token_id, + "bos_token_id": processing_class.bos_token_id, + "eos_token_id": processing_class.eos_token_id, + "temperature": self.temperature, + "top_p": self.top_p, + "top_k": self.top_k, + "min_p": self.min_p, + "repetition_penalty": self.repetition_penalty, + "cache_implementation": args.cache_implementation, + } + if args.generation_kwargs is not None: + generation_kwargs.update(args.generation_kwargs) + self.generation_config = GenerationConfig(**generation_kwargs) + + # 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 to the model + self.model.add_model_tags(self._tag_names) + + if self.ref_model is not None: + 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: + self.add_callback(SyncRefModelCallback(ref_model=self.ref_model, accelerator=self.accelerator)) + + for i, reward_func in enumerate(self.reward_funcs): + if isinstance(reward_func, PreTrainedModel): + if self.is_deepspeed_enabled: + self.reward_funcs[i] = prepare_deepspeed(reward_func, self.accelerator) + else: + # set device placement to True to make `prepare_model` move `reward_func` to device when using fsdp + self.reward_funcs[i] = self.accelerator.prepare_model( + reward_func, evaluation_mode=True, device_placement=True + ) + + 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 GRPOTrainer, we preprocess data, so using the model's signature columns doesn't work. + # Instead, we set them to the columns expected by the `training_step` method, hence the override. + if self._signature_columns is None: + self._signature_columns = ["prompt"] + + # This method overrides `Trainer.get_train_dataloader` to support our custom batching strategy. + # Instead of returning a standard per-step batch (i.e., `per_device_batch_size), our dataloader loads an + # *generation* batch (i.e., `per_device_batch_size × steps_per_generation`). This allows us to generate completions + # once every steps_per_generation step—rather than once per accumulation step—which is significantly more + # efficient. The only change from the original implementation is multiplying the batch size by + # `steps_per_generation`. Thus, `_prepare_inputs` is called with this *generation* batch, and it handles the + # splitting internally. + # Maintenance note: This method is a copy-paste of the original `Trainer.get_train_dataloader` with only one line + # modification. As a result, some parts of the method aren't relevant to GRPO, but we keep them to stay one line + # apart from the super method, ensuring easier maintenance in the future. + def get_train_dataloader(self): + if self.train_dataset is None: + raise ValueError("Trainer: training requires a train_dataset.") + + train_dataset = self.train_dataset + data_collator = self.data_collator + if is_datasets_available() and isinstance(train_dataset, datasets.Dataset): + train_dataset = self._remove_unused_columns(train_dataset, description="training") + else: + data_collator = self._get_collator_with_removed_columns(data_collator, description="training") + + dataloader_params = { + "batch_size": self._train_batch_size * self.args.steps_per_generation, # < this is the change + "collate_fn": data_collator, + "num_workers": self.args.dataloader_num_workers, + "pin_memory": self.args.dataloader_pin_memory, + "persistent_workers": self.args.dataloader_persistent_workers, + } + + if not isinstance(train_dataset, torch.utils.data.IterableDataset): + dataloader_params["sampler"] = self._get_train_sampler() + dataloader_params["drop_last"] = self.args.dataloader_drop_last + if version.parse(transformers.__version__) >= version.parse("4.52.0"): + # from transformers 4.52.0, the `seed_worker` requires the `num_workers` and `rank` arguments + dataloader_params["worker_init_fn"] = partial( + seed_worker, num_workers=self.args.dataloader_num_workers, rank=self.args.process_index + ) + else: + dataloader_params["worker_init_fn"] = seed_worker + dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor + + return self.accelerator.prepare(DataLoader(train_dataset, **dataloader_params)) + + def _get_train_sampler(self, dataset: Optional[Dataset] = None) -> Sampler: + # Returns a sampler that + # 1. ensures each prompt is repeated across multiple processes. This guarantees that identical prompts are + # distributed to different GPUs, allowing rewards to be computed and normalized correctly within each prompt + # group. Using the same seed across processes ensures consistent prompt assignment, preventing discrepancies + # in group formation. + # 2. repeats the batch multiple times to allow reusing generations across multiple updates. Refer to + # _prepare_inputs to see how the generations are stored and reused. + + # In the following figure, the values are the prompt indices. The first row shows the first sampled batch, the + # second row shows the second sampled batch, and so on. + # + # | GPU 0 | GPU 1 | + # + # global_step step <-───> num_generations=2 + # <-───────> per_device_train_batch_size=3 + # grad_accum ▲ ▲ 0 0 0 0 1 1 2 2 <- Generate for the first `steps_per_generation` (prompts 0 to 11); store the completions; use the first slice to compute the loss + # =2 ▼ | 0 1 3 3 4 4 5 5 <- Take the stored generations and use the second slice to compute the loss + # | + # | 1 2 6 6 7 7 8 8 <- Take the stored generations and use the third slice to compute the loss + # steps_per_gen=4 ▼ 1 3 9 9 10 10 11 11 <- Take the stored generations and use the fourth slice to compute the loss + # + # 2 4 12 12 13 13 14 14 <- Generate for the second `steps_per_generation` (prompts 12 to 23); store the completions; use the first slice to compute the loss + # 2 5 15 15 16 16 17 17 <- Take the stored generations and use the second slice to compute the loss + # ... + if dataset is None: + dataset = self.train_dataset + return RepeatSampler( + data_source=dataset, + mini_repeat_count=self.num_generations, + batch_size=self.args.generation_batch_size // self.num_generations, + repeat_count=self.num_iterations * self.args.steps_per_generation, + shuffle=self.shuffle_dataset, + seed=self.args.seed, + ) + + def _get_eval_sampler(self, eval_dataset) -> Sampler: + # See _get_train_sampler for an explanation of the sampler. + return RepeatSampler( + data_source=eval_dataset, + mini_repeat_count=self.num_generations, + seed=self.args.seed, + ) + + def _enable_gradient_checkpointing(self, model: PreTrainedModel, args: GRPOConfig) -> PreTrainedModel: + """Enables gradient checkpointing for the model.""" + # Ensure use_cache is disabled + model.config.use_cache = False + + # Enable gradient checkpointing on the base model for PEFT + if is_peft_model(model): + model.base_model.gradient_checkpointing_enable() + # Enable gradient checkpointing for non-PEFT models + else: + model.gradient_checkpointing_enable() + + gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs or {} + use_reentrant = ( + "use_reentrant" not in gradient_checkpointing_kwargs or gradient_checkpointing_kwargs["use_reentrant"] + ) + + if use_reentrant: + model.enable_input_require_grads() + + return model + + @profiling_decorator + def _get_last_hidden_state(self, unwrapped_model, input_ids, attention_mask, logits_to_keep=None): + if is_peft_model(unwrapped_model): + unwrapped_model = unwrapped_model.base_model.model + last_hidden_state = unwrapped_model.model(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state + last_hidden_state = last_hidden_state[:, :-1, :] # (B, L-1, H) + if logits_to_keep is not None: + last_hidden_state = last_hidden_state[:, -logits_to_keep:, :] # (B, logits_to_keep, H) + return last_hidden_state + + # Get the per-token log probabilities for the completions for the model and the reference model + def _get_per_token_logps(self, model, input_ids, attention_mask, logits_to_keep): + if True: # os.environ.get('UNSLOTH_USE_NEW_MODEL', '0') == '0': + return None # Unsloth efficient GRPO + # Otherwise, calculate normally: + if not hasattr(self, '_autocast_dtype'): + self._autocast_dtype = torch.float16 if os.environ.get('ACCELERATE_MIXED_PRECISION', 'fp16') == 'fp16' else torch.bfloat16 + if os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1': self._autocast_dtype = torch.float16 + + os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" + with torch.amp.autocast(device_type = 'cuda', dtype = self._autocast_dtype): + # We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded + logits = model( + input_ids = input_ids, + attention_mask = attention_mask, + logits_to_keep = logits_to_keep + 1, + ).logits + # logits = logits[:, :-1, :] # (B, L-1, V), exclude the last logit: it corresponds to the next token pred + return logits + # input_ids = input_ids[:, -logits_to_keep:] + # For transformers<=4.48, logits_to_keep argument isn't supported, so here we drop logits ourselves. + # See https://github.com/huggingface/trl/issues/2770 + # logits = logits[:, -logits_to_keep:] + # return logits + # See https://huggingface.co/blog/the_n_implementation_details_of_rlhf_with_ppo#policy-training-implementation-details + # logits = logits / self.temperature + # logps = selective_log_softmax(logits, input_ids) + + # row_indices, col_indices = torch.where(logps < -20) + + # # Method 1: Check if tensors have elements + # if len(row_indices) > 0 and len(col_indices) > 0: + # breakpoint() # Breakpoint triggered here + # print("Found high values!") + # return logps # compute logprobs for the input tokens + pass + + def _sync_fsdp_params_to_vllm(self, module: nn.Module, prefix: str = "", visited=None): + """Memory-efficient post-order traversal of FSDP modules to extract full parameters and sync with vLLM.""" + if visited is None: + visited = set() + + for child_name, child_module in module.named_children(): + child_prefix = f"{prefix}.{child_name}" if prefix else child_name + self._sync_fsdp_params_to_vllm( + child_module, prefix=child_prefix, visited=visited + ) # recurse into the child + + if isinstance(module, FSDP): + with FSDP.summon_full_params(module, recurse=False, writeback=False): + for param_name, param in module.named_parameters(): + full_name = f"{prefix}.{param_name}" if prefix else param_name + for extra in ("_fsdp_wrapped_module.", "_checkpoint_wrapped_module."): + full_name = full_name.replace(extra, "") + + if full_name in visited: + continue # skip FSDP subtrees already traversed + visited.add(full_name) + + if self.vllm_mode == "server" and self.accelerator.is_main_process: + self.vllm_client.update_named_param(full_name, param.data) + elif self.vllm_mode == "colocate": + + pass + + pass + + def _move_model_to_vllm(self, *args, **kwargs): return None + + @profiling_decorator + def _prepare_inputs( + self, generation_batch: dict[str, Union[torch.Tensor, Any]] + ) -> dict[str, Union[torch.Tensor, Any]]: + # Prepares inputs for model training/evaluation by managing completion generation and batch handling. + # During training: + # - Receives the local generation batch (Per-GPU batch size × steps per generation) + # from the modified training dataloader instead of the standard local batch + # - Generates completions once for the entire generation batch and splits it into batches of size + # `per_device_train_batch_size` + # - Buffers these completions and returns the appropriate slice for the current accumulation step + # - Optimizes by regenerating completions only periodically (every steps_per_generation * num_iterations) + # During evaluation: + # - The input is treated as a standard local batch (no accumulation, no multiple iterations) + # - Completions are generated for each batch without buffering or reuse + # Returns a single local batch in both cases. + if hasattr(self, 'llm'): + if getattr(self.llm.llm_engine.vllm_config.model_config, 'enable_sleep_mode', False): + self.llm.wake_up() + + mode = "train" if self.model.training else "eval" + if mode == "train": + generate_every = self.args.steps_per_generation * self.num_iterations + if self._step % generate_every == 0 or self._buffered_inputs is None: + # self._buffered_inputs=None can occur when resuming from a checkpoint + generation_batch = self._generate_and_score_completions(generation_batch) + generation_batch = shuffle_tensor_dict(generation_batch) + self._buffered_inputs = split_tensor_dict(generation_batch, self.args.steps_per_generation) + inputs = self._buffered_inputs[self._step % self.args.steps_per_generation] + self._step += 1 + else: + # In evaluation, there is neither batch grouping for generation, nor multiple iterations, hence + # local generation batch == local eval batch + inputs = self._generate_and_score_completions(generation_batch) + if hasattr(self, 'llm'): + if getattr(self.llm.llm_engine.vllm_config.model_config, 'enable_sleep_mode', False): + self.llm.sleep(os.environ.get('VLLM_SLEEP_MODE', 1)) + return inputs + + @profiling_decorator + def _calculate_rewards(self, inputs, prompts, completions, completion_ids_list): + device = self.accelerator.device + rewards_per_func = torch.zeros(len(prompts), len(self.reward_funcs), device=device) + + # Repeat all input columns (but "prompt", "completion", and "completion_ids") to match the num of generations + keys = [key for key in inputs[0] if key not in ["prompt", "completion", "completion_ids"]] + reward_kwargs = {key: [example[key] for example in inputs] for key in keys} + + for i, (reward_func, reward_processing_class, reward_func_name) in enumerate( + zip(self.reward_funcs, self.reward_processing_classes, self.reward_func_names) + ): + with profiling_context(self, reward_func_name): + if isinstance(reward_func, nn.Module): # Module (no PretrainedModel) for compat with compiled models + if is_conversational(inputs[0]): + messages = [{"messages": p + c} for p, c in zip(prompts, completions)] + texts = [apply_chat_template(x, reward_processing_class)["text"] for x in messages] + else: + texts = [p + c for p, c in zip(prompts, completions)] + reward_inputs = reward_processing_class( + text=texts, return_tensors="pt", padding=True, padding_side="right", add_special_tokens=False + ) + reward_inputs = super()._prepare_inputs(reward_inputs) + with torch.inference_mode(): + rewards_per_func[:, i] = reward_func(**reward_inputs).logits[:, 0] # Shape (B*G,) + else: + output_reward_func = reward_func( + prompts=prompts, completions=completions, completion_ids=completion_ids_list, **reward_kwargs + ) + # Convert None values to NaN + output_reward_func = [reward if reward is not None else torch.nan for reward in output_reward_func] + + rewards_per_func[:, i] = torch.tensor(output_reward_func, dtype=torch.float32, device=device) + + # If all reward functions return None for a given row, issue a detailed warning + if torch.isnan(rewards_per_func).all(dim=1).any(): + nan_row_idx = torch.isnan(rewards_per_func).all(dim=1).nonzero(as_tuple=True)[0][0] + row_reward_kwargs = {key: value[nan_row_idx] for key, value in reward_kwargs.items()} + row_reward_kwargs["prompt"] = prompts[nan_row_idx] + row_reward_kwargs["completion"] = completions[nan_row_idx] + warnings.warn( + f"All reward functions returned None for the following kwargs: {row_reward_kwargs}. " + "Please ensure that at least one reward function returns a valid reward." + ) + + # Gather the reward per function: this part is crucial, because the rewards are normalized per group and the + # completions may be distributed across processes + rewards_per_func = gather(rewards_per_func) + return rewards_per_func + + def _generate_and_score_completions( + self, inputs: list[dict[str, Union[torch.Tensor, Any]]] + ) -> dict[str, Union[torch.Tensor, Any]]: + device = self.accelerator.device + mode = "train" if self.model.training else "eval" + + prompts = [x["prompt"] for x in inputs] + prompts_text = [maybe_apply_chat_template(example, self.processing_class)["prompt"] for example in inputs] + prompt_inputs = self.processing_class( + text=prompts_text, return_tensors="pt", padding=True, padding_side="left", add_special_tokens=False + ) + prompt_inputs = super()._prepare_inputs(prompt_inputs) + prompt_ids, prompt_mask = prompt_inputs["input_ids"], prompt_inputs["attention_mask"] + + if self.max_prompt_length is not None: + # If max_prompt_length is set, we trim the prompt to keep only the last `max_prompt_length` tokens. + # Then we decode those tokens back into text. We manually remove leading pad tokens from the decoded text, + # because we can't use `skip_special_tokens=True` (some special tokens are still needed for generation). + prompt_ids = prompt_ids[:, -self.max_prompt_length :] + prompt_mask = prompt_mask[:, -self.max_prompt_length :] + prompts_text = self.processing_class.batch_decode( + prompt_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False + ) + prompts_text = [ + re.sub(rf"^({re.escape(self.processing_class.pad_token)})+", "", text) for text in prompts_text + ] + + # Generate completions using either vLLM or regular generation + if self.use_vllm: + # First, update the vLLM weights if needed + if self.state.global_step != self._last_loaded_step: + self._move_model_to_vllm() + self._last_loaded_step = self.state.global_step + + # Generate completions using vLLM: gather all prompts and use them in a single call in the main process + if self.vllm_mode == "server": + all_prompts_text = gather_object(prompts_text) + if self.accelerator.is_main_process: + # Since 'prompts' contains 'num_generations' duplicates, we first take unique prompts, and generate + # num_generations outputs for each one. This is faster than generating outputs for each duplicate + # prompt individually. + ordered_set_of_prompts = all_prompts_text[:: self.num_generations] + with profiling_context(self, "vLLM.generate"): + completion_ids = self.vllm_client.generate( + prompts=ordered_set_of_prompts, + n=self.num_generations, + repetition_penalty=self.repetition_penalty, + temperature=self.temperature, + top_p=self.top_p, + top_k=-1 if self.top_k is None else self.top_k, + min_p=0.0 if self.min_p is None else self.min_p, + max_tokens=self.max_completion_length, + guided_decoding_regex=self.guided_decoding_regex, + generation_kwargs=self.args.generation_kwargs, + ) + else: + completion_ids = [None] * len(all_prompts_text) + # Broadcast the completions from the main process to all processes, ensuring each process receives its + # corresponding slice. + completion_ids = broadcast_object_list(completion_ids, from_process=0) + process_slice = slice( + self.accelerator.process_index * len(prompts), + (self.accelerator.process_index + 1) * len(prompts), + ) + completion_ids = completion_ids[process_slice] + + # Generate completions using colocated vLLM instances: each device holds vLLM copy and work on their own batch of prompts + elif self.vllm_mode == "colocate": + if self.guided_decoding_regex: + guided_decoding = GuidedDecodingParams(backend="outlines", regex=self.guided_decoding_regex) + else: + guided_decoding = None + + generation_kwargs = { + "n": 1, # vLLM on each GPU generates only 1 in colocate mode + "repetition_penalty": self.repetition_penalty, + "temperature": self.temperature, + "top_p": self.top_p, + "top_k": -1 if self.top_k is None else self.top_k, + "min_p": 0.0 if self.min_p is None else self.min_p, + "max_tokens": self.max_completion_length, + "guided_decoding": guided_decoding, + } + if self.args.generation_kwargs is not None: + generation_kwargs.update(self.args.generation_kwargs) + sampling_params = SamplingParams(**generation_kwargs) + + if self.vllm_tensor_parallel_size > 1: + # Gather prompts from all ranks in the TP group and flatten. + # Each rank starts with its own prompts; after gathering, all ranks see the full group set. + orig_size = len(prompts_text) + gathered_prompts = [None for _ in range(self.vllm_tensor_parallel_size)] + torch.distributed.all_gather_object(gathered_prompts, prompts_text, group=self.tp_group) + all_prompts_text = [p for sublist in gathered_prompts for p in sublist] + else: + all_prompts_text = prompts_text + + with profiling_context(self, "vLLM.generate"): + all_outputs = self.llm.generate(all_prompts_text, sampling_params=sampling_params, use_tqdm=False, lora_request = self.model.load_lora('grpo_trainer_lora_model', load_tensors = True)) + + completion_ids = [output.token_ids for outputs in all_outputs for output in outputs.outputs] + + if self.vllm_tensor_parallel_size > 1: + # Slice completions for this rank within its TP group. + # Each rank generates all outputs — we keep only our share. + local_rank_in_group = torch.distributed.get_rank(group=self.tp_group) + tp_slice = slice(local_rank_in_group * orig_size, (local_rank_in_group + 1) * orig_size) + completion_ids = completion_ids[tp_slice] + + # Pad the completions, and concatenate them with the prompts + completion_ids = [torch.tensor(ids, device=device) for ids in completion_ids] + completion_ids = pad(completion_ids, padding_value=self.processing_class.pad_token_id) + prompt_completion_ids = torch.cat([prompt_ids, completion_ids], dim=1) + else: + # Regular generation path + with unwrap_model_for_generation( + self.model_wrapped, self.accelerator, gather_deepspeed3_params=self.args.ds3_gather_for_generation + ) as unwrapped_model: + with ( + FSDP.summon_full_params(self.model_wrapped, recurse=False) + if self.is_fsdp_enabled + else nullcontext() + ): + prompt_completion_ids = unwrapped_model.generate( + prompt_ids, attention_mask=prompt_mask, generation_config=self.generation_config + ) + + # Compute prompt length and extract completion ids + prompt_length = prompt_ids.size(1) + prompt_ids = prompt_completion_ids[:, :prompt_length] + completion_ids = prompt_completion_ids[:, prompt_length:] + + # Mask everything after the first EOS token + is_eos = completion_ids == self.processing_class.eos_token_id + eos_idx = torch.full((is_eos.size(0),), is_eos.size(1), dtype=torch.long, device=device) + eos_idx[is_eos.any(dim=1)] = is_eos.int().argmax(dim=1)[is_eos.any(dim=1)] + sequence_indices = torch.arange(is_eos.size(1), device=device).expand(is_eos.size(0), -1) + completion_mask = (sequence_indices <= eos_idx.unsqueeze(1)).int() + + # Convert tensor to a list of lists of token IDs. This will be passed to the reward function, avoiding the need + # to re-tokenize completions if the reward is computed from tokens. + completion_ids_list = [ + [id.item() for id, m in zip(row, mask_row) if m] for row, mask_row in zip(completion_ids, completion_mask) + ] + + # Sum along sequence dimension (dim=1) to get completion length per sequence, used for logging + completion_lengths = completion_mask.sum(1) + + # If mask_truncated_completions is enabled, zero out truncated completions in completion_mask + if self.mask_truncated_completions: + truncated_completions = ~is_eos.any(dim=1) + completion_mask = completion_mask * (~truncated_completions).unsqueeze(1).int() + + # Concatenate prompt_mask with completion_mask for logit computation + attention_mask = torch.cat([prompt_mask, completion_mask], dim=1) # (B, P+C) + + logits_to_keep = completion_ids.size(1) # we only need to compute the logits for the completion tokens + batch_size = self.args.per_device_train_batch_size if mode == "train" else self.args.per_device_eval_batch_size + + with torch.no_grad(): + # When using num_iterations == 1 and steps_per_generation <= gradient_accumulation_steps + # old_per_token_logps == per_token_logps, so we can skip it's computation here, and use + # per_token_logps.detach() instead. + if self.num_iterations > 1 or self.args.steps_per_generation > self.args.gradient_accumulation_steps: + old_per_token_logps = self._get_per_token_logps( + self.model, prompt_completion_ids, attention_mask, logits_to_keep, batch_size + ) + else: + old_per_token_logps = None + + # Compute the per-token log probabilities for the reference model + if self.beta != 0.0: + if self.ref_model is not None: + ref_per_token_logps = self._get_per_token_logps( + self.ref_model, prompt_completion_ids, attention_mask, logits_to_keep + ) + else: + with self.accelerator.unwrap_model(self.model).disable_adapter(): + ref_per_token_logps = self._get_per_token_logps( + self.model, prompt_completion_ids, attention_mask, logits_to_keep + ) + else: + ref_per_token_logps = None + + # Decode the generated completions + completions_text = self.processing_class.batch_decode(completion_ids, skip_special_tokens=True) + if is_conversational(inputs[0]): + completions = [] + for prompt, completion in zip(prompts, completions_text): + bootstrap = prompt.pop()["content"] if prompt[-1]["role"] == "assistant" else "" + completions.append([{"role": "assistant", "content": bootstrap + completion}]) + else: + completions = completions_text + + # Calculate rewards for each reward function. rewards_per_func aggregates rewards across all processes. This is + # important because rewards will be normalized per group, and completions are distributed. We will later slice + # rewards_per_func to extract each process's subset. + rewards_per_func = self._calculate_rewards(inputs, prompts, completions, completion_ids_list) + + # Apply weights to each reward function's output and sum + rewards = (rewards_per_func * self.reward_weights.to(device).unsqueeze(0)).nansum(dim=1) + + # Compute grouped-wise rewards + mean_grouped_rewards = rewards.view(-1, self.num_generations).mean(dim=1) + std_grouped_rewards = rewards.view(-1, self.num_generations).std(dim=1) + is_std_zero = torch.isclose(std_grouped_rewards, torch.zeros_like(std_grouped_rewards)) + + # Normalize the rewards to compute the advantages + mean_grouped_rewards = mean_grouped_rewards.repeat_interleave(self.num_generations, dim=0) + std_grouped_rewards = std_grouped_rewards.repeat_interleave(self.num_generations, dim=0) + advantages = rewards - mean_grouped_rewards + if self.scale_rewards: + advantages = advantages / (std_grouped_rewards + 1e-4) + + # Slice to keep only the local part of the data + process_slice = slice( + self.accelerator.process_index * len(prompts), + (self.accelerator.process_index + 1) * len(prompts), + ) + all_process_advantages = advantages.clone() # keep the aggregated advantages for logging + advantages = advantages[process_slice] + + # Log the metrics + if mode == "train": + self.state.num_input_tokens_seen += self.accelerator.gather(attention_mask.sum()).sum().item() + self._metrics[mode]["num_tokens"] = [self.state.num_input_tokens_seen] + + # Log completion lengths, mean, min, max + agg_completion_lengths = self.accelerator.gather(completion_lengths) + self._metrics[mode]["completions/mean_length"].append(agg_completion_lengths.float().mean().item()) + self._metrics[mode]["completions/min_length"].append(agg_completion_lengths.float().min().item()) + self._metrics[mode]["completions/max_length"].append(agg_completion_lengths.float().max().item()) + + # Identify sequences that terminated with EOS and log their lengths + agg_terminated_with_eos = self.accelerator.gather(is_eos.any(dim=1)) + term_completion_lengths = agg_completion_lengths[agg_terminated_with_eos] + clipped_completions_ratio = 1 - len(term_completion_lengths) / len(agg_completion_lengths) + self._metrics[mode]["completions/clipped_ratio"].append(clipped_completions_ratio) + if len(term_completion_lengths) == 0: # edge case where no terminated sequences are found + term_completion_lengths = torch.zeros(1, device=device) + self._metrics[mode]["completions/mean_terminated_length"].append(term_completion_lengths.float().mean().item()) + self._metrics[mode]["completions/min_terminated_length"].append(term_completion_lengths.float().min().item()) + self._metrics[mode]["completions/max_terminated_length"].append(term_completion_lengths.float().max().item()) + + # Calculate mean reward per function, but only for samples where the function was applied (non-NaN values) + for i, reward_func_name in enumerate(self.reward_func_names): + mean_rewards = torch.nanmean(rewards_per_func[:, i]).item() + self._metrics[mode][f"rewards/{reward_func_name}/mean"].append(mean_rewards) + std_rewards = nanstd(rewards_per_func[:, i]).item() + self._metrics[mode][f"rewards/{reward_func_name}/std"].append(std_rewards) + self._metrics[mode]["reward"].append(mean_grouped_rewards.mean().item()) + self._metrics[mode]["reward_std"].append(std_grouped_rewards.mean().item()) + self._metrics[mode]["frac_reward_zero_std"].append(is_std_zero.float().mean().item()) + + # Log prompt and completion texts + self._textual_logs["prompt"].extend(gather_object(prompts_text)) + self._textual_logs["completion"].extend(gather_object(completions_text)) + for i, name in enumerate(self.reward_func_names): + self._textual_logs["rewards"][name].extend(rewards_per_func[:, i].tolist()) + self._textual_logs["advantages"].extend(all_process_advantages.tolist()) + + return { + "prompt_ids": prompt_ids, + "prompt_mask": prompt_mask, + "completion_ids": completion_ids, + "completion_mask": completion_mask, + "advantages": advantages, + "old_per_token_logps": old_per_token_logps, + "ref_per_token_logps": ref_per_token_logps, + } + + def compute_liger_loss(self, unwrapped_model, inputs): + # Compute the per-token log probabilities for the model + prompt_ids, prompt_mask = inputs["prompt_ids"], inputs["prompt_mask"] + completion_ids, completion_mask = inputs["completion_ids"], inputs["completion_mask"] + input_ids = torch.cat([prompt_ids, completion_ids], dim=1) + attention_mask = torch.cat([prompt_mask, completion_mask], dim=1) + logits_to_keep = completion_ids.size(1) # we only need to compute the logits for the completion tokens + + # get the last hidden state of the model + last_hidden_state = self._get_last_hidden_state(unwrapped_model, input_ids, attention_mask, logits_to_keep) + + # compute loss and metrics using liger grpo loss + loss, metrics = self.liger_grpo_loss( + _input=last_hidden_state, + lin_weight=unwrapped_model.lm_head.weight, + selected_token_ids=completion_ids, + attention_mask=completion_mask, + advantages=inputs["advantages"], + bias=unwrapped_model.lm_head.bias, + old_per_token_logps=inputs["old_per_token_logps"], + ref_per_token_logps=inputs["ref_per_token_logps"], + ) + # Extract metrics from the liger_grpo_loss output + # KL divergence is the first metric when beta is non-zero + mean_kl = metrics[0] if self.beta != 0.0 else None + clip_ratio = metrics[-1] + + mode = "train" if self.model.training else "eval" + if self.beta != 0.0: + self._metrics[mode]["kl"].append(self.accelerator.gather(mean_kl).mean().item()) + self._metrics[mode]["clip_ratio"].append(self.accelerator.gather(clip_ratio).mean().item()) + return loss + + def compute_loss(self, model, inputs, return_outputs = False, num_items_in_batch = None): + if return_outputs: + raise ValueError("The GRPOTrainer does not support returning outputs") + # Compute the per-token log probabilities for the model + + prompt_ids, prompt_mask = inputs["prompt_ids"], inputs["prompt_mask"] + completion_ids, completion_mask = inputs["completion_ids"], inputs["completion_mask"] + input_ids = torch.cat([prompt_ids, completion_ids], dim=1) + bsz, qlen = input_ids.shape + attention_mask = torch.cat([prompt_mask, completion_mask], dim=1) + # attention_mask = None + logits_to_keep = completion_ids.size(1) # we only need to compute the logits for the completion tokens + _input_ids = input_ids + _logits_to_keep = logits_to_keep + + per_token_logps = self._get_per_token_logps(model, input_ids, attention_mask, logits_to_keep) + + # Compute the KL divergence between the model and the reference model + # _prepare_inputs doesn't return reference log probs anymore. We need to calculate it ourselves. + # https://github.com/huggingface/trl/blob/05bc43e960396581e458195b8388efe6b82cae1f/trl/trainer/grpo_trainer.py#L1328 + if self.beta != 0.0: + with torch.inference_mode(), model.disable_adapter(): + ref_per_token_logps = self._get_per_token_logps(model, input_ids, attention_mask, logits_to_keep) + else: + ref_per_token_logps = None + # per_token_kl = torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1 + # x - x.detach() allows for preserving gradients from x + advantages = inputs["advantages"] + # per_token_loss = torch.exp(per_token_logps - per_token_logps.detach()) * advantages.unsqueeze(1) + # per_token_loss = -(per_token_loss - self.beta * per_token_kl) + # loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean() + old_hidden_states = inputs.get("old_per_token_logps", None) + input_ids = input_ids[:, -logits_to_keep:] + + # Get logit softcapping and logit scale + logit_softcapping = getattr(model.config, "final_logit_softcapping", 0) # Gemma + if logit_softcapping is None: logit_softcapping = 0 + logit_scale_multiply = getattr(model.config, "logit_scale", 0) # Cohere + if logit_scale_multiply is None: logit_scale_multiply = 0 + logit_scale_divide = getattr(model.config, "logits_scaling", 0) # Granite + if logit_scale_divide is None: logit_scale_divide = 0 + if per_token_logps is not None: + + if ref_per_token_logps is not None: + ref_per_token_logps = ref_per_token_logps[:, :-1, :] # (B, L-1, V), exclude the last logit: it corresponds to the next token pred + per_token_logps = per_token_logps[:, :-1, :] # (B, L-1, V), exclude the last logit: it corresponds to the next token pred + + loss, completion_length, mean_kl = grpo_compute_loss_slow( + ref_per_token_logps, + per_token_logps, + old_hidden_states, + input_ids, + completion_mask, + self.beta, + advantages, + loss_type = self.args.loss_type, + epsilon_low = self.epsilon_low, + epsilon_high = self.epsilon_high, + max_completion_length = self.args.max_completion_length, + delta = self.args.delta, + temperature = self.args.temperature, + logit_softcapping = logit_softcapping, + logit_scale_multiply = logit_scale_multiply, + logit_scale_divide = logit_scale_divide, + ) + else: + if hasattr(self.args, "loss_type"): + loss, completion_length, mean_kl = grpo_accumulated_loss( + trainer = self, + input_ids = _input_ids, + logits_to_keep = logits_to_keep, + completion_mask = completion_mask, + advantages = advantages, + old_hidden_states = old_hidden_states, + n_chunks = self.args.unsloth_num_chunks, + loss_type = self.args.loss_type, + epsilon_low = self.epsilon_low, + epsilon_high = self.epsilon_high, + max_completion_length = self.args.max_completion_length, + delta = self.args.delta, + temperature = self.args.temperature, + logit_softcapping = logit_softcapping, + logit_scale_multiply = logit_scale_multiply, + logit_scale_divide = logit_scale_divide, + attention_mask = attention_mask, + ) + else: + # to ensure backwards compatibility with trl 0.15.2 and maybe even 0.17 + loss, completion_length, mean_kl = grpo_accumulated_loss( + trainer = self, + input_ids = _input_ids, + logits_to_keep = logits_to_keep, + completion_mask = completion_mask, + advantages = advantages, + old_hidden_states = old_hidden_states, + n_chunks = self.args.unsloth_num_chunks, + temperature = self.args.temperature, + logit_softcapping = logit_softcapping, + logit_scale_multiply = logit_scale_multiply, + logit_scale_divide = logit_scale_divide, + attention_mask = attention_mask, + ) + pass + pass + # Log the metrics + # completion_length = self.accelerator.gather_for_metrics(completion_mask.sum(1)).float().mean().item() + # mean_kl = ((per_token_kl * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean() + # self._metrics["kl"].append(self.accelerator.gather_for_metrics(mean_kl).mean().item()) + if "train" in self._metrics: + mode = "eval" if self.control.should_evaluate else "train" + self._metrics[mode]["completion_length"].append(completion_length.item()) + self._metrics[mode]["kl"].append(mean_kl.item()) + else: + self._metrics["completion_length"].append(completion_length.item()) + self._metrics["kl"].append(mean_kl.item()) + return loss + + def _compute_loss(self, model, inputs): + # Compute the per-token log probabilities for the model + prompt_ids, prompt_mask = inputs["prompt_ids"], inputs["prompt_mask"] + completion_ids, completion_mask = inputs["completion_ids"], inputs["completion_mask"] + input_ids = torch.cat([prompt_ids, completion_ids], dim=1) + attention_mask = torch.cat([prompt_mask, completion_mask], dim=1) + logits_to_keep = completion_ids.size(1) # we only need to compute the logits for the completion tokens + + per_token_logps = self._get_per_token_logps(model, input_ids, attention_mask, logits_to_keep) + + # Compute the KL divergence between the model and the reference model + if self.beta != 0.0: + ref_per_token_logps = inputs["ref_per_token_logps"] + per_token_kl = ( + torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1 + ) + + # Compute the loss + advantages = inputs["advantages"] + # When using num_iterations == 1 and steps_per_generation <= gradient_accumulation_steps + # old_per_token_logps == per_token_logps, so we can skip it's computation + # (see _generate_and_score_completions) and use per_token_logps.detach() instead. + old_per_token_logps = ( + per_token_logps.detach() if inputs["old_per_token_logps"] is None else inputs["old_per_token_logps"] + ) + coef_1 = torch.exp(per_token_logps - old_per_token_logps) + coef_2 = torch.clamp(coef_1, 1 - self.epsilon_low, 1 + self.epsilon_high) + + # Two-sided clipping + if self.args.delta is not None: + coef_1 = torch.clamp(coef_1, max=self.args.delta) + + per_token_loss1 = coef_1 * advantages.unsqueeze(1) + per_token_loss2 = coef_2 * advantages.unsqueeze(1) + per_token_loss = -torch.min(per_token_loss1, per_token_loss2) + if self.beta != 0.0: + per_token_loss = per_token_loss + self.beta * per_token_kl + + if self.loss_type == "grpo": + loss = ((per_token_loss * completion_mask).sum(-1) / completion_mask.sum(-1).clamp(min=1.0)).mean() + elif self.loss_type == "bnpo": + loss = (per_token_loss * completion_mask).sum() / completion_mask.sum().clamp(min=1.0) + elif self.loss_type == "dr_grpo": + loss = (per_token_loss * completion_mask).sum() / (per_token_loss.size(0) * self.max_completion_length) + else: + raise ValueError(f"Unknown loss type: {self.loss_type}") + + # Log the metrics + mode = "train" if self.model.training else "eval" + + if self.beta != 0.0: + mean_kl = (per_token_kl * completion_mask).sum() / completion_mask.sum() + self._metrics[mode]["kl"].append(self.accelerator.gather(mean_kl).nanmean().item()) + + # Compute the clipped probability ratios + is_low_clipped = (coef_1 < 1 - self.epsilon_low) & (advantages.unsqueeze(1) < 0) + is_high_clipped = (coef_1 > 1 + self.epsilon_high) & (advantages.unsqueeze(1) > 0) + is_region_clipped = is_low_clipped | is_high_clipped + + low_clip = (is_low_clipped * completion_mask).sum() / completion_mask.sum() + high_clip = (is_high_clipped * completion_mask).sum() / completion_mask.sum() + clip_ratio = (is_region_clipped * completion_mask).sum() / completion_mask.sum() + + gathered_low_clip = self.accelerator.gather(low_clip) + self._metrics[mode]["clip_ratio/low_mean"].append(gathered_low_clip.nanmean().item()) + self._metrics[mode]["clip_ratio/low_min"].append(nanmin(gathered_low_clip).item()) + gathered_high_clip = self.accelerator.gather(high_clip) + self._metrics[mode]["clip_ratio/high_mean"].append(gathered_high_clip.nanmean().item()) + self._metrics[mode]["clip_ratio/high_max"].append(nanmax(gathered_high_clip).item()) + gathered_clip_ratio = self.accelerator.gather(clip_ratio) + self._metrics[mode]["clip_ratio/region_mean"].append(gathered_clip_ratio.nanmean().item()) + return loss + + def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys: Optional[list[str]] = None): + inputs = self._prepare_inputs(inputs) + with torch.no_grad(): + with self.compute_loss_context_manager(): + loss = self.compute_loss(model, inputs) + loss = loss.mean().detach() + return loss, None, None + + def log(self, logs: dict[str, float], start_time: Optional[float] = None) -> None: + mode = "train" if self.model.training else "eval" + metrics = {key: sum(val) / len(val) for key, val in self._metrics[mode].items()} # average the metrics + + # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs` + # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format. + if mode == "eval": + metrics = {f"eval_{key}": val for key, val in metrics.items()} + + logs = {**logs, **metrics} + super().log(logs, start_time) + self._metrics[mode].clear() + + if self.accelerator.is_main_process and self.log_completions: + if is_rich_available(): + print_prompt_completions_sample( + self._textual_logs["prompt"], + self._textual_logs["completion"], + self._textual_logs["rewards"], + self._textual_logs["advantages"], + self.state.global_step, + self.num_completions_to_print, + ) + + if self.args.report_to and "wandb" in self.args.report_to and wandb.run is not None: + import pandas as pd + + table = { + "step": [str(self.state.global_step)] * len(self._textual_logs["prompt"]), + "prompt": self._textual_logs["prompt"], + "completion": self._textual_logs["completion"], + **self._textual_logs["rewards"], + "advantage": self._textual_logs["advantages"], + } + df = pd.DataFrame(table) + if self.wandb_log_unique_prompts: + df = df.drop_duplicates(subset=["prompt"]) + wandb.log({"completions": wandb.Table(dataframe=df)}) + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) + + 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 + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + citation = textwrap.dedent( + """\ + @article{zhihong2024deepseekmath, + title = {{DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models}}, + author = {Zhihong Shao and Peiyi Wang and Qihao Zhu and Runxin Xu and Junxiao Song and Mingchuan Zhang and Y. K. Li and Y. Wu and Daya Guo}, + year = 2024, + eprint = {arXiv:2402.03300}, + } + """ + ) + + 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="GRPO", + trainer_citation=citation, + paper_title="DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models", + paper_id="2402.03300", + ) + + model_card.save(os.path.join(self.args.output_dir, "README.md")) +class UnslothGRPOTrainer(_UnslothGRPOTrainer): + """ + + Trainer for the Group Relative Policy Optimization (GRPO) method. This algorithm was initially proposed in the + paper [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language + Models](https://huggingface.co/papers/2402.03300). + + Example: + + ```python + from datasets import load_dataset + from trl import GRPOTrainer + + dataset = load_dataset("trl-lib/tldr", split="train") + def reward_func(completions, **kwargs): + # Dummy reward function that rewards completions with more unique letters. + return [float(len(set(completion))) for completion in completions] + trainer = GRPOTrainer( + model="Qwen/Qwen2-0.5B-Instruct", + reward_funcs=reward_func, + train_dataset=dataset, + ) + + trainer.train() + ``` + + Args: + model (`Union[str, PreTrainedModel]`): + Model to be trained. Can be either: + + - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or a + path to a *directory* containing model weights saved using + [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded + using [`~transformers.AutoModelForCausalLM.from_pretrained`] with the keyword arguments in + `args.model_init_kwargs`. + - A [`~transformers.PreTrainedModel`] object. Only causal language models are supported. + reward_funcs (`Union[RewardFunc, list[RewardFunc]]`): + Reward functions to be used for computing the rewards. To compute the rewards, we call all the reward + functions with the prompts and completions and sum the rewards. Can be either: + + - A single reward function, such as: + - A string: The *model ID* of a pretrained model hosted inside a model repo on huggingface.co, or a + path to a *directory* containing model weights saved using + [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded + using [`~transformers.AutoModelForSequenceClassification.from_pretrained`] with `num_labels=1` and the + keyword arguments in `args.model_init_kwargs`. + - A [`~transformers.PreTrainedModel`] object: Only sequence classification models are supported. + - A custom reward function: The function is provided with the prompts and the generated completions, + plus any additional columns in the dataset. It should return a list of rewards. Custom reward + functions can also return None when the reward is not applicable to those samples. This is useful for + multi-task training where different reward functions apply to different types of samples. When a + reward function returns None for a sample, that reward function is excluded from the reward + calculation for that sample. For more details, see [Using a custom reward + function](#using-a-custom-reward-function). + - A list of reward functions, where each item can independently be any of the above types. Mixing different + types within the list (e.g., a string model ID and a custom reward function) is allowed. + args ([`GRPOConfig`], *optional*, defaults to `None`): + Configuration for this trainer. If `None`, a default configuration is used. + train_dataset ([`~datasets.Dataset`] or [`~datasets.IterableDataset`]): + Dataset to use for training. It must include a column `"prompt"`. Any additional columns in the dataset is + ignored. The format of the samples can be either: + + - [Standard](dataset_formats#standard): Each sample contains plain text. + - [Conversational](dataset_formats#conversational): Each sample contains structured messages (e.g., role + and content). + eval_dataset ([`~datasets.Dataset`], [`~datasets.IterableDataset`] or `dict[str, Union[Dataset, IterableDataset]]`): + Dataset to use for evaluation. It must meet the same requirements as `train_dataset`. + processing_class ([`~transformers.PreTrainedTokenizerBase`], *optional*, defaults to `None`): + Processing class used to process the data. The padding side must be set to "left". If `None`, the + processing class is loaded from the model's name with [`~transformers.AutoTokenizer.from_pretrained`]. A + padding token, `processing_class.pad_token`, must be set. If the processing class has not set a padding + token, `processing_class.eos_token` will be used as the default. + reward_processing_classes (`Union[PreTrainedTokenizerBase, list[PreTrainedTokenizerBase]]`, *optional*, defaults to `None`): + Processing classes corresponding to the reward functions specified in `reward_funcs`. Can be either: + + - A single processing class: Used when `reward_funcs` contains only one reward function. + - A list of processing classes: Must match the order and length of the reward functions in `reward_funcs`. + If set to `None`, or if an element of the list corresponding to a [`~transformers.PreTrainedModel`] is + `None`, the tokenizer for the model is automatically loaded using + [`~transformers.AutoTokenizer.from_pretrained`]. For elements in `reward_funcs` that are custom reward + functions (not [`~transformers.PreTrainedModel`]), the corresponding entries in `reward_processing_classes` + are ignored. + callbacks (list of [`~transformers.TrainerCallback`], *optional*, defaults to `None`): + List of callbacks to customize the training loop. Will add those to the list of default callbacks detailed + in [here](https://huggingface.co/docs/transformers/main_classes/callback). + + If you want to remove one of the default callbacks used, use the [`~transformers.Trainer.remove_callback`] + method. + optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`, *optional*, defaults to `(None, None)`): + A tuple containing the optimizer and the scheduler to use. Will default to an instance of [`AdamW`] on your + model and a scheduler given by [`get_linear_schedule_with_warmup`] controlled by `args`. + peft_config ([`~peft.PeftConfig`], *optional*, defaults to `None`): + PEFT configuration used to wrap the model. If `None`, the model is not wrapped. + + """ + def __init__( + self, + model, + reward_funcs, + args = None, + train_dataset = None, + eval_dataset = None, + processing_class = None, + reward_processing_classes = None, + callbacks = None, + peft_config = None, + **kwargs + ): + if args is None: args = UnslothGRPOConfig() + 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 + 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) + 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 + 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' + other_metrics = [] + if not isinstance(reward_funcs, list): _reward_funcs = [reward_funcs] + else: _reward_funcs = reward_funcs + for reward_func in _reward_funcs: + try: + reward_func_name = reward_func.__name__ + if True: + other_metrics.append(f'rewards/{reward_func_name}/mean') + if True: + other_metrics.append(f'rewards/{reward_func_name}/std') + if False: + other_metrics.append(f'rewards/{reward_func_name}') + except: pass + + from unsloth_zoo.logging_utils import PatchRLStatistics + PatchRLStatistics('grpo_trainer', other_metrics) + + super().__init__( + model = model, + reward_funcs = reward_funcs, + args = args, + train_dataset = train_dataset, + eval_dataset = eval_dataset, + processing_class = processing_class, + reward_processing_classes = reward_processing_classes, + callbacks = callbacks, + 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 diff --git a/unsloth_compiled_cache/UnslothIterativeSFTTrainer.py b/unsloth_compiled_cache/UnslothIterativeSFTTrainer.py new file mode 100644 index 0000000..1db805b --- /dev/null +++ b/unsloth_compiled_cache/UnslothIterativeSFTTrainer.py @@ -0,0 +1,947 @@ +""" +2025.7.1 +2025.7.1 +4.53.1 +0.19.1 +__UNSLOTH_VERSIONING__ +""" +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from trl.trainer.iterative_sft_trainer import (AutoModelForCausalLM, AutoTokenizer, BaseImageProcessor, Callable, DataCollator, DataCollatorForLanguageModeling, DataCollatorForSeq2Seq, DataLoader, Dataset, EvalLoopOutput, FeatureExtractionMixin, IterativeSFTConfig, IterativeSFTTrainer, Optional, PPODecorators, Path, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, Trainer, TrainingArguments, Union, generate_model_card, get_comet_experiment_url, is_peft_available, is_wandb_available, os, torch, warnings, Optional, PeftModel, PreTrainedModel, Trainer, is_peft_available, os, torch) + + +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 UnslothIterativeSFTConfig(IterativeSFTConfig): + """ + + Configuration class for the [`IterativeSFTTrainer`]. + + This class includes only the parameters that are specific to Iterative SFT training. For a full list of training + arguments, please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this + class may differ from those in [`~transformers.TrainingArguments`]. + + 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 + + model_init_kwargs (`dict[str, Any]` or `None`, *optional*, defaults to `None`): + Keyword arguments for [`~transformers.AutoModelForCausalLM.from_pretrained`], used when the `model` + argument of the [`IterativeSFTTrainer`] is provided as a string. + + > Parameters that control the data preprocessing + + max_length (`int` or `None`, *optional*, defaults to `None`): + Maximum length of the tokenized sequence. Sequences longer than `max_length` are truncated. + truncation_mode (`str`, *optional*, defaults to `"keep_end"`): + The truncation mode to use, either `"keep_end"` or `"keep_start"`. + optimize_device_cache (`bool`, *optional*, defaults to `False`): + Whether to optimize accelerator cache for slightly more memory-efficient training. + + """ + 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, + hub_revision = None, + 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, + liger_kernel_config = None, + eval_use_gather_object = False, + average_tokens_across_devices = False, + model_init_kwargs = None, + max_length = None, + truncation_mode = 'keep_end', + optimize_device_cache = 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' + + 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, + 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, + model_init_kwargs = model_init_kwargs, + max_length = max_length, + truncation_mode = truncation_mode, + optimize_device_cache = optimize_device_cache,**kwargs) + self.vllm_sampling_params = vllm_sampling_params + self.unsloth_num_chunks = unsloth_num_chunks +pass + +class _UnslothIterativeSFTTrainer(Trainer): + """""" + + _tag_names = ["trl", "iterative-sft"] + + def __init__( + self, + model: Union[str, PreTrainedModel], + args: Optional[Union[IterativeSFTConfig, TrainingArguments]] = None, + data_collator: Optional[DataCollator] = None, + eval_dataset: Optional[Union[Dataset, dict[str, Dataset]]] = None, + processing_class: Optional[ + Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin] + ] = 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, + compute_metrics: Optional[Callable[[EvalLoopOutput], dict]] = None, + # Deprecated parameters + max_length: Optional[int] = None, + truncation_mode: Optional[str] = None, + optimize_device_cache: Optional[bool] = None, + ): + # Handle deprecated parameters + deprecated_params = {} + if max_length is not None: + deprecated_params["max_length"] = max_length + warnings.warn( + "The `max_length` parameter is deprecated and will be removed in version 0.20. " + "Pass it through the `args` parameter using `IterativeSFTConfig(max_length=...)` instead.", + DeprecationWarning, + ) + if truncation_mode is not None: + deprecated_params["truncation_mode"] = truncation_mode + warnings.warn( + "The `truncation_mode` parameter is deprecated and will be removed in version 0.20. " + "Pass it through the `args` parameter using `IterativeSFTConfig(truncation_mode=...)` instead.", + DeprecationWarning, + ) + if optimize_device_cache is not None: + deprecated_params["optimize_device_cache"] = optimize_device_cache + warnings.warn( + "The `optimize_device_cache` parameter is deprecated and will be removed in version 0.20 " + "Pass it through the `args` parameter using `IterativeSFTConfig(optimize_device_cache=...)` instead.", + DeprecationWarning, + ) + + # Args + model_id = model if isinstance(model, str) else model.config._name_or_path + if args is None: + model_name = model_id.split("/")[-1] + args = IterativeSFTConfig(f"{model_name}-IterativeSFT") + elif isinstance(args, TrainingArguments) and not isinstance(args, IterativeSFTConfig): + dict_args = args.to_dict() + dict_args["hub_token"] = args.hub_token # to_dict hides the hub_token + dict_args.pop("push_to_hub_token") + args = IterativeSFTConfig(**dict_args) + + # Update args with deprecated parameters if provided + if deprecated_params: + for key, value in deprecated_params.items(): + setattr(args, key, value) + + # Handle the tokenizer + if processing_class is None: + processing_class = AutoTokenizer.from_pretrained(model_id) + + # Model + if args.model_init_kwargs is not None and not isinstance(model, str): + warnings.warn( + "You passed model_init_kwargs to the `IterativeSFTConfig`, but your model is already instantiated. " + "The `model_init_kwargs` will be ignored." + ) + if isinstance(model, str): + model = self._create_model_from_path(model, args) + + # PEFT configuration and model wrapping + if is_peft_available() and isinstance(model, PeftModel): + self.is_peft_model = True + else: + self.is_peft_model = False + + self.processing_class = processing_class + self.is_encoder_decoder = getattr(model.config, "is_encoder_decoder", False) + + if data_collator is None: + if self.is_encoder_decoder: + self.data_collator = DataCollatorForSeq2Seq( + processing_class, label_pad_token_id=-100, pad_to_multiple_of=8 + ) + else: + self.data_collator = DataCollatorForLanguageModeling(self.processing_class, mlm=False) + else: + self.data_collator = data_collator + + self.max_length = args.max_length + self.truncation_mode = args.truncation_mode + self.optimize_device_cache = args.optimize_device_cache + + super().__init__( + model=model, + args=args, + data_collator=self.data_collator, + eval_dataset=eval_dataset, + processing_class=processing_class, + compute_metrics=compute_metrics, + optimizers=optimizers, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + ) + + # 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) + + self.create_optimizer_and_scheduler(self.args.max_steps) + + # prepare model, optimizer and lr_scheduler + self.model, self.optimizer, self.lr_scheduler = self.accelerator.prepare( + self.model, self.optimizer, self.lr_scheduler + ) + + self.processing_class.truncation_side = "left" if self.truncation_mode == "keep_end" else "right" + + if not hasattr(self, "accelerator"): + raise AttributeError( + "Your `Trainer` does not have an `accelerator` object. Consider upgrading `transformers`." + ) + + PPODecorators.optimize_device_cache = self.optimize_device_cache + + def _create_model_from_path(self, model_path: str, args: IterativeSFTConfig) -> PreTrainedModel: + """Creates a model from a path or model identifier.""" + model_init_kwargs = args.model_init_kwargs or {} + return AutoModelForCausalLM.from_pretrained(model_path, **model_init_kwargs) + + def prepare_model_inputs(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, labels: torch.Tensor): + if attention_mask is None: + attention_mask = [torch.ones_like(ids) for ids in input_ids] + + if self.is_encoder_decoder: + input_data = self.data_collator( + [ + {"input_ids": ids, "attention_mask": att, "labels": lab} + for ids, att, lab in zip(input_ids, attention_mask, labels) + ] + ).to(self.model.device) + + input_data.pop("decoder_input_ids", None) # This is directly computed inside the model + + input_data["labels"][input_data["labels"] == self.processing_class.pad_token_id] = -100 + + else: + input_data = self.data_collator( + [{"input_ids": ids, "attention_mask": att} for ids, att in zip(input_ids, attention_mask)] + ).to(self.model.device) + + # truncate in case the user has provided input_ids, attention_mask and labels + if self.max_length is not None: + if self.truncation_mode == "keep_start": + input_data = {k: v[: self.max_length] for k, v in input_data.items()} + elif self.truncation_mode == "keep_end": + input_data = {k: v[-self.max_length :] for k, v in input_data.items()} + else: + raise ValueError(f"Unknown truncation mode: {self.truncation_mode}") + + return input_data + + @staticmethod + def _step_safety_checker( + input_ids: list[torch.LongTensor], + attention_mask: list[torch.LongTensor], + labels: list[torch.LongTensor], + texts: list[str], + texts_labels: list[str], + ): + """ + Check if the input data is valid for training. + + Args: + input_ids (list[`torch.LongTensor`]): + List of tensors containing the input_ids + attention_mask (list[`torch.LongTensor`]): + List of tensors containing the attention_mask + labels (list[`torch.FloatTensor`]): + List of tensors containing the labels + texts (list[`str`]): + List of string containing the text input. + texts_labels (list[`str`]): + List of string containing the text labels. + + Returns: + `tuple`: The input data. + """ + if texts is None: + if attention_mask is None: + for name, tensor_list in zip(["input_ids", "labels"], [input_ids, labels]): + if not isinstance(tensor_list, list): + raise ValueError(f"{name} must be a list of tensors - got {type(tensor_list)}") + if not isinstance(tensor_list[0], torch.Tensor): + raise ValueError(f"Elements in {name} must be tensors - got {type(tensor_list[0])}") + else: + for name, tensor_list in zip( + ["input_ids", "attention_mask", "labels"], [input_ids, attention_mask, labels] + ): + if not isinstance(tensor_list, list): + raise ValueError(f"{name} must be a list of tensors - got {type(tensor_list)}") + if not isinstance(tensor_list[0], torch.Tensor): + raise ValueError(f"Elements in {name} must be tensors - got {type(tensor_list[0])}") + else: + if not isinstance(texts, list): + raise ValueError(f"'text' must be a list of strings - got {type(texts)}") + if not isinstance(texts[0], str): + raise ValueError(f"Elements in 'text' must be strings - got {type(texts[0])}") + if texts_labels is not None: + if not isinstance(texts_labels, list): + raise ValueError(f"'text_labels' must be a list of strings - got {type(texts_labels)}") + if not isinstance(texts_labels[0], str): + raise ValueError(f"Elements in 'text_labels' must be strings - got {type(texts_labels[0])}") + + return input_ids, attention_mask, labels, texts, texts_labels + + @PPODecorators.empty_device_cache() + def step( + self, + input_ids: Optional[list[torch.LongTensor]] = None, + attention_mask: Optional[list[torch.LongTensor]] = None, + labels: Optional[list[torch.LongTensor]] = None, + texts: Optional[list[str]] = None, + texts_labels: Optional[list[str]] = None, + ): + """ + Run an optimisation step given a list of input_ids, attention_mask, and labels or a list of text and + text_labels. + + Args: + input_ids (list[`torch.LongTensor`]): + List of tensors containing the input_ids (if not provided, text will be used) + attention_mask (list[`torch.LongTensor`], , *optional*): + List of tensors containing the attention_mask + labels (list[`torch.FloatTensor`], *optional*): + List of tensors containing the labels (if set to None, will default to input_ids) + texts (list[`str`], *optional*): + List of strings containing the text input (if not provided, input_ids will directly be used) + texts_labels (list[`str`], *optional*): + List of strings containing the text labels (if set to None, will default to text) + + Returns: + `dict[str, Any]`: A summary of the training statistics + """ + self.model.train() + + if self.state.global_step == 0: + self.tr_loss = torch.tensor(0.0).to(self.args.device) + self._globalstep_last_logged = self.state.global_step + + if input_ids is None and texts is None: + raise ValueError("Step should include `input_ids` or `texts` as keyword arguments.") + elif input_ids is not None and texts is not None: + warnings.warn( + "Both `input_ids` and `texts` argument are provided. `input_ids` will be ignored. " + "Please provide only one of the two.", + UserWarning, + ) + + if labels is None and texts_labels is None and self.is_encoder_decoder: + raise ValueError( + "No 'labels' or 'text_labels' are provided. When using an encoder-decoder architecture, 'labels' or 'text_labels' must be passed." + ) + + # Convert Column to list if not already + input_ids = input_ids[:] if input_ids is not None else None + attention_mask = attention_mask[:] if attention_mask is not None else None + labels = labels[:] if labels is not None else None + texts = texts[:] if texts is not None else None + texts_labels = texts_labels[:] if texts_labels is not None else None + + input_ids, attention_mask, labels, texts, texts_labels = self._step_safety_checker( + input_ids, attention_mask, labels, texts, texts_labels + ) + + if texts is not None: + model_inputs = self.processing_class( + texts, max_length=self.max_length, truncation=True, padding=True, return_tensors="pt" + ) + + input_ids, attention_mask = model_inputs["input_ids"], model_inputs["attention_mask"] + + if texts_labels is not None: + labels = self.processing_class( + texts, max_length=self.max_length, truncation=True, padding=True, return_tensors="pt" + )["input_ids"] + + if labels is None: + labels = input_ids + + model_inputs = self.prepare_model_inputs(input_ids, attention_mask, labels) + + model_inputs_names = list(model_inputs.keys()) + + batch_dict = {} + batch_dict.update(model_inputs) + + def collator(data): + return_dict = dict() + for key in data[0]: + if key in ["input_ids", "attention_mask", "labels"]: + return_dict[key] = torch.stack([d[key] for d in data]).to(self.model.device) + return return_dict + + batch_data = Dataset.from_dict(batch_dict) + batch_data.set_format("torch") + + step_dataloader = DataLoader( + batch_data, + batch_size=self.args.per_device_train_batch_size, + shuffle=True, + collate_fn=collator, + ) + + for _, batch in enumerate(step_dataloader): + with self.accelerator.accumulate(self.model): + model_inputs = {k: batch[k] for k in model_inputs_names} + loss = self.compute_loss(self.model, model_inputs) + + if self.args.n_gpu > 1: + loss = loss.mean() + + tr_loss_step = loss.detach() + + self.accelerator.backward(loss) + + if self.accelerator.sync_gradients and self.args.max_grad_norm is not None: + self.accelerator.clip_grad_norm_( + self.model.parameters(), + self.args.max_grad_norm, + ) + + self.optimizer.step() + self.optimizer.zero_grad() + if self.lr_scheduler is not None: + self.lr_scheduler.step() + + self.state.global_step += 1 + + # update stats etc + self.tr_loss += tr_loss_step + + self._maybe_log_save_evaluate() + + def _maybe_log_save_evaluate(self): + # check if eval is required + if self.args.eval_steps is not None: + if self.state.global_step % self.args.eval_steps == 0 and self.state.global_step != 0: + self.evaluate(self.eval_dataset) + + # check if logging is required + if self.args.logging_steps is not None: + if self.state.global_step % self.args.logging_steps == 0 and self.state.global_step != 0: + logs: dict[str, float] = {} + + tr_loss_scalar = self._nested_gather(self.tr_loss).mean().item() + + # reset tr_loss to zero + self.tr_loss -= self.tr_loss + + logs["loss"] = round(tr_loss_scalar / (self.state.global_step - self._globalstep_last_logged), 4) + logs["learning_rate"] = self._get_learning_rate() + + self._globalstep_last_logged = self.state.global_step + + self.log(logs) + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) + + 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 + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + 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="Iterative SFT", + ) + + model_card.save(os.path.join(self.args.output_dir, "README.md")) +class UnslothIterativeSFTTrainer(_UnslothIterativeSFTTrainer): + """ + + The IterativeSFTTrainer can be used to finetune models with methods that requires some steps between optimization. + + Args: + model (`Union[str, PreTrainedModel]`): + Model to be trained. Can be either: + + - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or a + path to a *directory* containing model weights saved using + [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded + using [`~transformers.AutoModelForCausalLM.from_pretrained`] with the keyword arguments in + `args.model_init_kwargs`. + - A [`~transformers.PreTrainedModel`] object. Only causal language models are supported. + args ([`IterativeSFTConfig`], *optional*, defaults to `None`): + Configuration for this trainer. If `None`, a default configuration is used. + data_collator (`DataCollator`, *optional*): + Function to use to form a batch from a list of elements of the processed `train_dataset` or `eval_dataset`. + Will default to [`~transformers.default_data_collator`] if no `processing_class` is provided, an instance + of [`~transformers.DataCollatorWithPadding`] otherwise if the processing_class is a feature extractor or + tokenizer. + eval_dataset (`datasets.Dataset`): + The dataset to use for evaluation. + processing_class ([`~transformers.PreTrainedTokenizerBase`], *optional*, defaults to `None`): + Processing class used to process the data. If `None`, the processing class is loaded from the model's name + with [`~transformers.AutoTokenizer.from_pretrained`]. + 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. + 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. + max_length (`int`, *optional*, deprecated): + Maximum length of the tokenized sequence. Use `args.max_length` instead. + truncation_mode (`str`, *optional*, deprecated): + The truncation mode to use. Use `args.truncation_mode` instead. + optimize_device_cache (`bool`, *optional*, deprecated): + Whether to optimize accelerator cache. Use `args.optimize_device_cache` instead. + + """ + def __init__( + self, + model, + args = None, + data_collator = None, + eval_dataset = None, + processing_class = None, + preprocess_logits_for_metrics = None, + compute_metrics = None, + max_length = None, + truncation_mode = None, + optimize_device_cache = None, + **kwargs + ): + if args is None: args = UnslothIterativeSFTConfig() + 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 + 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) + 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 + 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' + other_metrics = [] + + from unsloth_zoo.logging_utils import PatchRLStatistics + PatchRLStatistics('iterative_sft_trainer', other_metrics) + + super().__init__( + model = model, + args = args, + data_collator = data_collator, + eval_dataset = eval_dataset, + processing_class = processing_class, + preprocess_logits_for_metrics = preprocess_logits_for_metrics, + compute_metrics = compute_metrics, + max_length = max_length, + truncation_mode = truncation_mode, + optimize_device_cache = optimize_device_cache,**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 diff --git a/unsloth_compiled_cache/UnslothKTOTrainer.py b/unsloth_compiled_cache/UnslothKTOTrainer.py new file mode 100644 index 0000000..216c7b4 --- /dev/null +++ b/unsloth_compiled_cache/UnslothKTOTrainer.py @@ -0,0 +1,2034 @@ +""" +2025.7.1 +2025.7.1 +4.53.1 +0.19.1 +__UNSLOTH_VERSIONING__ +""" +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from trl.trainer.kto_trainer import (Any, AutoModelForCausalLM, BaseImageProcessor, Callable, DPODataCollatorWithPadding, DataCollator, DataLoader, Dataset, EvalLoopOutput, F, FeatureExtractionMixin, KTOConfig, KTOTrainer, Literal, Optional, PartialState, Path, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, SequentialSampler, Trainer, TrainerCallback, TrainingArguments, Union, _get_kl_dataset, _process_tokens, _tokenize, autocast, concatenate_datasets, contextmanager, create_reference_model, defaultdict, disable_dropout_in_model, generate_model_card, get_comet_experiment_url, has_length, inspect, is_comet_available, is_liger_kernel_available, is_peft_available, is_wandb_available, itemgetter, log_table_to_comet_experiment, maybe_apply_chat_template, maybe_extract_prompt, maybe_unpair_preference_dataset, nn, np, nullcontext, os, pad_to_length, pd, peft_module_casting_to_bf16, prepare_deepspeed, prepare_model_for_kbit_training, random, textwrap, torch, tqdm, warnings, F, Optional, PeftModel, PreTrainedModel, Trainer, is_peft_available, os, torch) + + +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 UnslothKTOConfig(KTOConfig): + """ + + Configuration class for the [`KTOTrainer`]. + + This class includes only the parameters that are specific to KTO training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may + differ from those in [`~transformers.TrainingArguments`]. + + 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: + max_length (`int` or `None`, *optional*, defaults to `1024`): + Maximum length of the sequences (prompt + completion) in the batch. This argument is required if you want + to use the default data collator. + max_prompt_length (`int` or `None`, *optional*, defaults to `512`): + Maximum length of the prompt. This argument is required if you want to use the default data collator. + max_completion_length (`int` or `None`, *optional*, defaults to `None`): + Maximum length of the completion. This argument is required if you want to use the default data collator + and your model is an encoder-decoder. + beta (`float`, *optional*, defaults to `0.1`): + Parameter controlling the deviation from the reference model. Higher β means less deviation from the + reference model. + loss_type (`str`, *optional*, defaults to `"kto"`): + Type of loss to use. Possible values are: + + - `"kto"`: KTO loss from the [KTO](https://huggingface.co/papers/2402.01306) paper. + - `"apo_zero_unpaired"`: Unpaired variant of APO-zero loss from the + [APO](https://huggingface.co/papers/2408.06266) paper. + + desirable_weight (`float`, *optional*, defaults to `1.0`): + Desirable losses are weighed by this factor to counter unequal number of desirable and undesirable paris. + undesirable_weight (`float`, *optional*, defaults to `1.0`): + Undesirable losses are weighed by this factor to counter unequal number of desirable and undesirable pairs. + label_pad_token_id (`int`, *optional*, defaults to `-100`): + Label pad token id. This argument is required if you want to use the default data collator. + padding_value (`int` or `None`, *optional*, defaults to `None`): + Padding value to use. If `None`, the padding value of the tokenizer is used. + truncation_mode (`str`, *optional*, defaults to `"keep_end"`): + Truncation mode to use when the prompt is too long. Possible values are `"keep_end"` or `"keep_start"`. + This argument is required if you want to use the default data collator. + generate_during_eval (`bool`, *optional*, defaults to `False`): + If `True`, generates and logs completions from both the model and the reference model to W&B or Comet + during evaluation. + is_encoder_decoder (`bool` or `None`, *optional*, defaults to `None`): + When using the `model_init` argument (callable) to instantiate the model instead of the `model` argument, + you need to specify if the model returned by the callable is an encoder-decoder model. + precompute_ref_log_probs (`bool`, *optional*, defaults to `False`): + Whether to precompute reference model log probabilities for training and evaluation datasets. This is + useful when training without the reference model to reduce the total GPU memory needed. + model_init_kwargs (`dict[str, Any]` or `None`, *optional*, defaults to `None`): + Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model from a + string. + ref_model_init_kwargs (`dict[str, Any]` or `None`, *optional*, defaults to `None`): + Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the reference model + from a string. + dataset_num_proc: (`int` or `None`, *optional*, defaults to `None`): + Number of processes to use for processing the dataset. + disable_dropout (`bool`, *optional*, defaults to `True`): + Whether to disable dropout in the model and reference model. + use_liger_loss (`bool`, *optional*, defaults to `False`): + Whether to use Liger loss. It requires liger-kernel to be installed. + base_model_attribute_name (`str`, *optional*, defaults to `"model"`): + Name of the attribute in the model that contains the base model. This is used to get the base model from + the model when the model does not have a `get_decoder` method in the case when `use_liger_loss` is `True`. + + """ + 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, + hub_revision = None, + 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, + liger_kernel_config = None, + eval_use_gather_object = False, + average_tokens_across_devices = False, + max_length = 1024, + max_prompt_length = 512, + max_completion_length = None, + beta = 0.1, + loss_type = 'kto', + desirable_weight = 1.0, + undesirable_weight = 1.0, + label_pad_token_id = -100, + padding_value = None, + truncation_mode = 'keep_end', + generate_during_eval = False, + is_encoder_decoder = None, + disable_dropout = True, + precompute_ref_log_probs = False, + model_init_kwargs = None, + ref_model_init_kwargs = None, + dataset_num_proc = None, + use_liger_loss = False, + base_model_attribute_name = 'model', + 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, + 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, + max_length = max_length, + max_prompt_length = max_prompt_length, + max_completion_length = max_completion_length, + beta = beta, + loss_type = loss_type, + desirable_weight = desirable_weight, + undesirable_weight = undesirable_weight, + label_pad_token_id = label_pad_token_id, + padding_value = padding_value, + truncation_mode = truncation_mode, + generate_during_eval = generate_during_eval, + is_encoder_decoder = is_encoder_decoder, + disable_dropout = disable_dropout, + precompute_ref_log_probs = precompute_ref_log_probs, + model_init_kwargs = model_init_kwargs, + ref_model_init_kwargs = ref_model_init_kwargs, + dataset_num_proc = dataset_num_proc, + use_liger_loss = use_liger_loss, + base_model_attribute_name = base_model_attribute_name,**kwargs) + self.vllm_sampling_params = vllm_sampling_params + self.unsloth_num_chunks = unsloth_num_chunks +pass + +class _UnslothKTOTrainer(Trainer): + r"""""" + + _tag_names = ["trl", "kto"] + + def __init__( + self, + model: Union[PreTrainedModel, nn.Module, str] = None, + ref_model: Optional[Union[PreTrainedModel, nn.Module, str]] = None, + args: KTOConfig = None, + train_dataset: Optional[Dataset] = None, + eval_dataset: Optional[Union[Dataset, dict[str, Dataset]]] = None, + processing_class: Optional[ + Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin] + ] = None, + data_collator: Optional[DataCollator] = None, + model_init: Optional[Callable[[], PreTrainedModel]] = 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, + compute_metrics: Optional[Callable[[EvalLoopOutput], dict]] = None, + model_adapter_name: Optional[str] = None, + ref_adapter_name: Optional[str] = None, + ): + if type(args) is TrainingArguments: + raise ValueError("Please use `KTOConfig` instead TrainingArguments.") + + 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_kwargs to the KTOTrainer. 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 KTOConfig. 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_kwargs to the KTOTrainer. 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 KTOConfig. 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 with `pip install peft` 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 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 args.gradient_checkpointing: + # 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 args.gradient_checkpointing: + # 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." + ) + + if model is not None: + self.is_encoder_decoder = model.config.is_encoder_decoder + elif args.is_encoder_decoder is None: + raise ValueError("When no model is provided, you need to pass the parameter is_encoder_decoder.") + else: + self.is_encoder_decoder = args.is_encoder_decoder + + self.is_peft_model = is_peft_available() and isinstance(model, PeftModel) + self.model_adapter_name = model_adapter_name + self.ref_adapter_name = ref_adapter_name + + 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( + "max_length or a processing_class must be specified when using the default DPODataCollatorWithPadding" + ) + if args.max_length is None: + warnings.warn( + "When using DPODataCollatorWithPadding, you should set `max_length` in the KTOTrainer's init" + " it will be set to `512` by default, but you should do it yourself in the future.", + UserWarning, + ) + max_length = 512 + if args.max_length is not None: + max_length = args.max_length + + if args.max_prompt_length is None: + warnings.warn( + "When using DPODataCollatorWithPadding, you should set `max_prompt_length` in the KTOTrainer's init" + " it will be set to `128` by default, but you should do it yourself in the future.", + UserWarning, + ) + max_prompt_length = 128 + if args.max_prompt_length is not None: + max_prompt_length = args.max_prompt_length + + max_completion_length = None + if args.max_completion_length is None and self.is_encoder_decoder: + warnings.warn( + "When using DPODataCollatorWithPadding with an encoder decoder architecture, you should set `max_completion_length` in the KTOTrainer's init" + " it will be set to `128` by default, but you should do it yourself in the future.", + UserWarning, + ) + max_completion_length = 128 + if args.max_completion_length is not None and self.is_encoder_decoder: + max_completion_length = args.max_completion_length + + if data_collator is None: + data_collator = DPODataCollatorWithPadding( + pad_token_id=processing_class.pad_token_id, + label_pad_token_id=args.label_pad_token_id, + is_encoder_decoder=self.is_encoder_decoder, + ) + + if args.remove_unused_columns: + args.remove_unused_columns = False + # warn users + warnings.warn( + "When using DPODataCollatorWithPadding, you should set `remove_unused_columns=False` in your KTOConfig" + " we have set it for you, but you should do it yourself in the future.", + UserWarning, + ) + + self.use_dpo_data_collator = True + else: + self.use_dpo_data_collator = False + + # 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.loss_type = args.loss_type + self.max_length = max_length + self.generate_during_eval = args.generate_during_eval + self.label_pad_token_id = args.label_pad_token_id + self.padding_value = args.padding_value if args.padding_value is not None else processing_class.pad_token_id + self.max_prompt_length = max_prompt_length + self.truncation_mode = args.truncation_mode + self.max_completion_length = max_completion_length + self.processing_class = processing_class + self.precompute_ref_log_probs = args.precompute_ref_log_probs + + # Not all losses require a KL calculation + self.calculate_KL = True + if self.loss_type in ["apo_zero_unpaired"]: + self.calculate_KL = False + + # 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 + + # metric + self._stored_metrics = defaultdict(lambda: defaultdict(list)) + + # KTO parameter + self.beta = args.beta + self.desirable_weight = args.desirable_weight + self.undesirable_weight = args.undesirable_weight + self.aux_loss_enabled = getattr(model.config, "output_router_logits", False) + 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, + ) + + # 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 KTO, the sampled data does not include the + # "input_ids" key. Instead, the available keys are "prompt_input_ids" and "completion_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 + + # Compute that only on the main process for faster data processing. + # see: https://github.com/huggingface/trl/pull/1255 + with PartialState().main_process_first(): + # Extract the prompt if needed + train_dataset = train_dataset.map( + maybe_extract_prompt, num_proc=args.dataset_num_proc, desc="Extracting prompt from train dataset" + ) + # Unpair the dataset if needed + train_dataset = maybe_unpair_preference_dataset( + train_dataset, args.dataset_num_proc, desc="Unpairing train dataset" + ) + # Apply the chat template if needed + train_dataset = train_dataset.map( + maybe_apply_chat_template, + fn_kwargs={"tokenizer": processing_class}, + num_proc=args.dataset_num_proc, + desc="Applying chat template to train dataset", + ) + if eval_dataset is not None: + eval_dataset = eval_dataset.map( + maybe_extract_prompt, num_proc=args.dataset_num_proc, desc="Extracting prompt from eval dataset" + ) + eval_dataset = maybe_unpair_preference_dataset( + eval_dataset, args.dataset_num_proc, desc="Unpairing eval dataset" + ) + eval_dataset = eval_dataset.map( + maybe_apply_chat_template, + fn_kwargs={"tokenizer": processing_class}, + num_proc=args.dataset_num_proc, + desc="Applying chat template to eval dataset", + ) + + # Tokenize and prepare the training datasets + train_dataset = train_dataset.map( + _tokenize, + batched=True, + fn_kwargs={"tokenizer": self.processing_class}, + num_proc=args.dataset_num_proc, + desc="Tokenizing train dataset", + ) + + fn_kwargs = { + "prefix": "", + "is_encoder_decoder": self.is_encoder_decoder, + "tokenizer": self.processing_class, + "max_length": self.max_length, + "truncation_mode": self.truncation_mode, + "label_pad_token_id": self.label_pad_token_id, + "max_prompt_length": self.max_prompt_length, + "max_completion_length": self.max_completion_length, + } + + train_dataset = train_dataset.map( + _process_tokens, + fn_kwargs=fn_kwargs, + num_proc=args.dataset_num_proc, + desc="Processing tokenized train dataset", + ) + + # Tokenize and prepare the eval datasets + if eval_dataset is not None: + eval_dataset = eval_dataset.map( + _tokenize, + fn_kwargs={"tokenizer": self.processing_class}, + batched=True, + num_proc=args.dataset_num_proc, + desc="Tokenizing eval dataset", + ) + + eval_dataset = eval_dataset.map( + _process_tokens, + fn_kwargs=fn_kwargs, + num_proc=args.dataset_num_proc, + desc="Processing tokenized eval dataset", + ) + + # Get KL datasets if needed + if self.calculate_KL: + if args.per_device_train_batch_size <= 1: + raise ValueError( + "Actual (not effective) batch size must be > 1. KTO will not work properly because the KL term will be equivalent to the implied reward." + ) + + # create pairs for estimating the KL term by flipping the matched pairs in each batch of size total_batch_size + # i.e., [x_1, y_1], ..., [x_n, y_n] --> [x_1, y_n], ..., [x_n, y_1] = [x'_1, y'_1], ..., [x'_n, y'_n] + train_kl_dataset = train_dataset.map( + _get_kl_dataset, + batched=True, + batch_size=args.per_device_train_batch_size, + num_proc=args.dataset_num_proc, + desc="Extracting KL train dataset", + ) + + fn_kwargs["prefix"] = "KL_" + train_kl_dataset = train_kl_dataset.map( + _process_tokens, + fn_kwargs=fn_kwargs, + num_proc=args.dataset_num_proc, + remove_columns=[c for c in train_kl_dataset.column_names if c in train_dataset.column_names], + desc="Processing tokenized train KL dataset", + ) + + # merge the datasets + train_dataset = concatenate_datasets([train_dataset, train_kl_dataset], axis=1) + + if eval_dataset is not None: + # Get KL dataset + eval_kl_dataset = eval_dataset.map( + _get_kl_dataset, + batched=True, + batch_size=args.per_device_train_batch_size, + num_proc=args.dataset_num_proc, + desc="Extracting eval KL dataset", + ) + + eval_kl_dataset = eval_kl_dataset.map( + _process_tokens, + fn_kwargs=fn_kwargs, + num_proc=args.dataset_num_proc, + remove_columns=[c for c in eval_kl_dataset.column_names if c in eval_dataset.column_names], + desc="Processing tokenized eval KL dataset", + ) + + # merge the datasets + eval_dataset = concatenate_datasets([eval_dataset, eval_kl_dataset], axis=1) + + # calculate dataset desirability balance + num_desirable = max(sum(train_dataset["label"]), 1) + num_undesirable = max(len(train_dataset["label"]) - num_desirable, 1) # "label" is binary + + if num_desirable != num_undesirable: + # The lower and upper bounds come from Eq. [8] of https://huggingface.co/papers/2402.01306 + des_weight_lower_bound = round((num_undesirable * self.undesirable_weight / num_desirable) * 1, 2) + des_weight_upper_bound = round((num_undesirable * self.undesirable_weight / num_desirable) * 1.33, 2) + und_weight_lower_bound = round((num_desirable * self.desirable_weight / num_undesirable) / 1.33, 2) + und_weight_upper_bound = round((num_desirable * self.desirable_weight / num_undesirable) / 1, 2) + + des_weight_in_range = des_weight_lower_bound <= self.desirable_weight <= des_weight_upper_bound + und_weight_in_range = und_weight_lower_bound <= self.undesirable_weight <= und_weight_upper_bound + + if not (des_weight_in_range or und_weight_in_range): + warnings.warn( + "You have different amounts of desirable/positive and undesirable/negative examples but the " + "weights on the desirable and undesirable losses don't seem to be in an ideal range. Based " + f"on your data, we recommend EITHER " + f"desirable_weight in [{des_weight_lower_bound}, {des_weight_upper_bound}] or " + f"undesirable_weight in [{und_weight_lower_bound}, {und_weight_upper_bound}] (but NOT BOTH). " + "See the documentation on how to optimally set these weights.", + UserWarning, + ) + + 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`" + ) + else: + if self.is_deepspeed_enabled: + self.ref_model = prepare_deepspeed(self.ref_model, self.accelerator) + else: + self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True) + + # Import Liger loss if enabled + if self.args.use_liger_loss: + if not is_liger_kernel_available(): + raise ImportError( + "You set `use_liger_loss=True` but the liger kernel is not available. " + "Please install liger-kernel first: `pip install liger-kernel`" + ) + if self.loss_type in ["apo_zero_unpaired"]: + raise ValueError( + "You cannot set `loss_type='apo_zero_unpaired'` with liger-kernel." + "Only KTO loss is supported with liger-kernel." + ) + if self.precompute_ref_log_probs: + raise ValueError( + "You cannot use `precompute_ref_log_probs=True` with liger kernel. Please set " + "`precompute_ref_log_probs=False`." + ) + if self.is_peft_model or self.ref_adapter_name is not None: + raise ValueError( + "You cannot use `use_liger_loss=True` with Peft models. Please set `use_liger_loss=False`." + ) + self.kto_loss_fn = LigerFusedLinearKTOLoss( + ignore_index=self.label_pad_token_id, beta=self.beta, use_ref_model=(self.ref_model is not None) + ) + + @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 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: + dataloader_params = { + "batch_size": self.args.per_device_train_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)) + reference_completion_logps = [] + reference_KL_logps = [] + + for padded_batch in tqdm(iterable=data_loader, desc="Train dataset reference log probs"): + reference_completion_logp, reference_KL_logp = self.compute_reference_log_probs(padded_batch) + + reference_completion_logp = self.accelerator.gather_for_metrics(reference_completion_logp) + reference_completion_logps.append(reference_completion_logp.cpu()) + + if self.calculate_KL: + reference_KL_logp = self.accelerator.gather_for_metrics(reference_KL_logp) + reference_KL_logps.append(reference_KL_logp.cpu()) + + self.train_dataset = self.train_dataset.add_column( + name="reference_logps", column=torch.cat(reference_completion_logps).float().numpy() + ) + + if self.calculate_KL: + self.train_dataset = self.train_dataset.add_column( + name="reference_KL_logps", column=torch.cat(reference_KL_logps).float().numpy() + ) + + 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: + dataloader_params = { + "batch_size": self.args.per_device_eval_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)) + + reference_completion_logps = [] + reference_KL_logps = [] + + for padded_batch in tqdm(iterable=data_loader, desc="Eval dataset reference log probs"): + reference_completion_logp, reference_KL_logp = self.compute_reference_log_probs(padded_batch) + + reference_completion_logp = self.accelerator.gather_for_metrics(reference_completion_logp) + reference_completion_logps.append(reference_completion_logp.cpu()) + + if self.calculate_KL: + reference_KL_logp = self.accelerator.gather_for_metrics(reference_KL_logp) + reference_KL_logps.append(reference_KL_logp.cpu()) + + eval_dataset = eval_dataset.add_column( + name="reference_logps", column=torch.cat(reference_completion_logps).float().numpy() + ) + if self.calculate_KL: + eval_dataset = eval_dataset.add_column( + name="reference_KL_logps", column=torch.cat(reference_KL_logps).float().numpy() + ) + + # Save calculated reference_chosen_logps and reference_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) + + def compute_reference_log_probs(self, padded_batch: dict) -> dict: + """Computes log probabilities of the reference model for a single padded batch of a KTO specific dataset.""" + with torch.no_grad(): + if self.ref_model is None: + with self.null_ref_context(): + if self.is_encoder_decoder: + completion_logits = self.model( + padded_batch["prompt_input_ids"], + attention_mask=padded_batch["prompt_attention_mask"], + decoder_input_ids=padded_batch.get("completion_decoder_input_ids"), + labels=padded_batch["completion_labels"], + ).logits + + if self.calculate_KL: + KL_logits = self.model( + padded_batch["KL_prompt_input_ids"], + attention_mask=padded_batch["KL_prompt_attention_mask"], + decoder_input_ids=padded_batch.get("KL_completion_decoder_input_ids"), + labels=padded_batch["KL_completion_labels"], + ).logits + else: + completion_logits = self.model( + padded_batch["completion_input_ids"], + attention_mask=padded_batch["completion_attention_mask"], + ).logits + + if self.calculate_KL: + KL_logits = self.model( + padded_batch["KL_completion_input_ids"], + attention_mask=padded_batch["KL_completion_attention_mask"], + ).logits + else: + if self.is_encoder_decoder: + completion_logits = self.ref_model( + padded_batch["prompt_input_ids"], + attention_mask=padded_batch["prompt_attention_mask"], + decoder_input_ids=padded_batch.get("completion_decoder_input_ids"), + labels=padded_batch["completion_labels"], + ).logits + + if self.calculate_KL: + KL_logits = self.ref_model( + padded_batch["KL_prompt_input_ids"], + attention_mask=padded_batch["KL_prompt_attention_mask"], + decoder_input_ids=padded_batch.get("KL_completion_decoder_input_ids"), + labels=padded_batch["KL_completion_labels"], + ).logits + else: + completion_logits = self.ref_model( + padded_batch["completion_input_ids"], attention_mask=padded_batch["completion_attention_mask"] + ).logits + + if self.calculate_KL: + KL_logits = self.ref_model( + padded_batch["KL_completion_input_ids"], + attention_mask=padded_batch["KL_completion_attention_mask"], + ).logits + + completion_logps = self.get_batch_logps( + completion_logits, + padded_batch["completion_labels"], + average_log_prob=False, + is_encoder_decoder=self.is_encoder_decoder, + label_pad_token_id=self.label_pad_token_id, + ) + + if self.calculate_KL: + KL_logps = self.get_batch_logps( + KL_logits, + padded_batch["KL_completion_labels"], + average_log_prob=False, + is_encoder_decoder=self.is_encoder_decoder, + label_pad_token_id=self.label_pad_token_id, + ) + else: + KL_logps = None + + return completion_logps, KL_logps + + @staticmethod + def get_batch_logps( + logits: torch.FloatTensor, + labels: torch.LongTensor, + average_log_prob: bool = False, + label_pad_token_id: int = -100, + is_encoder_decoder: bool = False, + ) -> torch.FloatTensor: + """Compute the log probabilities of the given labels under the given logits. + + Args: + logits: + Logits of the model (unnormalized). Shape: (batch_size, sequence_length, vocab_size) + labels: + Labels for which to compute the log probabilities. Label tokens with a value of label_pad_token_id are + ignored. Shape: (batch_size, sequence_length) + average_log_prob: + If True, return the average log probability per (non-masked) token. Otherwise, return the sum of the + log probabilities of the (non-masked) tokens. + + Returns: + A tensor of shape (batch_size,) containing the average/sum log probabilities of the given labels under the + given logits. + """ + if logits.shape[:-1] != labels.shape: + raise ValueError("Logits (batch and sequence length dim) and labels must have the same shape.") + + if not is_encoder_decoder: + labels = labels[:, 1:].clone() + logits = logits[:, :-1, :] + else: + # Fixes end-dec RuntimeError + labels = labels.clone() + + loss_mask = labels != label_pad_token_id + + # dummy token; we'll ignore the losses on these tokens later + labels[labels == label_pad_token_id] = 0 + + per_token_logps = selective_log_softmax(logits, labels) + + if average_log_prob: + return (per_token_logps * loss_mask).sum(-1) / loss_mask.sum(-1) + else: + return (per_token_logps * loss_mask).sum(-1) + + def forward( + self, model: nn.Module, batch: dict[str, Union[list, torch.LongTensor]] + ) -> tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: + KL_logps = self._compute_kl_logps(model, batch) + + model_kwargs = ( + { + "labels": batch["completion_labels"], + "decoder_input_ids": batch.get("completion_decoder_input_ids"), + } + if self.is_encoder_decoder + else {} + ) + if self.aux_loss_enabled: + model_kwargs["output_router_logits"] = True + + outputs = model( + batch["completion_input_ids"], + attention_mask=batch["completion_attention_mask"], + **model_kwargs, + ) + completion_logits = outputs.logits + + completion_logps = self.get_batch_logps( + completion_logits, + batch["completion_labels"], + average_log_prob=False, + is_encoder_decoder=self.is_encoder_decoder, + label_pad_token_id=self.label_pad_token_id, + ) + + if completion_logps.shape[0] != len(batch["label"]): + raise ValueError( + "There is a mismatch between the number of examples in this batch and the number of " + "examples for which an output sequence was predicted." + ) + + chosen_idx = [i for i in range(completion_logps.shape[0]) if batch["label"][i] is True] + rejected_idx = [i for i in range(completion_logps.shape[0]) if batch["label"][i] is False] + + chosen_logps = completion_logps[chosen_idx, ...] + rejected_logps = completion_logps[rejected_idx, ...] + + chosen_logits = completion_logits[chosen_idx, ...] + rejected_logits = completion_logits[rejected_idx, ...] + + if self.aux_loss_enabled: + return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, KL_logps, outputs.aux_loss) + else: + return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, KL_logps) + + def kto_loss( + self, + policy_chosen_logps: torch.FloatTensor, + policy_rejected_logps: torch.FloatTensor, + policy_KL_logps: torch.FloatTensor, + reference_chosen_logps: torch.FloatTensor, + reference_rejected_logps: torch.FloatTensor, + reference_KL_logps: torch.FloatTensor, + ) -> tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: + """Compute the KTO loss for a batch of policy and reference model log probabilities. + + Args: + policy_chosen_logps: + Log probabilities of the policy model for the chosen responses. Shape: (num(chosen) in batch_size,) + policy_rejected_logps: + Log probabilities of the policy model for the rejected responses. Shape: (num(rejected) in batch_size,) + policy_KL_logps: Log probabilities of the policy model for the KL responses. Shape: (batch_size,) + reference_chosen_logps: + Log probabilities of the reference model for the chosen responses. Shape: (num(chosen) in batch_size,) + reference_rejected_logps: + Log probabilities of the reference model for the rejected responses. Shape: (num(rejected) in + batch_size,) + reference_KL_logps: Log probabilities of the reference model for the KL responses. Shape: (batch_size,) + + Returns: + A tuple of four tensors: (losses, chosen_rewards, rejected_rewards, KL). The losses tensor contains the KTO + loss for each example in the batch. The chosen_rewards and rejected_rewards tensors contain the rewards for + the chosen and rejected responses, respectively. The KL tensor contains the detached KL divergence estimate + between the policy and reference models. + """ + if self.calculate_KL: + kl = (policy_KL_logps - reference_KL_logps).mean().detach() + kl = self.accelerator.gather_for_metrics(kl).mean().clamp(min=0) + else: + kl = torch.zeros(1).to(policy_chosen_logps.device) + + # Chosen losses + if policy_chosen_logps.shape[0] != 0 or reference_chosen_logps.shape[0] != 0: + chosen_logratios = policy_chosen_logps - reference_chosen_logps + + if self.loss_type == "kto": + # Eqn (7) of the KTO paper (https://huggingface.co/papers/2402.01306) + chosen_losses = 1 - F.sigmoid(self.beta * (chosen_logratios - kl)) + elif self.loss_type == "apo_zero_unpaired": + # Unpaired variant of 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 + chosen_losses = 1 - F.sigmoid(self.beta * chosen_logratios) + + chosen_rewards = self.beta * chosen_logratios.detach() + + else: + # lists can't be empty -- if they are, then accelerate.gather will hang + chosen_losses = torch.Tensor([]).to(self.accelerator.device) + chosen_rewards = torch.Tensor([]).to(self.accelerator.device) + + # Rejected losses + if policy_rejected_logps.shape[0] != 0 or reference_rejected_logps.shape[0] != 0: + rejected_logratios = policy_rejected_logps - reference_rejected_logps + + if self.loss_type == "kto": + rejected_losses = 1 - F.sigmoid(self.beta * (kl - rejected_logratios)) + elif self.loss_type == "apo_zero_unpaired": + rejected_losses = F.sigmoid(self.beta * rejected_logratios) + + rejected_rewards = self.beta * rejected_logratios.detach() + else: + # lists can't be empty -- if they are, then accelerate.gather will hang + rejected_losses = torch.Tensor([]).to(self.accelerator.device) + rejected_rewards = torch.Tensor([]).to(self.accelerator.device) + + losses = torch.cat( + (self.desirable_weight * chosen_losses, self.undesirable_weight * rejected_losses), + 0, + ) + + return losses, chosen_rewards, rejected_rewards, kl + + def _compute_kl_logps(self, model, batch): + """Compute KL log probabilities for a given batch.""" + KL_logps = None + if self.calculate_KL: + if self.is_encoder_decoder: + KL_model_kwargs = { + "input_ids": batch["KL_prompt_input_ids"], + "attention_mask": batch["KL_prompt_attention_mask"], + "labels": batch["KL_completion_labels"], + "decoder_input_ids": batch.get("KL_completion_decoder_input_ids"), + } + else: + KL_model_kwargs = { + "input_ids": batch["KL_completion_input_ids"], + "attention_mask": batch["KL_completion_attention_mask"], + } + + with torch.no_grad(): + KL_logits = model(**KL_model_kwargs).logits + + KL_logps = self.get_batch_logps( + KL_logits, + batch["KL_completion_labels"], + average_log_prob=False, + is_encoder_decoder=self.is_encoder_decoder, + label_pad_token_id=self.label_pad_token_id, + ) + return KL_logps + + def _compute_loss_liger(self, model, batch): + """ + Compute the KTO loss using the Liger-Kernel's LigerFusedLinearKTOLoss. + + Args: + model: + The policy model used for generating log probabilities and outputs. It could be an encoder-decoder + model or a regular language model. + batch: A dictionary containing the input data and labels for the batch. + + Returns: + A dictionary containing the following keys: + - "loss": The computed KTO loss for the batch. + - "chosen_logits_sum": Sum of the logits for the chosen responses from the policy model. + - "rejected_logits_sum": Sum of the logits for the rejected responses from the policy model. + - "chosen_logps": Log probabilities of the chosen responses from the policy model. + - "rejected_logps": Log probabilities of the rejected responses from the policy model. + - "chosen_rewards": Rewards for the chosen responses. + - "rejected_rewards": Rewards for the rejected responses. + - "kl": The KL divergence between the policy and reference models (detached). + + If auxiliary loss is enabled, the dictionary will also include: + - "aux_loss": The auxiliary loss from the model outputs. + """ + policy_KL_logps = self._compute_kl_logps(model, batch) + reference_KL_logps = self._compute_kl_logps(self.ref_model, batch) + if self.calculate_KL: + kl = (policy_KL_logps - reference_KL_logps).mean().detach() + kl = self.accelerator.gather_for_metrics(kl).mean().clamp(min=0) + else: + kl = torch.zeros(1).to(self.accelerator.device) + + model_kwargs = ( + { + "labels": batch["completion_labels"], + "decoder_input_ids": batch.get("completion_decoder_input_ids"), + } + if self.is_encoder_decoder + else {} + ) + if self.aux_loss_enabled: + model_kwargs["output_router_logits"] = True + + if self.is_encoder_decoder: + # 1. Get encoder outputs + encoder_outputs = model.get_encoder()( + batch["completion_input_ids"], + attention_mask=batch["completion_attention_mask"], + return_dict=True, + **model_kwargs, + ) + # 2. Get decoder outputs + outputs = model.get_decoder()( + input_ids=model_kwargs["decoder_input_ids"], + encoder_hidden_states=encoder_outputs.last_hidden_state, + use_cache=False, + **model_kwargs, + ) + # 1. Get reference encoder outputs + ref_encoder_outputs = self.ref_model.get_encoder()( + batch["completion_input_ids"], + attention_mask=batch["completion_attention_mask"], + return_dict=True, + **model_kwargs, + ) + # 2. Get reference decoder outputs + ref_outputs = self.ref_model.get_decoder()( + input_ids=model_kwargs["decoder_input_ids"], + encoder_hidden_states=ref_encoder_outputs.last_hidden_state, + use_cache=False, + **model_kwargs, + ) + else: + # skip the lm head and get the last hidden state + if hasattr(model, "get_decoder"): + base_model = model.get_decoder() + else: + base_model = getattr(model, self.args.base_model_attribute_name) + outputs = base_model( + batch["completion_input_ids"], + attention_mask=batch["completion_attention_mask"], + use_cache=False, + **model_kwargs, + ) + + # reference model + if hasattr(self.ref_model, "get_decoder"): + ref_base_model = self.ref_model.get_decoder() + else: + ref_base_model = getattr(self.ref_model, self.args.base_model_attribute_name) + ref_outputs = ref_base_model( + batch["completion_input_ids"], + attention_mask=batch["completion_attention_mask"], + use_cache=False, + **model_kwargs, + ) + lm_head = model.get_output_embeddings() + ref_lm_head = self.ref_model.get_output_embeddings() + + ( + loss, + ( + chosen_logps_sum, + rejected_logps_sum, + chosen_logits_sum, + rejected_logits_sum, + chosen_rewards_sum, + rejected_rewards_sum, + ), + ) = self.kto_loss_fn( + _input=outputs.last_hidden_state[:, :-1] if not self.is_encoder_decoder else outputs.last_hidden_state, + lin_weight=lm_head.weight, + target=batch["completion_labels"][:, 1:], + bias=lm_head.bias if hasattr(lm_head, "bias") else None, + preference_labels=torch.tensor(batch["label"], dtype=torch.bool).to(self.accelerator.device), + ref_input=ref_outputs.last_hidden_state[:, :-1] + if not self.is_encoder_decoder + else outputs.last_hidden_state, + ref_weight=ref_lm_head.weight, + ref_bias=ref_lm_head.bias if hasattr(lm_head, "bias") else None, + kl=kl, + ) + + output = { + "loss": loss, + "chosen_logits_sum": chosen_logits_sum, + "rejected_logits_sum": rejected_logits_sum, + "chosen_logps_sum": chosen_logps_sum, + "rejected_logps_sum": rejected_logps_sum, + "chosen_rewards_sum": chosen_rewards_sum, + "rejected_rewards_sum": rejected_rewards_sum, + "kl": kl, + } + 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]], + ): + """Compute the KTO loss and other metrics for the given batch of inputs for train or test.""" + metrics = {} + batch = {k: (v.to(self.accelerator.device) if isinstance(v, torch.Tensor) else v) for k, v in batch.items()} + + labels = torch.tensor(batch["label"]) + num_chosen = labels.sum().to(self.accelerator.device) + num_rejected = (len(labels) - num_chosen).to(self.accelerator.device) + + if self.args.use_liger_loss: + model_output = self._compute_loss_liger(model, batch) + losses = model_output["loss"] + policy_chosen_logits = model_output["chosen_logits_sum"] + policy_rejected_logits = model_output["rejected_logits_sum"] + policy_chosen_logps = model_output["chosen_logps_sum"] + policy_rejected_logps = model_output["rejected_logps_sum"] + chosen_rewards = model_output["chosen_rewards_sum"] + rejected_rewards = model_output["rejected_rewards_sum"] + kl = model_output["kl"] + if self.aux_loss_enabled: + aux_loss = model_output["aux_loss"] + else: + forward_output = self.forward(model, batch) + ( + policy_chosen_logps, + policy_rejected_logps, + policy_chosen_logits, + policy_rejected_logits, + policy_KL_logps, + ) = forward_output[:5] + if self.aux_loss_enabled: + aux_loss = forward_output[5] + + # if reference_logps in batch use them, otherwise use the reference model + if "reference_logps" in batch: + chosen_idx = [i for i in range(batch["reference_logps"].shape[0]) if batch["label"][i] is True] + rejected_idx = [i for i in range(batch["reference_logps"].shape[0]) if batch["label"][i] is False] + + reference_chosen_logps = batch["reference_logps"][chosen_idx, ...] + reference_rejected_logps = batch["reference_logps"][rejected_idx, ...] + if self.calculate_KL: + reference_KL_logps = batch["reference_KL_logps"] + else: + reference_KL_logps = None + else: + with torch.no_grad(): + if self.ref_model is None: + with self.null_ref_context(): + ( + reference_chosen_logps, + reference_rejected_logps, + _, + _, + reference_KL_logps, + ) = self.forward(self.model, batch)[:5] + else: + ( + reference_chosen_logps, + reference_rejected_logps, + _, + _, + reference_KL_logps, + ) = self.forward(self.ref_model, batch)[:5] + + losses, chosen_rewards, rejected_rewards, kl = self.kto_loss( + policy_chosen_logps, + policy_rejected_logps, + policy_KL_logps, + reference_chosen_logps, + reference_rejected_logps, + reference_KL_logps, + ) + + metrics["kl"] = kl.item() + + all_num_chosen = self.accelerator.gather_for_metrics(num_chosen).sum().item() + all_num_rejected = self.accelerator.gather_for_metrics(num_rejected).sum().item() + + if all_num_chosen > 0: + metrics["rewards/chosen_sum"] = ( + self.accelerator.gather_for_metrics(chosen_rewards.nansum()).nansum().item() + ) + metrics["logps/chosen_sum"] = ( + self.accelerator.gather_for_metrics(policy_chosen_logps.nansum()).nansum().item() + ) + metrics["logits/chosen_sum"] = ( + self.accelerator.gather_for_metrics(policy_chosen_logits.nansum()).nansum().item() + ) + metrics["count/chosen"] = all_num_chosen + + if all_num_rejected > 0: + metrics["rewards/rejected_sum"] = ( + self.accelerator.gather_for_metrics(rejected_rewards.nansum()).nansum().item() + ) + metrics["logps/rejected_sum"] = ( + self.accelerator.gather_for_metrics(policy_rejected_logps.nansum()).nansum().item() + ) + metrics["logits/rejected_sum"] = ( + self.accelerator.gather_for_metrics(policy_rejected_logits.nansum()).nansum().item() + ) + metrics["count/rejected"] = all_num_rejected + + loss = losses.nanmean() + if self.aux_loss_enabled: + loss += self.aux_loss_coef * aux_loss + + return loss, 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]]]: + compute_loss_context_manager = ( + autocast(self.accelerator.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) + + # 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 + if self.accelerator.is_main_process: + self.store_metrics(metrics, train_eval="train") + + if return_outputs: + return (loss, metrics) + return loss + + 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 _get_train_sampler(self, dataset: Optional[Dataset] = None) -> Optional[torch.utils.data.Sampler]: + if dataset is None: + dataset = self.train_dataset + if dataset is None or not has_length(dataset): + return None + return SequentialSampler(dataset) + + 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. + generate_context_manager = ( + autocast(self.accelerator.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.processing_class.pad_token_id, + ) + + # if reference_output in batch use that otherwise use the reference model + if "reference_output" in batch: + reference_output = batch["reference_output"] + else: + if self.ref_model is None: + with self.null_ref_context(): + reference_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.processing_class.pad_token_id, + ) + else: + reference_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.processing_class.pad_token_id, + ) + + policy_output = pad_to_length(policy_output, self.max_length, self.processing_class.pad_token_id) + policy_output_decoded = self.processing_class.batch_decode(policy_output, skip_special_tokens=True) + + reference_output = pad_to_length(reference_output, self.max_length, self.processing_class.pad_token_id) + reference_output_decoded = self.processing_class.batch_decode(reference_output, skip_special_tokens=True) + + return policy_output_decoded, reference_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 = [] + + prediction_context_manager = ( + autocast(self.accelerator.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) + + # force log the metrics + if self.accelerator.is_main_process: + 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 = {} + if "logits/chosen_sum" in metrics: + logits_dict["eval_logits/chosen"] = metrics["logits/chosen_sum"] + if "logits/rejected_sum" in metrics: + logits_dict["eval_logits/rejected"] = metrics["logits/rejected_sum"] + 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 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) + + target_indicies = [i for i in range(len(random_batch["label"])) if random_batch["label"][i] is False] + target_batch = { + "prompt_input_ids": random_batch["prompt_input_ids"][target_indicies], + "prompt_attention_mask": random_batch["prompt_attention_mask"][target_indicies], + "prompt": itemgetter(*target_indicies)(random_batch["prompt"]), + } + policy_output_decoded, ref_output_decoded = self.generate_from_model_and_ref(self.model, target_batch) + + table = pd.DataFrame( + columns=["Prompt", "Policy", "Ref Model"], + data=[ + [prompt, pol[len(prompt) :], ref[len(prompt) :]] + for prompt, pol, ref in zip(target_batch["prompt"], policy_output_decoded, ref_output_decoded) + ], + ) + if "wandb" in self.args.report_to: + 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" + # train metrics should have no prefix, eval should have 'eval_' + prefix = "eval_" if train_eval == "eval" else "" + # accumulate average metrics from sums and lengths + for split in ["chosen", "rejected"]: + if f"count/{split}" in self._stored_metrics[train_eval]: + count_sum = torch.Tensor(self._stored_metrics[train_eval][f"count/{split}"]).sum().item() + for metric in ["rewards", "logps", "logits"]: + logs[f"{prefix}{metric}/{split}"] = ( + torch.Tensor(self._stored_metrics[train_eval][f"{metric}/{split}_sum"]).sum().item() + / count_sum + ) + # delete obsolete metric + del self._stored_metrics[train_eval][f"{metric}/{split}_sum"] + del self._stored_metrics[train_eval][f"count/{split}"] + # calculate reward margin + if f"{prefix}rewards/chosen" in logs and f"{prefix}rewards/rejected" in logs: + logs[f"{prefix}rewards/margins"] = logs[f"{prefix}rewards/chosen"] - logs[f"{prefix}rewards/rejected"] + # Add averaged stored metrics to logs + for key, metrics in self._stored_metrics[train_eval].items(): + logs[f"{prefix}{key}"] = torch.Tensor(metrics).mean().item() + del self._stored_metrics[train_eval] + return super().log(logs, start_time) + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) + + 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 + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + citation = textwrap.dedent("""\ + @article{ethayarajh2024kto, + title = {{KTO: Model Alignment as Prospect Theoretic Optimization}}, + author = {Kawin Ethayarajh and Winnie Xu and Niklas Muennighoff and Dan Jurafsky and Douwe Kiela}, + year = 2024, + eprint = {arXiv:2402.01306}, + }""") + + 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="KTO", + trainer_citation=citation, + paper_title="KTO: Model Alignment as Prospect Theoretic Optimization", + paper_id="2402.01306", + ) + + model_card.save(os.path.join(self.args.output_dir, "README.md")) +class UnslothKTOTrainer(_UnslothKTOTrainer): + """ + + Initialize KTOTrainer. + + 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 (`KTOConfig`): + The arguments to use for training. + 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. + data_collator (`transformers.DataCollator`, *optional*, defaults to `None`): + 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. + model_init (`Callable[[], transformers.PreTrainedModel]`): + The model initializer to use for training. If None is specified, the default model initializer will be + used. + 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. + 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. + model_adapter_name (`str`, defaults to `None`): + Name of the train target PEFT adapter, when using LoRA with multiple adapters. + ref_adapter_name (`str`, defaults to `None`): + Name of the reference PEFT adapter, when using LoRA with multiple adapters. + + """ + def __init__( + self, + model = None, + ref_model = None, + args = None, + train_dataset = None, + eval_dataset = None, + processing_class = None, + data_collator = None, + model_init = None, + callbacks = None, + preprocess_logits_for_metrics = None, + peft_config = None, + compute_metrics = None, + model_adapter_name = None, + ref_adapter_name = None, + **kwargs + ): + if args is None: args = UnslothKTOConfig() + 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 + 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) + 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 + 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('kto_trainer', other_metrics) + + super().__init__( + model = model, + ref_model = ref_model, + args = args, + train_dataset = train_dataset, + eval_dataset = eval_dataset, + processing_class = processing_class, + data_collator = data_collator, + model_init = model_init, + callbacks = callbacks, + preprocess_logits_for_metrics = preprocess_logits_for_metrics, + peft_config = peft_config, + compute_metrics = compute_metrics, + model_adapter_name = model_adapter_name, + ref_adapter_name = ref_adapter_name,**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 diff --git a/unsloth_compiled_cache/UnslothNashMDTrainer.py b/unsloth_compiled_cache/UnslothNashMDTrainer.py new file mode 100644 index 0000000..5220a70 --- /dev/null +++ b/unsloth_compiled_cache/UnslothNashMDTrainer.py @@ -0,0 +1,993 @@ +""" +2025.7.1 +2025.7.1 +4.53.1 +0.19.1 +__UNSLOTH_VERSIONING__ +""" +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +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, generate_model_card, get_comet_experiment_url, get_reward, is_conversational, is_peft_available, is_wandb_available, jinja2, maybe_apply_chat_template, nn, os, 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 +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 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.'}, + ) + 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, + hub_revision = None, + 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, + liger_kernel_config = None, + eval_use_gather_object = False, + average_tokens_across_devices = False, + reward_model_path = None, + judge = None, + max_new_tokens = 64, + max_length = 512, + temperature = 0.9, + missing_eos_penalty = None, + loss_type = 'sigmoid', + dataset_num_proc = None, + disable_dropout = True, + use_vllm = False, + gpu_memory_utilization = 0.55, + ds3_gather_for_generation = True, + 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() + if temperature <= 0: + raise MathError('Unsloth: Please set a positive non-zero temperature since your results will be wrong.') + elif temperature >= 10: + raise MathError('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, + 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, + 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, + missing_eos_penalty = missing_eos_penalty, + loss_type = loss_type, + dataset_num_proc = dataset_num_proc, + disable_dropout = disable_dropout, + use_vllm = use_vllm, + gpu_memory_utilization = gpu_memory_utilization, + ds3_gather_for_generation = ds3_gather_for_generation,**kwargs) + self.vllm_sampling_params = vllm_sampling_params + self.unsloth_num_chunks = unsloth_num_chunks +pass + +class _UnslothNashMDTrainer(OnlineDPOTrainer): + r"""""" + + _tag_names = ["trl", "nash-md"] + + def __init__( + self, + model: Union[PreTrainedModel, nn.Module] = None, + ref_model: Union[PreTrainedModel, nn.Module] = None, + reward_model: 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, + ) -> None: + super().__init__( + model=model, + ref_model=ref_model, + reward_model=reward_model, + judge=judge, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + reward_processing_class=processing_class, # for now, NashMDTrainer can't use any reward model + peft_config=peft_config, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + ) + + 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_model is not None: + 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_model, model_data["input_ids"], self.processing_class.pad_token_id, context_length + ) + _, mixture_scores, _ = get_reward( + self.reward_model, 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_model 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_model 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 + + if self.use_apex: + with amp.scale_loss(loss, self.optimizer) as scaled_loss: + scaled_loss.backward() + else: + self.accelerator.backward(loss, **kwargs) + + return loss.detach() / self.args.gradient_accumulation_steps + + 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 + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + 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} + }""") + + 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="Nash-MD", + trainer_citation=citation, + paper_title="Nash Learning from Human Feedback", + paper_id="2312.00886", + ) + + model_card.save(os.path.join(self.args.output_dir, "README.md")) +class UnslothNashMDTrainer(_UnslothNashMDTrainer): + """ + + Initialize NashMDTrainer as a subclass of [`OnlineDPOConfig`]. + + 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_model (`transformers.PreTrainedModel`): + The reward model to score completions with, preferably an `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 (`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. + 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. + + """ + def __init__( + self, + model = None, + ref_model = None, + reward_model = 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, + **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 + 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) + 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 + 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('nash_md_trainer', other_metrics) + + super().__init__( + model = model, + ref_model = ref_model, + reward_model = reward_model, + 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,**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 diff --git a/unsloth_compiled_cache/UnslothORPOTrainer.py b/unsloth_compiled_cache/UnslothORPOTrainer.py new file mode 100644 index 0000000..1aef3e9 --- /dev/null +++ b/unsloth_compiled_cache/UnslothORPOTrainer.py @@ -0,0 +1,1547 @@ +""" +2025.7.1 +2025.7.1 +4.53.1 +0.19.1 +__UNSLOTH_VERSIONING__ +""" +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from trl.trainer.orpo_trainer import (Any, AutoModelForCausalLM, BaseImageProcessor, Callable, DPODataCollatorWithPadding, DataCollator, DataLoader, Dataset, EvalLoopOutput, F, FeatureExtractionMixin, Literal, ORPOConfig, ORPOTrainer, Optional, PartialState, Path, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, Trainer, TrainerCallback, Union, add_bos_token_if_needed, add_eos_token_if_needed, autocast, defaultdict, disable_dropout_in_model, generate_model_card, get_comet_experiment_url, inspect, is_comet_available, is_peft_available, is_torch_fx_proxy, is_torch_xla_available, is_wandb_available, log_table_to_comet_experiment, maybe_apply_chat_template, maybe_extract_prompt, nn, np, nullcontext, os, pad_to_length, pd, peft_module_casting_to_bf16, prepare_model_for_kbit_training, random, textwrap, torch, warnings, F, Optional, PeftModel, PreTrainedModel, Trainer, is_peft_available, os, torch) + + +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 UnslothORPOConfig(ORPOConfig): + """ + + Configuration class for the [`ORPOTrainer`]. + + This class includes only the parameters that are specific to ORPO training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may + differ from those in [`~transformers.TrainingArguments`]. + + 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: + max_length (`int` or `None`, *optional*, defaults to `1024`): + Maximum length of the sequences (prompt + completion) in the batch. This argument is required if you want + to use the default data collator. + max_prompt_length (`int` or `None`, *optional*, defaults to `512`): + Maximum length of the prompt. This argument is required if you want to use the default data collator. + max_completion_length (`int` or `None`, *optional*, defaults to `None`): + Maximum length of the completion. This argument is required if you want to use the default data collator + and your model is an encoder-decoder. + beta (`float`, *optional*, defaults to `0.1`): + Parameter controlling the relative ratio loss weight in the ORPO loss. In the + [paper](https://huggingface.co/papers/2403.07691), it is denoted by λ. In the + [code](https://github.com/xfactlab/orpo), it is denoted by `alpha`. + disable_dropout (`bool`, *optional*, defaults to `True`): + Whether to disable dropout in the model. + label_pad_token_id (`int`, *optional*, defaults to `-100`): + Label pad token id. This argument is required if you want to use the default data collator. + padding_value (`int` or `None`, *optional*, defaults to `None`): + Padding value to use. If `None`, the padding value of the tokenizer is used. + truncation_mode (`str`, *optional*, defaults to `"keep_end"`): + Truncation mode to use when the prompt is too long. Possible values are `"keep_end"` or `"keep_start"`. + This argument is required if you want to use the default data collator. + generate_during_eval (`bool`, *optional*, defaults to `False`): + If `True`, generates and logs completions from the model to W&B or Comet during evaluation. + is_encoder_decoder (`bool` or `None`, *optional*, defaults to `None`): + When using the `model_init` argument (callable) to instantiate the model instead of the `model` argument, + you need to specify if the model returned by the callable is an encoder-decoder model. + model_init_kwargs (`dict[str, Any]` or `None`, *optional*, defaults to `None`): + Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model from a + string. + dataset_num_proc (`int` or `None`, *optional*, defaults to `None`): + Number of processes to use for processing the dataset. + + """ + 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, + hub_revision = None, + 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, + liger_kernel_config = None, + eval_use_gather_object = False, + average_tokens_across_devices = False, + max_length = 1024, + max_prompt_length = 512, + max_completion_length = None, + beta = 0.1, + disable_dropout = True, + label_pad_token_id = -100, + padding_value = None, + truncation_mode = 'keep_end', + generate_during_eval = False, + is_encoder_decoder = None, + model_init_kwargs = None, + dataset_num_proc = None, + 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, + 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, + max_length = max_length, + max_prompt_length = max_prompt_length, + max_completion_length = max_completion_length, + beta = beta, + disable_dropout = disable_dropout, + label_pad_token_id = label_pad_token_id, + padding_value = padding_value, + truncation_mode = truncation_mode, + generate_during_eval = generate_during_eval, + is_encoder_decoder = is_encoder_decoder, + model_init_kwargs = model_init_kwargs, + dataset_num_proc = dataset_num_proc,**kwargs) + self.vllm_sampling_params = vllm_sampling_params + self.unsloth_num_chunks = unsloth_num_chunks +pass + +class _UnslothORPOTrainer(Trainer): + r"""""" + + _tag_names = ["trl", "orpo"] + + def __init__( + self, + model: Optional[Union[PreTrainedModel, nn.Module, str]] = None, + args: Optional[ORPOConfig] = 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, + 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, + compute_metrics: Optional[Callable[[EvalLoopOutput], dict]] = None, + ): + if args.model_init_kwargs is None: + model_init_kwargs = {} + elif not isinstance(model, str): + raise ValueError("You passed model_kwargs to the ORPOTrainer. 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 ORPOConfig. Expected a string with either `torch.dtype` or 'auto', but got {torch_dtype}." + ) + model_init_kwargs["torch_dtype"] = torch_dtype + + if isinstance(model, str): + model = AutoModelForCausalLM.from_pretrained(model, **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 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 args.gradient_checkpointing: + # 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 args.gradient_checkpointing: + # 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." + ) + + if model is not None: + self.is_encoder_decoder = model.config.is_encoder_decoder + elif args.is_encoder_decoder is None: + raise ValueError("When no model is provided, you need to pass the parameter is_encoder_decoder.") + else: + self.is_encoder_decoder = args.is_encoder_decoder + + if self.is_encoder_decoder: + self.decoder_start_token_id = model.config.decoder_start_token_id + self.pad_token_id = model.config.pad_token_id + + if processing_class is None: + raise ValueError("processing_class must be specified to tokenize a ORPO dataset.") + if args.max_length is None: + warnings.warn( + "`max_length` is not set in the ORPOConfig's init" + " it will default to `512` by default, but you should do it yourself in the future.", + UserWarning, + ) + max_length = 512 + else: + max_length = args.max_length + if args.max_prompt_length is None: + warnings.warn( + "`max_prompt_length` is not set in the ORPOConfig's init" + " it will default to `128` by default, but you should do it yourself in the future.", + UserWarning, + ) + max_prompt_length = 128 + else: + max_prompt_length = args.max_prompt_length + + if args.max_completion_length is None and self.is_encoder_decoder: + warnings.warn( + "When using an encoder decoder architecture, you should set `max_completion_length` in the ORPOConfig's init" + " it will default to `128` by default, but you should do it yourself in the future.", + UserWarning, + ) + self.max_completion_length = 128 + else: + self.max_completion_length = args.max_completion_length + + if data_collator is None: + data_collator = DPODataCollatorWithPadding( + pad_token_id=processing_class.pad_token_id, + label_pad_token_id=args.label_pad_token_id, + is_encoder_decoder=self.is_encoder_decoder, + ) + + if args.remove_unused_columns: + args.remove_unused_columns = False + # warn users + warnings.warn( + "When using DPODataCollatorWithPadding, you should set `remove_unused_columns=False` in your TrainingArguments" + " we have set it for you, but you should do it yourself in the future.", + UserWarning, + ) + + self.use_dpo_data_collator = True + else: + self.use_dpo_data_collator = False + + # Disable dropout in the model and reference model + if args.disable_dropout: + disable_dropout_in_model(model) + + self.max_length = max_length + self.generate_during_eval = args.generate_during_eval + self.label_pad_token_id = args.label_pad_token_id + self.padding_value = args.padding_value if args.padding_value is not None else processing_class.pad_token_id + self.max_prompt_length = max_prompt_length + self.truncation_mode = args.truncation_mode + self.processing_class = processing_class + + self.beta = args.beta + self.aux_loss_enabled = getattr(model.config, "output_router_logits", False) + 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)) + + # 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 ORPO, 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 + + # Compute that only on the main process for faster data processing. + # see: https://github.com/huggingface/trl/pull/1255 + with PartialState().main_process_first(): + # Extract the prompt if needed, and apply the chat template if needed + train_dataset = train_dataset.map(maybe_extract_prompt, num_proc=args.dataset_num_proc) + train_dataset = train_dataset.map( + maybe_apply_chat_template, fn_kwargs={"tokenizer": processing_class}, num_proc=args.dataset_num_proc + ) + train_dataset = train_dataset.map(self.tokenize_row, num_proc=args.dataset_num_proc) + if eval_dataset is not None: + eval_dataset = eval_dataset.map(maybe_extract_prompt, num_proc=args.dataset_num_proc) + eval_dataset = eval_dataset.map( + maybe_apply_chat_template, + fn_kwargs={"tokenizer": processing_class}, + num_proc=args.dataset_num_proc, + ) + eval_dataset = eval_dataset.map(self.tokenize_row, num_proc=args.dataset_num_proc) + + 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, + ) + + # 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`." + ) + + def build_tokenized_answer(self, prompt, answer): + """ + Llama tokenizer does satisfy `enc(a + b) = enc(a) + enc(b)`. It does ensure `enc(a + b) = enc(a) + enc(a + + b)[len(enc(a)):]`. Reference: + https://github.com/EleutherAI/lm-evaluation-harness/pull/531#issuecomment-1595586257 + """ + + full_tokenized = self.processing_class(prompt + answer, add_special_tokens=False) + prompt_input_ids = self.processing_class(prompt, add_special_tokens=False)["input_ids"] + + answer_input_ids = full_tokenized["input_ids"][len(prompt_input_ids) :] + answer_attention_mask = full_tokenized["attention_mask"][len(prompt_input_ids) :] + + # Concat tokens to form `enc(a) + enc(a + b)[len(enc(a)):]` + full_concat_input_ids = np.concatenate([prompt_input_ids, answer_input_ids]) + + # Prepare input tokens for token by token comparison + full_input_ids = np.array(full_tokenized["input_ids"]) + + if len(full_input_ids) != len(full_concat_input_ids): + raise ValueError("Prompt input ids and answer input ids should have the same length.") + + # On some tokenizers, like Llama-2 tokenizer, there are occasions where tokens + # can be merged together when tokenizing prompt+answer. This could result + # on the last token from the prompt being different when tokenized on its own + # vs when done as prompt+answer. + response_token_ids_start_idx = len(prompt_input_ids) + + # If tokenized prompt is different than both prompt+answer, then it means the + # last token has changed due to merging. + if prompt_input_ids != full_tokenized["input_ids"][:response_token_ids_start_idx]: + response_token_ids_start_idx -= 1 + + prompt_input_ids = full_tokenized["input_ids"][:response_token_ids_start_idx] + prompt_attention_mask = full_tokenized["attention_mask"][:response_token_ids_start_idx] + + if len(prompt_input_ids) != len(prompt_attention_mask): + raise ValueError("Prompt input ids and attention mask should have the same length.") + + answer_input_ids = full_tokenized["input_ids"][response_token_ids_start_idx:] + answer_attention_mask = full_tokenized["attention_mask"][response_token_ids_start_idx:] + + return dict( + prompt_input_ids=prompt_input_ids, + prompt_attention_mask=prompt_attention_mask, + input_ids=answer_input_ids, + attention_mask=answer_attention_mask, + ) + + def tokenize_row(self, feature, model: Optional[Union[PreTrainedModel, nn.Module]] = None) -> dict: + """Tokenize a single row from a ORPO specific dataset. + + At this stage, we don't convert to PyTorch tensors yet; we just handle the truncation in case the prompt + + chosen or prompt + rejected responses is/are too long. First we truncate the prompt; if we're still too long, + we truncate the chosen/rejected. + + We also create the labels for the chosen/rejected responses, which are of length equal to the sum of the length + of the prompt and the chosen/rejected response, with label_pad_token_id for the prompt tokens. + """ + batch = {} + prompt = feature["prompt"] + chosen = feature["chosen"] + rejected = feature["rejected"] + + if not self.is_encoder_decoder: + # Check issues below for more details + # 1. https://github.com/huggingface/trl/issues/907 + # 2. https://github.com/EleutherAI/lm-evaluation-harness/pull/531#issuecomment-1595586257 + # 3. https://github.com/LianjiaTech/BELLE/issues/337 + + if not isinstance(prompt, str): + raise ValueError(f"prompt should be an str but got {type(prompt)}") + prompt_tokens = self.processing_class(prompt, add_special_tokens=False) + prompt_tokens = {f"prompt_{k}": v for k, v in prompt_tokens.items()} + + if not isinstance(chosen, str): + raise ValueError(f"chosen should be an str but got {type(chosen)}") + chosen_tokens = self.build_tokenized_answer(prompt, chosen) + + if not isinstance(rejected, str): + raise ValueError(f"rejected should be an str but got {type(rejected)}") + rejected_tokens = self.build_tokenized_answer(prompt, rejected) + + # Last prompt token might get merged by tokenizer and + # it should not be included for generation if that happens + prompt_len_input_ids = len(prompt_tokens["prompt_input_ids"]) + + chosen_prompt_len_input_ids = len(chosen_tokens["prompt_input_ids"]) + rejected_prompt_len_input_ids = len(rejected_tokens["prompt_input_ids"]) + prompt_len_input_ids = min(chosen_prompt_len_input_ids, rejected_prompt_len_input_ids) + + for k, v in prompt_tokens.items(): + prompt_tokens[k] = v[:prompt_len_input_ids] + + # Make sure prompts only have one different token at most an + # and length only differs by 1 at most + num_diff_tokens = sum( + [a != b for a, b in zip(chosen_tokens["prompt_input_ids"], rejected_tokens["prompt_input_ids"])] + ) + num_diff_len = abs(chosen_prompt_len_input_ids - rejected_prompt_len_input_ids) + if num_diff_tokens > 1 or num_diff_len > 1: + raise ValueError( + "Chosen and rejected prompt_input_ids might only differ on the " + "last token due to tokenizer merge ops." + ) + + # add BOS token to head of prompt. Avoid adding if it's already there + prompt_tokens, chosen_tokens, rejected_tokens = add_bos_token_if_needed( + self.processing_class.bos_token_id, + prompt_len_input_ids, + prompt_tokens, + chosen_prompt_len_input_ids, + chosen_tokens, + rejected_prompt_len_input_ids, + rejected_tokens, + ) + + # add EOS token to end of answer. Avoid adding if it's already there + chosen_tokens, rejected_tokens = add_eos_token_if_needed( + self.processing_class.eos_token_id, chosen_tokens, rejected_tokens + ) + + longer_response_length = max(len(chosen_tokens["input_ids"]), len(rejected_tokens["input_ids"])) + + # if combined sequence is too long, truncate the prompt + for answer_tokens in [chosen_tokens, rejected_tokens, prompt_tokens]: + if len(answer_tokens["prompt_input_ids"]) + longer_response_length > self.max_length: + if self.truncation_mode == "keep_start": + for k in ["prompt_input_ids", "prompt_attention_mask"]: + answer_tokens[k] = answer_tokens[k][: self.max_prompt_length] + elif self.truncation_mode == "keep_end": + for k in ["prompt_input_ids", "prompt_attention_mask"]: + answer_tokens[k] = answer_tokens[k][-self.max_prompt_length :] + else: + raise ValueError(f"Unknown truncation mode: {self.truncation_mode}") + + # if that's still too long, truncate the response + for answer_tokens in [chosen_tokens, rejected_tokens]: + if len(answer_tokens["prompt_input_ids"]) + longer_response_length > self.max_length: + for k in ["input_ids", "attention_mask"]: + answer_tokens[k] = answer_tokens[k][: self.max_length - self.max_prompt_length] + + # Create labels + chosen_sequence_tokens = { + k: chosen_tokens[f"prompt_{k}"] + chosen_tokens[k] for k in ["input_ids", "attention_mask"] + } + rejected_sequence_tokens = { + k: rejected_tokens[f"prompt_{k}"] + rejected_tokens[k] for k in ["input_ids", "attention_mask"] + } + chosen_sequence_tokens["labels"] = chosen_sequence_tokens["input_ids"][:] + chosen_sequence_tokens["labels"][: len(chosen_tokens["prompt_input_ids"])] = [ + self.label_pad_token_id + ] * len(chosen_tokens["prompt_input_ids"]) + rejected_sequence_tokens["labels"] = rejected_sequence_tokens["input_ids"][:] + rejected_sequence_tokens["labels"][: len(rejected_tokens["prompt_input_ids"])] = [ + self.label_pad_token_id + ] * len(rejected_tokens["prompt_input_ids"]) + + for k, toks in { + "chosen_": chosen_sequence_tokens, + "rejected_": rejected_sequence_tokens, + "": prompt_tokens, + }.items(): + for type_key, tokens in toks.items(): + if type_key == "token_type_ids": + continue + batch[f"{k}{type_key}"] = tokens + + else: + chosen_tokens = self.processing_class( + chosen, truncation=True, max_length=self.max_completion_length, add_special_tokens=True + ) + rejected_tokens = self.processing_class( + rejected, truncation=True, max_length=self.max_completion_length, add_special_tokens=True + ) + prompt_tokens = self.processing_class( + prompt, truncation=True, max_length=self.max_prompt_length, add_special_tokens=True + ) + + batch["chosen_labels"] = chosen_tokens["input_ids"] + batch["rejected_labels"] = rejected_tokens["input_ids"] + batch["prompt_input_ids"] = prompt_tokens["input_ids"] + batch["prompt_attention_mask"] = prompt_tokens["attention_mask"] + + if model is not None and hasattr(model, "prepare_decoder_input_ids_from_labels"): + batch["rejected_decoder_input_ids"] = model.prepare_decoder_input_ids_from_labels( + labels=torch.tensor(batch["rejected_labels"]) + ) + batch["chosen_decoder_input_ids"] = model.prepare_decoder_input_ids_from_labels( + labels=torch.tensor(batch["chosen_labels"]) + ) + + if is_torch_xla_available(): + # Pad the sequences to global max_length to avoid TorchXLA recompilation + for k in batch: + if "labels" in k or self.is_encoder_decoder: + pad_value = self.label_pad_token_id + elif k.endswith("_input_ids"): + pad_value = self.padding_value + elif k.endswith("_attention_mask"): + pad_value = 0 + batch[k] = batch[k] + [pad_value] * (self.max_length - len(batch[k])) + return batch + + @staticmethod + def concatenated_inputs( + batch: dict[str, Union[list, torch.LongTensor]], + is_encoder_decoder: bool = False, + label_pad_token_id: int = -100, + padding_value: int = 0, + device: Optional[torch.device] = None, + ) -> dict[str, torch.LongTensor]: + """Concatenate the chosen and rejected inputs into a single tensor. + + Args: + batch: + A batch of data. Must contain the keys 'chosen_input_ids' and 'rejected_input_ids', which are tensors + of shape (batch_size, sequence_length). + is_encoder_decoder: + Whether the model is an encoder-decoder model. + label_pad_token_id: + The label pad token id. + padding_value: + The padding value to use for the concatenated inputs_ids. + device: + The device for the concatenated inputs. + + Returns: + A dictionary containing the concatenated inputs under the key 'concatenated_input_ids'. + """ + concatenated_batch = {} + + if is_encoder_decoder: + max_length = max(batch["chosen_labels"].shape[1], batch["rejected_labels"].shape[1]) + else: + max_length = max(batch["chosen_input_ids"].shape[1], batch["rejected_input_ids"].shape[1]) + + for k in batch: + if k.startswith("chosen") and isinstance(batch[k], torch.Tensor): + if "labels" in k or is_encoder_decoder: + pad_value = label_pad_token_id + elif k.endswith("_input_ids"): + pad_value = padding_value + elif k.endswith("_attention_mask"): + pad_value = 0 + concatenated_key = k.replace("chosen", "concatenated") + concatenated_batch[concatenated_key] = pad_to_length(batch[k], max_length, pad_value=pad_value) + for k in batch: + if k.startswith("rejected") and isinstance(batch[k], torch.Tensor): + if "labels" in k or is_encoder_decoder: + pad_value = label_pad_token_id + elif k.endswith("_input_ids"): + pad_value = padding_value + elif k.endswith("_attention_mask"): + pad_value = 0 + concatenated_key = k.replace("rejected", "concatenated") + concatenated_batch[concatenated_key] = torch.cat( + ( + concatenated_batch[concatenated_key], + pad_to_length(batch[k], max_length, pad_value=pad_value), + ), + dim=0, + ).to(device=device) + + if is_encoder_decoder: + concatenated_batch["concatenated_input_ids"] = batch["prompt_input_ids"].repeat(2, 1).to(device=device) + concatenated_batch["concatenated_attention_mask"] = ( + batch["prompt_attention_mask"].repeat(2, 1).to(device=device) + ) + + return concatenated_batch + + def odds_ratio_loss( + self, + policy_chosen_logps: torch.FloatTensor, + policy_rejected_logps: torch.FloatTensor, + ) -> tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: + """Compute ORPO's odds ratio (OR) loss for a batch of policy and reference model log probabilities. + + Args: + policy_chosen_logps: + Log probabilities of the policy model for the chosen responses. Shape: (batch_size,) + policy_rejected_logps: + Log probabilities of the policy model for the rejected responses. Shape: (batch_size,) + + Returns: + A tuple of three tensors: (losses, chosen_rewards, rejected_rewards). The losses tensor contains the ORPO + loss for each example in the batch. The chosen_rewards and rejected_rewards tensors contain the rewards for + the chosen and rejected responses, respectively. The log odds ratio of the chosen responses over the + rejected responses ratio for logging purposes. The `log(sigmoid(log_odds_chosen))` for logging purposes. + """ + + # Derived from Eqs. (4) and (7) from https://huggingface.co/papers/2403.07691 by using log identities and exp(log(P(y|x)) = P(y|x) + log_odds = (policy_chosen_logps - policy_rejected_logps) - ( + torch.log1p(-torch.exp(policy_chosen_logps)) - torch.log1p(-torch.exp(policy_rejected_logps)) + ) + ratio = F.logsigmoid(log_odds) + losses = self.beta * ratio + + chosen_rewards = self.beta * (policy_chosen_logps.to(self.accelerator.device)).detach() + rejected_rewards = self.beta * (policy_rejected_logps.to(self.accelerator.device)).detach() + + return losses, chosen_rewards, rejected_rewards, torch.mean(ratio), torch.mean(log_odds) + + @staticmethod + def get_batch_logps( + logits: torch.FloatTensor, + labels: torch.LongTensor, + average_log_prob: bool = False, + label_pad_token_id: int = -100, + is_encoder_decoder: bool = False, + ) -> torch.FloatTensor: + """Compute the log probabilities of the given labels under the given logits. + + Args: + logits: Logits of the model (unnormalized). Shape: (batch_size, sequence_length, vocab_size) + labels: + Labels for which to compute the log probabilities. Label tokens with a value of label_pad_token_id are + ignored. Shape: (batch_size, sequence_length) + average_log_prob: + If True, return the average log probability per (non-masked) token. Otherwise, return the sum of the + log probabilities of the (non-masked) tokens. + label_pad_token_id: The label pad token id. + is_encoder_decoder: Whether the model is an encoder-decoder model. + + Returns: + A tensor of shape (batch_size,) containing the average/sum log probabilities of the given labels under the + given logits. + """ + if logits.shape[:-1] != labels.shape: + raise ValueError("Logits (batch and sequence length dim) and labels must have the same shape.") + + if not is_encoder_decoder: + labels = labels[:, 1:].clone() + logits = logits[:, :-1, :] + loss_mask = labels != label_pad_token_id + + # dummy token; we'll ignore the losses on these tokens later + labels = torch.where(labels == label_pad_token_id, 0, labels) + + per_token_logps = selective_log_softmax(logits, labels) + + if average_log_prob: + return (per_token_logps * loss_mask).sum(-1) / loss_mask.sum(-1) + else: + return (per_token_logps * loss_mask).sum(-1) + + def concatenated_forward( + self, model: nn.Module, batch: dict[str, Union[list, torch.LongTensor]] + ) -> tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: + """Run 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. + """ + concatenated_batch = self.concatenated_inputs( + batch, + is_encoder_decoder=self.is_encoder_decoder, + label_pad_token_id=self.label_pad_token_id, + padding_value=self.padding_value, + device=self.accelerator.device, + ) + len_chosen = batch["chosen_labels"].shape[0] + + model_kwargs = ( + { + "decoder_input_ids": self._shift_right(concatenated_batch["concatenated_labels"]), + } + if self.is_encoder_decoder + else {} + ) + + if self.aux_loss_enabled: + model_kwargs["output_router_logits"] = True + + outputs = model( + concatenated_batch["concatenated_input_ids"], + attention_mask=concatenated_batch["concatenated_attention_mask"], + use_cache=False, + **model_kwargs, + ) + all_logits = outputs.logits + + def cross_entropy_loss(logits, labels): + if not self.is_encoder_decoder: + # Shift so that tokens < n predict n + logits = logits[..., :-1, :].contiguous() + labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = nn.CrossEntropyLoss() + logits = logits.view(-1, logits.shape[-1]) + labels = labels.view(-1) + # Enable model parallelism + labels = labels.to(logits.device) + loss = loss_fct(logits, labels) + return loss + + if self.is_encoder_decoder: + labels = concatenated_batch["concatenated_labels"].clone() + else: + labels = concatenated_batch["concatenated_input_ids"].clone() + attention_mask = concatenated_batch["concatenated_attention_mask"] + labels = torch.where(attention_mask == 1, labels, self.label_pad_token_id) + # orpo chosen nll loss is computed over the full prompt and response + chosen_nll_loss = cross_entropy_loss(all_logits[:len_chosen], labels[:len_chosen]) + + all_logps = self.get_batch_logps( + all_logits, + concatenated_batch["concatenated_labels"], + average_log_prob=True, + is_encoder_decoder=self.is_encoder_decoder, + label_pad_token_id=self.label_pad_token_id, + ) + + chosen_logps = all_logps[:len_chosen] + rejected_logps = all_logps[len_chosen:] + + if not self.is_encoder_decoder: + chosen_logits = all_logits[:len_chosen, :-1, :] + rejected_logits = all_logits[len_chosen:, :-1, :] + else: + chosen_logits = all_logits[:len_chosen] + rejected_logits = all_logits[len_chosen:] + + if self.aux_loss_enabled: + return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, chosen_nll_loss, outputs.aux_loss) + + return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, chosen_nll_loss) + + def get_batch_loss_metrics( + self, + model, + batch: dict[str, Union[list, torch.LongTensor]], + train_eval: Literal["train", "eval"] = "train", + ): + """Compute the ORPO loss and other metrics for the given batch of inputs for train or test.""" + metrics = {} + + forward_output = self.concatenated_forward(model, batch) + ( + policy_chosen_logps, + policy_rejected_logps, + policy_chosen_logits, + policy_rejected_logits, + policy_nll_loss, + ) = forward_output[:5] + if self.aux_loss_enabled: + aux_loss = forward_output[5] + + losses, chosen_rewards, rejected_rewards, log_odds_ratio, log_odds_chosen = self.odds_ratio_loss( + policy_chosen_logps, policy_rejected_logps + ) + # full ORPO loss + loss = policy_nll_loss - losses.mean() + + reward_accuracies = (chosen_rewards > rejected_rewards).float() + + prefix = "eval_" if train_eval == "eval" else "" + metrics[f"{prefix}rewards/chosen"] = self.accelerator.gather_for_metrics(chosen_rewards).mean() + metrics[f"{prefix}rewards/rejected"] = self.accelerator.gather_for_metrics(rejected_rewards).mean() + metrics[f"{prefix}rewards/accuracies"] = self.accelerator.gather_for_metrics(reward_accuracies).mean() + metrics[f"{prefix}rewards/margins"] = self.accelerator.gather_for_metrics( + chosen_rewards - rejected_rewards + ).mean() + metrics[f"{prefix}logps/rejected"] = self.accelerator.gather_for_metrics(policy_rejected_logps).detach().mean() + metrics[f"{prefix}logps/chosen"] = self.accelerator.gather_for_metrics(policy_chosen_logps).detach().mean() + metrics[f"{prefix}logits/rejected"] = self.accelerator.gather_for_metrics( + policy_rejected_logits.detach().mean() + ).mean() + metrics[f"{prefix}logits/chosen"] = self.accelerator.gather_for_metrics( + policy_chosen_logits.detach().mean() + ).mean() + metrics[f"{prefix}nll_loss"] = self.accelerator.gather_for_metrics(policy_nll_loss).detach().mean() + metrics[f"{prefix}log_odds_ratio"] = self.accelerator.gather_for_metrics(log_odds_ratio).detach().mean() + metrics[f"{prefix}log_odds_chosen"] = self.accelerator.gather_for_metrics(log_odds_chosen).detach().mean() + if is_torch_xla_available(): + xm.mark_step() # needed because .item() calls + for k, v in metrics.items(): + metrics[k] = v.item() + if self.aux_loss_enabled: + loss += self.aux_loss_coef * aux_loss + + return loss, 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]]]: + compute_loss_context_manager = ( + autocast(self.accelerator.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(self, model, batch: dict[str, torch.LongTensor]) -> 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. + generate_context_manager = ( + autocast(self.accelerator.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.processing_class.pad_token_id, + ) + + policy_output = pad_to_length(policy_output, self.max_length, self.processing_class.pad_token_id) + policy_output_decoded = self.processing_class.batch_decode(policy_output, skip_special_tokens=True) + + return policy_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 not self.use_dpo_data_collator: + warnings.warn( + "prediction_step is only implemented for DPODataCollatorWithPadding, and you passed a datacollator that is different than " + "DPODataCollatorWithPadding - you might see unexpected behavior. Alternatively, you can implement your own prediction_step method if you are using a custom data collator" + ) + if ignore_keys is None: + if hasattr(model, "config"): + ignore_keys = getattr(model.config, "keys_to_ignore_at_inference", []) + else: + ignore_keys = [] + + prediction_context_manager = ( + autocast(self.accelerator.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 = self.generate_from_model(self.model, random_batch) + + table = pd.DataFrame( + columns=["Prompt", "Policy"], + data=[ + [prompt, pol[len(prompt) :]] for prompt, pol in zip(random_batch["prompt"], policy_output_decoded) + ], + ) + if "wandb" in self.args.report_to: + 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] + return super().log(logs, start_time) + + def _shift_right(self, input_ids): + if self.decoder_start_token_id is None: + raise ValueError( + "model.config.decoder_start_token_id has to be defined. It is usually set to the pad_token_id." + ) + + # shift inputs to the right + if is_torch_fx_proxy(input_ids): + # Item assignment is not supported natively for proxies. + shifted_input_ids = torch.full(input_ids.shape[:-1] + (1,), self.decoder_start_token_id) + shifted_input_ids = torch.cat([shifted_input_ids, input_ids[..., :-1]], dim=-1) + else: + shifted_input_ids = input_ids.new_zeros(input_ids.shape) + shifted_input_ids[..., 1:] = input_ids[..., :-1].clone() + shifted_input_ids[..., 0] = self.decoder_start_token_id + + if self.pad_token_id is None: + raise ValueError("model.config.pad_token_id has to be defined.") + # replace possible -100 values in labels by `pad_token_id` + shifted_input_ids.masked_fill_(shifted_input_ids == -100, self.pad_token_id) + + return shifted_input_ids + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) + + 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 + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + citation = textwrap.dedent("""\ + @article{hong2024orpo, + title = {{ORPO: Monolithic Preference Optimization without Reference Model}}, + author = {Jiwoo Hong and Noah Lee and James Thorne}, + year = 2024, + eprint = {arXiv:2403.07691} + }""") + + 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="ORPO", + trainer_citation=citation, + paper_title="ORPO: Monolithic Preference Optimization without Reference Model", + paper_id="2403.07691", + ) + + model_card.save(os.path.join(self.args.output_dir, "README.md")) +class UnslothORPOTrainer(_UnslothORPOTrainer): + """ + + Initialize ORPOTrainer. + + Args: + model (`transformers.PreTrainedModel`): + The model to train, preferably an `AutoModelForSequenceClassification`. + args (`ORPOConfig`): + The ORPO 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 (`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. + 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. + 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. + + """ + def __init__( + self, + model = None, + args = None, + data_collator = None, + train_dataset = None, + eval_dataset = None, + processing_class = None, + model_init = None, + callbacks = None, + preprocess_logits_for_metrics = None, + peft_config = None, + compute_metrics = None, + **kwargs + ): + if args is None: args = UnslothORPOConfig() + 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 + 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) + 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 + 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('orpo_trainer', other_metrics) + + 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, + callbacks = callbacks, + preprocess_logits_for_metrics = preprocess_logits_for_metrics, + peft_config = peft_config, + compute_metrics = compute_metrics,**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 diff --git a/unsloth_compiled_cache/UnslothOnlineDPOTrainer.py b/unsloth_compiled_cache/UnslothOnlineDPOTrainer.py new file mode 100644 index 0000000..1547c70 --- /dev/null +++ b/unsloth_compiled_cache/UnslothOnlineDPOTrainer.py @@ -0,0 +1,1260 @@ +""" +2025.7.1 +2025.7.1 +4.53.1 +0.19.1 +__UNSLOTH_VERSIONING__ +""" +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from trl.trainer.online_dpo_trainer import (Any, BaseImageProcessor, BasePairwiseJudge, Callable, DPODataCollatorWithPadding, DataCollator, DataLoader, Dataset, EvalPrediction, F, FeatureExtractionMixin, GenerationConfig, IterableDataset, OnlineDPOConfig, OnlineDPOTrainer, OptimizerNames, Optional, Path, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, SIMPLE_CHAT_TEMPLATE, Trainer, TrainerCallback, Union, apply_chat_template, create_reference_model, datasets, disable_dropout_in_model, empty_cache, generate_model_card, get_comet_experiment_url, get_reward, is_conversational, is_peft_available, is_wandb_available, jinja2, logging, maybe_apply_chat_template, nn, os, prepare_deepspeed, seed_worker, textwrap, torch, truncate_right, unwrap_model_for_generation, version, warnings, wraps, F, is_conversational, os, torch, F, Optional, PeftModel, PreTrainedModel, Trainer, is_peft_available, os, torch) + + +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 +def vLLMSamplingParams(**kwargs): + from vllm import SamplingParams + sampling_params = SamplingParams(**kwargs) + sampling_params._set_kwargs = kwargs + return sampling_params +@dataclass +class UnslothOnlineDPOConfig(OnlineDPOConfig): + """ + + Configuration class for the [`OnlineDPOTrainer`]. + + This class includes only the parameters that are specific to Online DPO training. For a full list of training + arguments, please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this + class may differ from those in [`~transformers.TrainingArguments`]. + + 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: + reward_model_path (`str` or `None`, *optional*, defaults to `None`): + Path to the reward model. Either `judge` or `reward_model_path` must be set, but not both. + judge (`str` or `None`, *optional*, defaults to `None`): + Name of the judge to use. Either `judge` or `reward_model_path` must be set, but not both. + max_new_tokens (`int`, *optional*, defaults to `64`): + Maximum number of tokens to generate per completion. + max_length (`int`, *optional*, defaults to `256`): + Maximum total length of the sequence (prompt + completion) used to compute log probabilities. If the + sequence exceeds this limit, the leftmost tokens will be truncated to preserve as much of the completion as + possible. + temperature (`float`, *optional*, defaults to `0.9`): + Temperature for sampling. The higher the temperature, the more random the completions. + missing_eos_penalty (`float` or `None`, *optional*, defaults to `None`): + Penalty applied to the score when the model fails to generate an EOS token. This is useful to encourage to + generate completions shorter than the maximum length (`max_new_tokens`). The penalty must be a positive + value. + beta (`float` or `list[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). If a list of floats is provided then the β is + selected for each new epoch and the last β is used for the rest of the epochs. + 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. + - `"ipo"`: IPO loss from the [IPO](https://huggingface.co/papers/2310.12036) paper. + + dataset_num_proc (`int` or `None`, *optional*, defaults to `None`): + Number of processes to use for processing the dataset. + disable_dropout (`bool`, *optional*, defaults to `True`): + Whether to disable dropout in the model and reference model. + use_vllm (`bool`, *optional*, defaults to `False`): + Whether to use vLLM for generating completions. Requires vLLM to be installed (`pip install vllm`). + gpu_memory_utilization (`float`, *optional*, defaults to `0.55`): + The vLLM memory utilization. The default value is 0.55. + ds3_gather_for_generation (`bool`, *optional*, defaults to `True`): + This setting applies to DeepSpeed ZeRO-3. If enabled, the policy model weights are gathered for generation, + improving generation speed. However, disabling this option allows training models that exceed the VRAM + capacity of a single GPU, albeit at the cost of slower generation. + + """ + 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, + hub_revision = None, + 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, + liger_kernel_config = None, + eval_use_gather_object = False, + average_tokens_across_devices = False, + reward_model_path = None, + judge = None, + max_new_tokens = 64, + max_length = 512, + temperature = 0.9, + missing_eos_penalty = None, + loss_type = 'sigmoid', + dataset_num_proc = None, + disable_dropout = True, + use_vllm = False, + gpu_memory_utilization = 0.55, + ds3_gather_for_generation = True, + 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() + if temperature <= 0: + raise MathError('Unsloth: Please set a positive non-zero temperature since your results will be wrong.') + elif temperature >= 10: + raise MathError('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, + 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, + 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, + missing_eos_penalty = missing_eos_penalty, + loss_type = loss_type, + dataset_num_proc = dataset_num_proc, + disable_dropout = disable_dropout, + use_vllm = use_vllm, + gpu_memory_utilization = gpu_memory_utilization, + ds3_gather_for_generation = ds3_gather_for_generation,**kwargs) + self.vllm_sampling_params = vllm_sampling_params + self.unsloth_num_chunks = unsloth_num_chunks +pass + +class _UnslothOnlineDPOTrainer(Trainer): + r"""""" + + _tag_names = ["trl", "online-dpo"] + + def __init__( + self, + model: Union[PreTrainedModel, nn.Module], + ref_model: Union[PreTrainedModel, nn.Module, None] = None, + reward_model: Union[PreTrainedModel, nn.Module, None] = None, + judge: Optional[BasePairwiseJudge] = None, + args: Optional[OnlineDPOConfig] = None, + data_collator: Optional[DataCollator] = None, + train_dataset: Optional[Union[Dataset, IterableDataset, "datasets.Dataset"]] = None, + eval_dataset: Optional[Union[Dataset, dict[str, Dataset], "datasets.Dataset"]] = None, + processing_class: Optional[ + Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin] + ] = None, + reward_processing_class: Optional[PreTrainedTokenizerBase] = 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, + ) -> None: + + if hasattr(model, 'vllm_engine') and hasattr(args, 'use_vllm'): + if (getattr(args, 'use_vllm', False) == False): + args.use_vllm = True + if 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`, either omit the `ref_model` argument or pass `None`." + ) + + self.ref_model = ref_model + + if reward_model is not None and judge is not None: + warnings.warn( + "Both `reward_model` and `judge` are provided. Please choose provide only one of them. " + "Ignoring `judge` and using `reward_model`.", + UserWarning, + ) + judge = None + elif reward_model is None and judge is None: + raise ValueError("Either `reward_model` or `judge` must be provided.") + + self.reward_model = reward_model + self.reward_processing_class = reward_processing_class + self.judge = judge + self.is_encoder_decoder = model.config.is_encoder_decoder + + if args.missing_eos_penalty is not None and judge is not None: + raise ValueError("`missing_eos_penalty` is not supported when `judge` is provided.") + + if args is None: + raise ValueError("`args` must be provided.") + + # Check that the processing_class is provided + if processing_class is None: + raise ValueError("`processing_class` must be provided.") + + # Convert to PEFT model if peft_config is provided + if False: + # Check if PEFT is available + if not is_peft_available(): + raise ImportError( + "PEFT is not available and passed `peft_config`. Please install PEFT with " + "`pip install peft` to use it." + ) + + # If the model is already a PeftModel, we need to merge and unload it. + # Further information here: https://huggingface.co/docs/trl/dpo_trainer#reference-model-considerations-with-peft + if isinstance(model, PeftModel): + model = model.merge_and_unload() + + # Get peft model with the given config + model = model + + # 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) + + # Handle the ref_model + # Usually, the user wants the ref model to be the initial version of the model. When using PEFT, it's easy to + # get the ref model, as it's just the model with a disabled adapter. When not using PEFT, we need to create + # the ref model from the model by copying it and disable the gradients and set it in evaluation mode. + if ref_model is None: # No ref model provided, the most common case + if False: + self.ref_model = create_reference_model(model) # copy, disable gradients, set eval mode + else: + self.ref_model = None # we don't need a ref model here, we can just disable the adapter. + else: # rare case, the user provided a ref model + self.ref_model = ref_model + self.ref_model.eval() + + # Disable the gradient and set the reward model in eval mode + if self.reward_model is not None: + self.reward_model.eval() + + # Define the collator is not provided + if data_collator is None: + data_collator = DPODataCollatorWithPadding(pad_token_id=processing_class.pad_token_id) + + self.max_length = args.max_length + + self.stats = { + "objective/kl": [], + "objective/entropy": [], + "objective/non_score_reward": [], + "rewards/chosen": [], + "rewards/rejected": [], + "rewards/accuracies": [], + "rewards/margins": [], + "logps/chosen": [], + "logps/rejected": [], + "val/contain_eos_token": [], + "beta": [], + } + if self.reward_model is not None: + self.stats["objective/rlhf_reward"] = [] + self.stats["objective/scores_margin"] = [] + self.stats["objective/scores"] = [] + + if args.use_vllm: + self.llm = model.vllm_engine; self._last_loaded_step = 0; self.generation_config = SamplingParams( + n=2, + max_tokens=args.max_new_tokens, + temperature=args.temperature, + top_k=50, + top_p=1.0, + detokenize=False, + **getattr(getattr(args, 'vllm_sampling_params', vLLMSamplingParams()), '_set_kwargs', {}), + ) + else: + self.generation_config = GenerationConfig( + max_new_tokens=args.max_new_tokens, + temperature=args.temperature, + top_k=50, + top_p=1.0, + do_sample=True, + use_cache=False if args.gradient_checkpointing else True, + ) + + # 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 Online DPO, the sampled data does not include + # the "input_ids" key. 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 + + super().__init__( + model=model, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + ) + + # 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) + + self._beta = args.beta + + # Placed after the super[].__init__ because we need self.is_deepspeed_enabled and self.accelerator + if self.is_deepspeed_enabled: + if self.reward_model is not None: + self.reward_model = prepare_deepspeed( + self.reward_model, args.per_device_train_batch_size, args.fp16, args.bf16 + ) + if self.ref_model is not None: + self.ref_model = prepare_deepspeed( + self.ref_model, args.per_device_train_batch_size, args.fp16, args.bf16 + ) + else: + if self.ref_model is not None: + self.ref_model = self.ref_model.to(self.accelerator.device) + if self.reward_model is not None: + self.reward_model = self.reward_model.to(self.accelerator.device) + + @property + def beta(self): + if isinstance(self._beta, list): + epoch = self.state.epoch + return self._beta[epoch] if epoch < len(self._beta) else self._beta[-1] + else: + return self._beta + + @staticmethod + def tokenize_row(feature, is_encoder_decoder: bool, tokenizer: PreTrainedTokenizerBase) -> dict[str, Any]: + """Tokenize a single row from a DPO specific dataset.""" + if not is_encoder_decoder: + batch = tokenizer(feature["prompt"], add_special_tokens=False) + # Add BOS token to head of prompt. Avoid adding if it's already there + if tokenizer.bos_token_id is not None: + prompt_len_input_ids = len(batch["input_ids"]) + if prompt_len_input_ids == 0 or tokenizer.bos_token_id != batch["input_ids"][0]: + batch["input_ids"] = [tokenizer.bos_token_id] + batch["input_ids"] + batch["attention_mask"] = [1] + batch["attention_mask"] + else: + batch = tokenizer(feature["prompt"], add_special_tokens=True) + batch = {f"prompt_{key}": value for key, value in batch.items()} + return batch + + # Same as Trainer.get_train_dataloader but skip the "remove_unused_columns". + @wraps(Trainer.get_train_dataloader) + def get_train_dataloader(self) -> DataLoader: + if self.train_dataset is None: + raise ValueError("Trainer: training requires a train_dataset.") + + train_dataset = self.train_dataset + data_collator = self.data_collator + dataloader_params = { + "batch_size": self._train_batch_size, + "collate_fn": data_collator, + "num_workers": self.args.dataloader_num_workers, + "pin_memory": self.args.dataloader_pin_memory, + "persistent_workers": self.args.dataloader_persistent_workers, + } + + if not isinstance(train_dataset, torch.utils.data.IterableDataset): + dataloader_params["sampler"] = self._get_train_sampler() + dataloader_params["drop_last"] = self.args.dataloader_drop_last + dataloader_params["worker_init_fn"] = seed_worker + dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor + + return self.accelerator.prepare(DataLoader(train_dataset, **dataloader_params)) + + # Same as Trainer.get_eval_dataloader but skip the "remove_unused_columns". + @wraps(Trainer.get_eval_dataloader) + def get_eval_dataloader(self, eval_dataset: Optional[Union[str, Dataset]] = None) -> DataLoader: + if eval_dataset is None and self.eval_dataset is None: + raise ValueError("Trainer: evaluation requires an eval_dataset.") + + # If we have persistent workers, don't do a fork bomb especially as eval datasets + # don't change during training + dataloader_key = eval_dataset if isinstance(eval_dataset, str) else "eval" + if ( + hasattr(self, "_eval_dataloaders") + and dataloader_key in self._eval_dataloaders + and self.args.dataloader_persistent_workers + ): + return self.accelerator.prepare(self._eval_dataloaders[dataloader_key]) + + eval_dataset = ( + self.eval_dataset[eval_dataset] + if isinstance(eval_dataset, str) + else eval_dataset + if eval_dataset is not None + else self.eval_dataset + ) + data_collator = self.data_collator + + dataloader_params = { + "batch_size": self.args.eval_batch_size, + "collate_fn": data_collator, + "num_workers": self.args.dataloader_num_workers, + "pin_memory": self.args.dataloader_pin_memory, + "persistent_workers": self.args.dataloader_persistent_workers, + } + + if not isinstance(eval_dataset, torch.utils.data.IterableDataset): + dataloader_params["sampler"] = self._get_eval_sampler(eval_dataset) + dataloader_params["drop_last"] = self.args.dataloader_drop_last + dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor + + # accelerator.free_memory() will destroy the references, so + # we need to store the non-prepared version + eval_dataloader = DataLoader(eval_dataset, **dataloader_params) + if self.args.dataloader_persistent_workers: + if hasattr(self, "_eval_dataloaders"): + self._eval_dataloaders[dataloader_key] = eval_dataloader + else: + self._eval_dataloaders = {dataloader_key: eval_dataloader} + + return self.accelerator.prepare(eval_dataloader) + + def _generate_vllm(self, model, prompts): + eos_token_id = self.processing_class.eos_token_id + pad_token_id = self.processing_class.pad_token_id + + # Load the latest weights + + pass + + pass + + if is_conversational({"prompt": prompts[0]}): + outputs = self.llm.chat(prompts, self.generation_config, use_tqdm=False, lora_request = self.model.load_lora('online_dpo_trainer_lora_model', load_tensors = True)) + else: + outputs = self.llm.generate(prompts, self.generation_config, use_tqdm=False, lora_request = self.model.load_lora('online_dpo_trainer_lora_model', load_tensors = True)) + + completion_ids = [list(output.outputs[i].token_ids) for i in range(2) for output in outputs] + prompt_ids = [list(output.prompt_token_ids) for _ in range(2) for output in outputs] + + # Create mask and pad the prompt and completion + max_prompt_length = max(len(ids) for ids in prompt_ids) + prompt_mask = [[0] * (max_prompt_length - len(ids)) + [1] * len(ids) for ids in prompt_ids] + prompt_ids = [[pad_token_id] * (max_prompt_length - len(ids)) + ids for ids in prompt_ids] + max_tokens = self.generation_config.max_tokens + completion_mask = [[1] * len(ids) + [0] * (max_tokens - len(ids)) for ids in completion_ids] + completion_ids = [ + ids + [eos_token_id] if ids[-1] != eos_token_id and len(ids) < max_tokens else ids + for ids in completion_ids + ] + completion_ids = [ids + [pad_token_id] * (max_tokens - len(ids)) for ids in completion_ids] + + # Convert to tensors + prompt_ids = torch.tensor(prompt_ids, device=self.accelerator.device) + prompt_mask = torch.tensor(prompt_mask, device=self.accelerator.device) + completion_ids = torch.tensor(completion_ids, device=self.accelerator.device) + completion_mask = torch.tensor(completion_mask, device=self.accelerator.device) + + return prompt_ids, prompt_mask, completion_ids, completion_mask + + def _generate(self, model, prompts): + eos_token_id = self.processing_class.eos_token_id + pad_token_id = self.processing_class.pad_token_id + + # Apply chat template and tokenize the input. We do this on-the-fly to enable the use of reward models and + # policies with different tokenizers / chat templates. + inputs = [{"prompt": prompt} for prompt in prompts] + inputs = [maybe_apply_chat_template(x, self.processing_class) for x in inputs] + inputs = [self.tokenize_row(x, self.is_encoder_decoder, self.processing_class) for x in inputs] + inputs = self.data_collator(inputs) + + # Sample 2 completions per prompt of size `max_new_tokens` from the model + inputs = self._prepare_inputs(inputs) + prompt_ids = inputs["prompt_input_ids"].repeat(2, 1) + prompt_mask = inputs["prompt_attention_mask"].repeat(2, 1) + with unwrap_model_for_generation( + model, self.accelerator, gather_deepspeed3_params=self.args.ds3_gather_for_generation + ) as unwrapped_model: + output = unwrapped_model.generate( + input_ids=prompt_ids, + attention_mask=prompt_mask, + generation_config=self.generation_config, + ) + + completion_ids = output[:, prompt_ids.size(1) :] + completion_ids, completion_mask = truncate_right(completion_ids, eos_token_id, pad_token_id) + + return prompt_ids, prompt_mask, completion_ids, completion_mask + + def _forward(self, model, prompt_ids, prompt_mask, completion_ids, completion_mask): + # Get the number of tokens to truncate from prompt + num_tokens_to_truncate = max(prompt_ids.size(1) + completion_ids.size(1) - self.max_length, 0) + + # Truncate left to avoid oom + prompt_ids = prompt_ids[:, num_tokens_to_truncate:] + prompt_mask = prompt_mask[:, num_tokens_to_truncate:] + + # Concat the prompt and completion + prompt_completion_ids = torch.cat((prompt_ids, completion_ids), dim=1) + prompt_completion_mask = torch.cat((prompt_mask, completion_mask), dim=1) + + # Get the logprobs of the completions from the model + output = model(prompt_completion_ids, attention_mask=prompt_completion_mask) + + # There is 1 offset, because the model predict the next token + logits = output.logits[:, prompt_ids.size(1) - 1 : -1] + + # Take the completion tokens logprob + logprobs = torch.take_along_dim(logits.log_softmax(dim=-1), completion_ids.unsqueeze(-1), dim=2).squeeze(-1) + return logprobs + + 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() + + prompts = inputs["prompt"] + batch_size = len(prompts) + + if self.args.use_vllm: + prompt_ids, prompt_mask, completion_ids, completion_mask = self._generate_vllm(model, prompts) + else: + prompt_ids, prompt_mask, completion_ids, completion_mask = self._generate(model, prompts) + + contain_eos_token = torch.any(completion_ids == self.processing_class.eos_token_id, dim=-1) + + logprobs = self._forward(model, prompt_ids, prompt_mask, completion_ids, completion_mask) + with torch.no_grad(): + if self.ref_model is not None: + ref_logprobs = self._forward(self.ref_model, prompt_ids, prompt_mask, completion_ids, completion_mask) + else: # peft case: we just need to disable the adapter + with self.model.disable_adapter(): + ref_logprobs = self._forward(self.model, prompt_ids, prompt_mask, completion_ids, completion_mask) + + # Decode the completions, and format them if the input is conversational + device = logprobs.device + completions = self.processing_class.batch_decode(completion_ids, skip_special_tokens=True) + if is_conversational({"prompt": prompts[0]}): + completions = [[{"role": "assistant", "content": completion}] for completion in completions] + + # Get the reward from the reward model or judge + if self.judge is not None: + # Once formatted, conversational data may contain special tokens (such as <|im_start|>) that are not + # directly understandable by the judge and could alter its judgment. To avoid this and make the judge + # independent of the model's chat template, we use the raw conversation data, and apply our own chat + # template to it. + if is_conversational({"prompt": prompts[0]}): + environment = jinja2.Environment() + template = environment.from_string(SIMPLE_CHAT_TEMPLATE) + prompts = [template.render(messages=prompt) for prompt in prompts] + completions = [template.render(messages=completion) for completion in completions] + + ranks_of_first_completion = self.judge.judge( + prompts, list(zip(completions[:batch_size], completions[batch_size:])) + ) + + # convert ranks to a True/False mask: + # when rank == 0, it means the first completion is the best + # when rank == 1, it means the second completion is the best + mask = torch.tensor([rank == 0 for rank in ranks_of_first_completion], device=device) + else: + # The reward model may not have the same chat template or tokenizer as the model, so we need to use the + # raw data (string), apply the chat template (if needed), and tokenize it with the reward processing class. + prompts = 2 * prompts # repeat the prompt: [prompt0, prompt1] -> [prompt0, prompt1, prompt0, prompt1] + if is_conversational({"prompt": prompts[0]}): + examples = [{"prompt": p, "completion": c} for p, c in zip(prompts, completions)] + examples = [apply_chat_template(example, self.reward_processing_class) for example in examples] + prompts = [example["prompt"] for example in examples] + completions = [example["completion"] for example in examples] + + # Tokenize the prompts + prompts_ids = self.reward_processing_class( + prompts, padding=True, return_tensors="pt", padding_side="left" + )["input_ids"].to(device) + context_length = prompts_ids.shape[1] + + # Tokenize the completions + completions_ids = self.reward_processing_class( + completions, padding=True, return_tensors="pt", padding_side="right" + )["input_ids"].to(device) + + # Concatenate the prompts and completions and get the reward + prompt_completion_ids = torch.cat((prompts_ids, completions_ids), dim=1) + with torch.inference_mode(): + _, scores, _ = get_reward( + self.reward_model, prompt_completion_ids, self.reward_processing_class.pad_token_id, context_length + ) + + # Filter completion. Ensure that the sample contains stop_token_id + # Completions not passing that filter will receive a lower score. + if self.args.missing_eos_penalty is not None: + scores[~contain_eos_token] -= self.args.missing_eos_penalty + + # Split the scores in 2 (the prompts of the first half are the same as the second half) + first_half, second_half = scores.split(batch_size) + + # Get the indices of the chosen and rejected examples + mask = first_half >= second_half + + batch_range = torch.arange(batch_size, device=device) + chosen_indices = batch_range + (~mask * batch_size) + rejected_indices = batch_range + (mask * batch_size) + + # Build tensor so that the first half is the chosen examples and the second half the rejected examples + cr_indices = torch.cat((chosen_indices, rejected_indices), dim=0) # cr = chosen and rejected + cr_logprobs = logprobs[cr_indices] + cr_ref_logprobs = ref_logprobs[cr_indices] + + # mask out the padding tokens + padding_mask = ~completion_mask.bool() + cr_padding_mask = padding_mask[cr_indices] + + cr_logprobs_sum = (cr_logprobs * ~cr_padding_mask).sum(1) + cr_ref_logprobs_sum = (cr_ref_logprobs * ~cr_padding_mask).sum(1) + + # Split the chosen and rejected examples + chosen_logprobs_sum, rejected_logprobs_sum = torch.split(cr_logprobs_sum, batch_size) + chosen_ref_logprobs_sum, rejected_ref_logprobs_sum = torch.split(cr_ref_logprobs_sum, batch_size) + pi_logratios = chosen_logprobs_sum - rejected_logprobs_sum + ref_logratios = chosen_ref_logprobs_sum - rejected_ref_logprobs_sum + + logits = pi_logratios - ref_logratios + + if self.args.loss_type == "sigmoid": + losses = -F.logsigmoid(self.beta * logits) + elif self.args.loss_type == "ipo": + losses = (logits - 1 / (2 * self.beta)) ** 2 + else: + raise NotImplementedError(f"invalid loss type {self.loss_type}") + + loss = losses.mean() + + # Log everything + if self.reward_model is not None: + scores_margin = scores[chosen_indices] - scores[rejected_indices] + self.stats["objective/scores_margin"].append( + self.accelerator.gather_for_metrics(scores_margin.mean()).mean().item() + ) + self.stats["objective/scores"].append(self.accelerator.gather_for_metrics(scores.mean()).mean().item()) + self.stats["val/contain_eos_token"].append(contain_eos_token.float().mean().item()) + self.stats["logps/chosen"].append(self.accelerator.gather_for_metrics(chosen_logprobs_sum).mean().item()) + self.stats["logps/rejected"].append(self.accelerator.gather_for_metrics(rejected_logprobs_sum).mean().item()) + + kl = logprobs - ref_logprobs + mean_kl = kl.sum(1).mean() + self.stats["objective/kl"].append(self.accelerator.gather_for_metrics(mean_kl).mean().item()) + non_score_reward = (-self.beta * kl).sum(1) + mean_non_score_reward = non_score_reward.mean() + self.stats["objective/non_score_reward"].append( + self.accelerator.gather_for_metrics(mean_non_score_reward).mean().item() + ) + if self.reward_model is not None: + rlhf_reward = scores + non_score_reward + self.stats["objective/rlhf_reward"].append(self.accelerator.gather_for_metrics(rlhf_reward).mean().item()) + mean_entropy = -logprobs.sum(1).mean() + self.stats["objective/entropy"].append(self.accelerator.gather_for_metrics(mean_entropy).mean().item()) + chosen_rewards = self.beta * (chosen_logprobs_sum - chosen_ref_logprobs_sum) + gathered_chosen_rewards = self.accelerator.gather_for_metrics(chosen_rewards) + self.stats["rewards/chosen"].append(gathered_chosen_rewards.mean().item()) + rejected_rewards = self.beta * (rejected_logprobs_sum - rejected_ref_logprobs_sum) + gathered_rejected_rewards = self.accelerator.gather_for_metrics(rejected_rewards) + self.stats["rewards/rejected"].append(gathered_rejected_rewards.mean().item()) + margin = gathered_chosen_rewards - gathered_rejected_rewards + self.stats["rewards/margins"].append(margin.mean().item()) + accuracy = margin > 0 + self.stats["rewards/accuracies"].append(accuracy.float().mean().item()) + self.stats["beta"].append(self.beta) + + 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 + + if self.use_apex: + with amp.scale_loss(loss, self.optimizer) as scaled_loss: + scaled_loss.backward() + else: + self.accelerator.backward(loss, **kwargs) + + return loss.detach() / self.args.gradient_accumulation_steps + + # Same as Trainer._maybe_log_save_evaluate but log our metrics + def _maybe_log_save_evaluate( + self, tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval, start_time, learning_rate=None + ): + if self.control.should_log and self.state.global_step > self._globalstep_last_logged: + logs: dict[str, float] = {} + + # all_gather + mean() to get average loss over all processes + tr_loss_scalar = self._nested_gather(tr_loss).mean().item() + + # reset tr_loss to zero + tr_loss -= tr_loss + + logs["loss"] = round(tr_loss_scalar / (self.state.global_step - self._globalstep_last_logged), 4) + if grad_norm is not None: + logs["grad_norm"] = grad_norm.detach().item() if isinstance(grad_norm, torch.Tensor) else grad_norm + if learning_rate is not None: + logs["learning_rate"] = learning_rate + else: + logs["learning_rate"] = self._get_learning_rate() + + # Add our metrics + for key, val in self.stats.items(): + logs[key] = sum(val) / len(val) + self.stats = {key: [] for key in self.stats} # reset stats + + self._total_loss_scalar += tr_loss_scalar + self._globalstep_last_logged = self.state.global_step + self.store_flos() + self.log(logs, start_time) + + metrics = None + if self.control.should_evaluate: + metrics = self._evaluate(trial, ignore_keys_for_eval) + is_new_best_metric = self._determine_best_metric(metrics=metrics, trial=trial) + + if self.args.save_strategy == "best": + self.control.should_save = is_new_best_metric + + if self.control.should_save: + self._save_checkpoint(model, trial) + self.control = self.callback_handler.on_save(self.args, self.state, self.control) + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) + + 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 + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + citation = textwrap.dedent("""\ + @article{guo2024direct, + title = {{Direct Language Model Alignment from Online AI Feedback}}, + author = {Shangmin Guo and Biao Zhang and Tianlin Liu and Tianqi Liu and Misha Khalman and Felipe Llinares and Alexandre Ram{\'{e}} and Thomas Mesnard and Yao Zhao and Bilal Piot and Johan Ferret and Mathieu Blondel}, + year = 2024, + eprint = {arXiv:2402.04792} + }""") + + 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="Online DPO", + trainer_citation=citation, + paper_title="Direct Language Model Alignment from Online AI Feedback", + paper_id="2402.04792", + ) + model_card.save(os.path.join(self.args.output_dir, "README.md")) +class UnslothOnlineDPOTrainer(_UnslothOnlineDPOTrainer): + """ + + Initialize OnlineDPOTrainer. + + Args: + model (`transformers.PreTrainedModel` or `torch.nn.Module`): + The model to train, preferably an `AutoModelForCausalLM`. + ref_model (`transformers.PreTrainedModel` or `torch.nn.Module` or `None`): + The reference model to use for training. If None is specified, the reference model will be created from the + model. + reward_model (`transformers.PreTrainedModel` or `torch.nn.Module` or `None`): + The reward model to score completions with, preferably an `AutoModelForSequenceClassification`. + judge (`BasePairwiseJudge`): + The judge to use for pairwise comparison of model completions. + args (`OnlineDPOConfig`): + The online 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 + (`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 (`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. + 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. + + """ + def __init__( + self, + model, + ref_model = None, + reward_model = None, + judge = None, + args = None, + data_collator = None, + train_dataset = None, + eval_dataset = None, + processing_class = None, + reward_processing_class = None, + peft_config = None, + compute_metrics = None, + callbacks = None, + preprocess_logits_for_metrics = None, + **kwargs + ): + if args is None: args = UnslothOnlineDPOConfig() + 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 + 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) + 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 + 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('online_dpo_trainer', other_metrics) + + super().__init__( + model = model, + ref_model = ref_model, + reward_model = reward_model, + judge = judge, + args = args, + data_collator = data_collator, + train_dataset = train_dataset, + eval_dataset = eval_dataset, + processing_class = processing_class, + reward_processing_class = reward_processing_class, + peft_config = peft_config, + compute_metrics = compute_metrics, + callbacks = callbacks, + preprocess_logits_for_metrics = preprocess_logits_for_metrics,**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 diff --git a/unsloth_compiled_cache/UnslothPPOTrainer.py b/unsloth_compiled_cache/UnslothPPOTrainer.py new file mode 100644 index 0000000..b8ef376 --- /dev/null +++ b/unsloth_compiled_cache/UnslothPPOTrainer.py @@ -0,0 +1,1300 @@ +""" +2025.7.1 +2025.7.1 +4.53.1 +0.19.1 +__UNSLOTH_VERSIONING__ +""" +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from trl.trainer.ppo_trainer import (Accelerator, BaseImageProcessor, CallbackHandler, DEFAULT_CALLBACKS, DEFAULT_PROGRESS_CALLBACK, DataCollatorWithPadding, DataLoader, Dataset, ExportableState, FeatureExtractionMixin, GenerationConfig, INVALID_LOGPROB, OnlineTrainerState, Optional, PPOConfig, PPOTrainer, Path, PeftConfig, PeftModel, PolicyAndValueWrapper, PreTrainedTokenizerBase, PrinterCallback, ProcessorMixin, Trainer, TrainerCallback, TrainerControl, Union, batch_generation, broadcast, contextmanager, create_reference_model, defaultdict, disable_dropout_in_model, empty_cache, exact_div, first_true_indices, forward, gather_object, gc, generate_model_card, get_comet_experiment_url, get_peft_model, get_reporting_integration_callbacks, get_reward, is_peft_available, is_rich_available, is_wandb_available, log_table_to_comet_experiment, masked_mean, masked_whiten, math, nn, np, nullcontext, os, pd, peft_module_casting_to_bf16, prepare_deepspeed, print_rich_table, textwrap, time, torch, truncate_response, unwrap_model_for_generation, Optional, PeftModel, Trainer, is_peft_available, os, torch) + + +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 UnslothPPOConfig(PPOConfig): + """ + + Configuration class for the [`PPOTrainer`]. + + This class includes only the parameters that are specific to PPO training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] and [`OnPolicyConfig`] documentation. Note that default + values in this class may differ from those in [`~transformers.TrainingArguments`]. + + 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: + exp_name (`str`, *optional*, defaults to `os.path.basename(__file__)[:-3]`): + Name of this experiment. + reward_model_path (`str`, *optional*, defaults to `"EleutherAI/pythia-160m"`): + Path to the reward model. + 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. + num_ppo_epochs (`int`, *optional*, defaults to `4`): + Number of epochs to train. + whiten_rewards (`bool`, *optional*, defaults to `False`): + Whether to whiten the rewards. + kl_coef (`float`, *optional*, defaults to `0.05`): + KL coefficient. + kl_estimator (`Literal["k1", "k3"]`, *optional*, defaults to `"k1"`): + Which estimator for KL-Divergence to use from [Approximating KL + Divergence](http://joschu.net/blog/kl-approx.html). Defaults to "k1", a straightforward, unbiased + estimator. Can be set to "k3", an unbiased estimator with lower variance which "appears to be a strictly + better estimator". Cannot be set to "k2", as it is used for logging purposes. + cliprange (`float`, *optional*, defaults to `0.2`): + Clip range. + vf_coef (`float`, *optional*, defaults to `0.1`): + Value function coefficient. + cliprange_value (`float`, *optional*, defaults to `0.2`): + Clip range for the value function. + gamma (`float`, *optional*, defaults to `1.0`): + Discount factor. + lam (`float`, *optional*, defaults to `0.95`): + Lambda value for GAE. + ds3_gather_for_generation (`bool`, *optional*, defaults to `True`): + This setting applies to DeepSpeed ZeRO-3. If enabled, the policy model weights are gathered for generation, + improving generation speed. However, disabling this option allows training models that exceed the VRAM + capacity of a single GPU, albeit at the cost of slower generation. + + """ + 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, + hub_revision = None, + 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, + liger_kernel_config = None, + eval_use_gather_object = False, + average_tokens_across_devices = False, + dataset_num_proc = None, + num_mini_batches = 1, + total_episodes = None, + local_rollout_forward_batch_size = 64, + num_sample_generations = 10, + response_length = 53, + stop_token = None, + stop_token_id = None, + temperature = 0.7, + missing_eos_penalty = None, + sft_model_path = 'EleutherAI/pythia-160m', + world_size = None, + num_total_batches = None, + micro_batch_size = None, + local_batch_size = None, + batch_size = None, + local_mini_batch_size = None, + mini_batch_size = None, + exp_name = 'ppo_config', + reward_model_path = 'EleutherAI/pythia-160m', + model_adapter_name = None, + ref_adapter_name = None, + num_ppo_epochs = 4, + whiten_rewards = False, + kl_coef = 0.05, + kl_estimator = 'k1', + cliprange = 0.2, + vf_coef = 0.1, + cliprange_value = 0.2, + gamma = 1.0, + lam = 0.95, + ds3_gather_for_generation = True, + 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() + if temperature <= 0: + raise MathError('Unsloth: Please set a positive non-zero temperature since your results will be wrong.') + elif temperature >= 10: + raise MathError('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, + 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, + 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, + dataset_num_proc = dataset_num_proc, + num_mini_batches = num_mini_batches, + total_episodes = total_episodes, + local_rollout_forward_batch_size = local_rollout_forward_batch_size, + num_sample_generations = num_sample_generations, + response_length = response_length, + stop_token = stop_token, + stop_token_id = stop_token_id, + temperature = temperature, + missing_eos_penalty = missing_eos_penalty, + sft_model_path = sft_model_path, + world_size = world_size, + num_total_batches = num_total_batches, + micro_batch_size = micro_batch_size, + local_batch_size = local_batch_size, + batch_size = batch_size, + local_mini_batch_size = local_mini_batch_size, + mini_batch_size = mini_batch_size, + exp_name = exp_name, + reward_model_path = reward_model_path, + model_adapter_name = model_adapter_name, + ref_adapter_name = ref_adapter_name, + num_ppo_epochs = num_ppo_epochs, + whiten_rewards = whiten_rewards, + kl_coef = kl_coef, + kl_estimator = kl_estimator, + cliprange = cliprange, + vf_coef = vf_coef, + cliprange_value = cliprange_value, + gamma = gamma, + lam = lam, + ds3_gather_for_generation = ds3_gather_for_generation,**kwargs) + self.vllm_sampling_params = vllm_sampling_params + self.unsloth_num_chunks = unsloth_num_chunks +pass + +class _UnslothPPOTrainer(Trainer): + _tag_names = ["trl", "ppo"] + + def __init__( + self, + args: PPOConfig, + processing_class: Optional[ + Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin] + ], + model: nn.Module, + ref_model: Optional[nn.Module], + reward_model: nn.Module, + train_dataset: Dataset, + value_model: nn.Module, + data_collator: Optional[DataCollatorWithPadding] = None, + eval_dataset: Optional[Union[Dataset, dict[str, Dataset]]] = None, + # less commonly used + optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), + callbacks: Optional[list[TrainerCallback]] = None, + peft_config: Optional["PeftConfig"] = None, + ) -> None: + if 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 make a copy of it, or `None` if you use peft." + ) + + self.args = args + self.processing_class = processing_class + self.policy_model = model + + # Define the collator if not provided + if data_collator is None: + data_collator = DataCollatorWithPadding(self.processing_class) + + # Handle stop token settings: update policy model's generation_config to use provided stop token + if args.stop_token and args.stop_token_id: + raise ValueError("You cannot set both `stop_token` and `stop_token_id`.") + elif args.stop_token: + if args.stop_token == "eos": + self.policy_model.generation_config.eos_token_id = self.stop_token_id = processing_class.eos_token_id + else: + raise ValueError( + f"Unknown `stop_token` {args.stop_token}. Allowed values are: `'eos'` and `None` (no stop token)." + ) + else: + self.policy_model.generation_config.eos_token_id = self.stop_token_id = args.stop_token_id # None or int + + # Check that the kl estimator is valid + if self.args.kl_estimator not in {"k1", "k3"}: + raise ValueError( + "kl_estimator must be either 'k1' (straightforward, unbiased) or 'k3' (lower variance, unbiased, " + "appears to be a strictly better estimator). See " + "[Approximating KL Divergence](http://joschu.net/blog/kl-approx.html) for details." + ) + + # peft support + if not is_peft_available() and peft_config is not None: + raise ImportError( + "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_confg, we merge and unload it first + if isinstance(self.policy_model, PeftModel): + self.policy_model = self.policy_model.merge_and_unload() + + # get peft model with the given config + self.policy_model = get_peft_model(self.policy_model, peft_config) + if args.bf16 and getattr(self.policy_model, "is_loaded_in_4bit", False): + peft_module_casting_to_bf16(self.policy_model) + + self.is_peft_model = is_peft_available() and isinstance(self.policy_model, PeftModel) + self.model_adapter_name = args.model_adapter_name + self.ref_adapter_name = args.ref_adapter_name + + if ref_model: + self.ref_model = ref_model + elif self.is_peft_model: + self.ref_model = None + else: + self.ref_model = create_reference_model(self.policy_model) + + self.reward_model = reward_model + self.train_dataset = train_dataset + self.train_dataset_len = len(train_dataset) + self.value_model = value_model + self.data_collator = data_collator + self.eval_dataset = eval_dataset + self.optimizer, self.lr_scheduler = optimizers + self.optimizer_cls_and_kwargs = None # needed for transformers >= 4.47 + + ######### + # calculate various batch sizes + ######### + if args.total_episodes is None: # allow the users to define episodes in terms of epochs. + args.total_episodes = int(args.num_train_epochs * self.train_dataset_len) + accelerator = Accelerator(gradient_accumulation_steps=args.gradient_accumulation_steps) + self.accelerator = accelerator + args.world_size = accelerator.num_processes + args.local_batch_size = args.per_device_train_batch_size * args.gradient_accumulation_steps + args.micro_batch_size = int(args.per_device_train_batch_size * args.world_size) + args.batch_size = int(args.local_batch_size * args.world_size) + args.mini_batch_size = exact_div( + args.batch_size, args.num_mini_batches, "`batch_size` must be a multiple of `num_mini_batches`" + ) + args.local_mini_batch_size = exact_div( + args.local_batch_size, args.num_mini_batches, "`local_batch_size` must be a multiple of `num_mini_batches`" + ) + if args.whiten_rewards: + assert args.local_mini_batch_size >= 8, ( + f"Per-rank minibatch size {args.local_mini_batch_size} is insufficient for whitening" + ) + # `per_rank_rollout_batch_size` is our `args.local_batch_size` + # `per_rank_minibatch_size` is our `args.local_mini_batch_size` + args.num_total_batches = math.ceil( + args.total_episodes / args.batch_size + ) # we may train for more than `total_episodes` + time_tensor = torch.tensor(int(time.time()), device=accelerator.device) + time_int = broadcast(time_tensor, 0).item() # avoid different timestamps across processes + args.run_name = f"{args.exp_name}__{args.seed}__{time_int}" + self.local_seed = args.seed + accelerator.process_index * 100003 # Prime + if args.num_sample_generations > 0: + self.sample_generations_freq = max(1, args.num_total_batches // args.num_sample_generations) + self.local_dataloader_batch_size = args.local_batch_size + + ######### + # setup model, optimizer, and others + ######### + for module in [self.policy_model, self.ref_model, self.value_model, self.reward_model]: + if module is not None: + disable_dropout_in_model(module) + self.model = PolicyAndValueWrapper(self.policy_model, self.value_model) + self.model.config = self.policy_model.config # needed for pushing to hub + self.create_optimizer_and_scheduler( + num_training_steps=args.num_total_batches + ) # note that we are calling `self.lr_scheduler.step[]` manually only at the batch level + + ######### + ### trainer specifics + ######### + default_callbacks = DEFAULT_CALLBACKS + get_reporting_integration_callbacks(self.args.report_to) + self.callbacks = default_callbacks if callbacks is None else default_callbacks + callbacks + self.callback_handler = CallbackHandler( + self.callbacks, self.model, self.processing_class, self.optimizer, self.lr_scheduler + ) + self.add_callback(PrinterCallback if self.args.disable_tqdm else DEFAULT_PROGRESS_CALLBACK) + self.control = TrainerControl() + self.state = OnlineTrainerState( + is_local_process_zero=self.is_local_process_zero(), + is_world_process_zero=self.is_world_process_zero(), + stateful_callbacks=[ + cb for cb in self.callback_handler.callbacks + [self.control] if isinstance(cb, ExportableState) + ], + ) + self.current_flos = 0 + self.hp_search_backend = None + self.is_deepspeed_enabled = getattr(self.accelerator.state, "deepspeed_plugin", None) is not None + self.is_fsdp_enabled = getattr(self.accelerator.state, "fsdp_plugin", None) is not None + # Create distant repo and output directory if needed + self.hub_model_id = None + if self.args.push_to_hub: + self.init_hf_repo() + if self.args.should_save: + os.makedirs(self.args.output_dir, exist_ok=True) + + # 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) + + ######### + ### setup dataloader + ######### + self.dataloader = DataLoader( + self.train_dataset, + batch_size=self.local_dataloader_batch_size, + shuffle=True, + collate_fn=self.data_collator, + drop_last=True, # needed; otherwise the last batch will be of ragged shape + ) + # sync random states for DataLoader[shuffle=True] before `accelerator.prepare` + # see https://gist.github.com/vwxyzjn/2581bff1e48e185e0b85b6dfe1def79c + torch.manual_seed(args.seed) + self.model, self.optimizer, self.dataloader = accelerator.prepare(self.model, self.optimizer, self.dataloader) + torch.manual_seed(self.local_seed) # reset the local seed again + + self.eval_dataloader = DataLoader( + self.eval_dataset, + batch_size=args.per_device_eval_batch_size, + collate_fn=self.data_collator, + drop_last=True, + ) # no need to shuffle eval dataset + self.eval_dataloader = accelerator.prepare(self.eval_dataloader) + + if self.is_deepspeed_enabled: + self.reward_model = prepare_deepspeed( + self.reward_model, args.per_device_train_batch_size, args.fp16, args.bf16 + ) + + if self.ref_model is None: + if not self.is_peft_model: + raise ValueError("No reference model and model is not a Peft model.") + else: + self.ref_model = prepare_deepspeed( + self.ref_model, args.per_device_train_batch_size, args.fp16, args.bf16 + ) + else: + if self.ref_model is None: + if not self.is_peft_model: + raise ValueError("No reference model and model is not a Peft model.") + else: + self.ref_model = self.ref_model.to(self.accelerator.device) + self.reward_model = self.reward_model.to(self.accelerator.device) + + def get_train_dataloader(self) -> DataLoader: + return self.dataloader + + def get_eval_dataloader(self) -> DataLoader: + return self.eval_dataloader + + @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.policy).disable_adapter() + if self.is_peft_model and not self.ref_adapter_name + else nullcontext() + ): + if self.ref_adapter_name: + self.model.policy.set_adapter(self.ref_adapter_name) + yield + if self.ref_adapter_name: + self.model.policy.set_adapter(self.model_adapter_name or "default") + + def save_model(self, output_dir: Optional[str] = None, _internal_call: bool = False): + backup_model = self.model + self.model = self.model.policy # save only the policy + + if self.is_deepspeed_enabled: + backup_deepspeed = self.deepspeed + self.deepspeed = self.model + + super().save_model(output_dir, _internal_call) + + self.model = backup_model + + if self.is_deepspeed_enabled: + self.deepspeed = backup_deepspeed + + def train(self): + args = self.args + accelerator = self.accelerator + optimizer = self.optimizer + model = self.model + ref_policy = self.ref_model + reward_model = self.reward_model + processing_class = self.processing_class + dataloader = self.dataloader + device = accelerator.device + + def repeat_generator(): + while True: + yield from dataloader + + iter_dataloader = iter(repeat_generator()) + generation_config = GenerationConfig( + max_new_tokens=args.response_length, + temperature=(args.temperature + 1e-7), + top_k=0.0, + top_p=1.0, + do_sample=True, + ) + + accelerator.print("===training policy===") + start_time = time.time() + stats_shape = (args.num_ppo_epochs, args.num_mini_batches, args.gradient_accumulation_steps) + approxkl_stats = torch.zeros(stats_shape, device=device) + pg_clipfrac_stats = torch.zeros(stats_shape, device=device) + pg_loss_stats = torch.zeros(stats_shape, device=device) + vf_loss_stats = torch.zeros(stats_shape, device=device) + vf_clipfrac_stats = torch.zeros(stats_shape, device=device) + entropy_stats = torch.zeros(stats_shape, device=device) + ratio_stats = torch.zeros(stats_shape, device=device) + model.train() + + # trainer state initialization + self.state.global_step = 0 + self.state.episode = 0 + self.state.max_steps = args.num_total_batches + self.state.num_train_epochs = args.total_episodes / self.train_dataset_len + # Compute absolute values for logging, eval, and save if given as ratio + if args.logging_steps is not None: + if args.logging_steps < 1: + self.state.logging_steps = math.ceil(self.state.max_steps * args.logging_steps) + else: + self.state.logging_steps = args.logging_steps + if args.eval_steps is not None: + if args.eval_steps < 1: + self.state.eval_steps = math.ceil(self.state.max_steps * args.eval_steps) + else: + self.state.eval_steps = args.eval_steps + if args.save_steps is not None: + if args.save_steps < 1: + self.state.save_steps = math.ceil(self.state.max_steps * args.save_steps) + else: + self.state.save_steps = args.save_steps + self.control = self.callback_handler.on_train_begin(args, self.state, self.control) + + # backward compatibility + if self.is_deepspeed_enabled: + self.deepspeed = self.model + self.model_wrapped = self.model + + for update in range(1, args.num_total_batches + 1): + self.state.episode += 1 * args.batch_size + data = next(iter_dataloader) + with torch.no_grad(): + queries = data["input_ids"].to(device) + context_length = queries.shape[1] + responses = [] + postprocessed_responses = [] + logprobs = [] + ref_logprobs = [] + scores = [] + sequence_lengths = [] + values = [] + with unwrap_model_for_generation( + self.model, self.accelerator, gather_deepspeed3_params=self.args.ds3_gather_for_generation + ) as unwrapped_model: + query_responses, logitss = batch_generation( + unwrapped_model.policy, + queries, + args.local_rollout_forward_batch_size, + processing_class.pad_token_id, + generation_config, + ) + + for i in range(0, queries.shape[0], args.local_rollout_forward_batch_size): + query = queries[i : i + args.local_rollout_forward_batch_size] + query_response = query_responses[i : i + args.local_rollout_forward_batch_size] + response = query_response[:, context_length:] + logits = logitss[i : i + args.local_rollout_forward_batch_size] + logprob = selective_log_softmax(logits, response) + del logits + empty_cache() + + if ref_policy is None: + with self.null_ref_context(): + ref_output = forward(model.policy, query_response, processing_class.pad_token_id) + else: + ref_output = forward(ref_policy, query_response, processing_class.pad_token_id) + ref_logits = ref_output.logits[:, context_length - 1 : -1] + ref_logits /= args.temperature + 1e-7 + ref_logprob = selective_log_softmax(ref_logits, response) + del ref_output, ref_logits + empty_cache() + + # Response Processing 1. truncate response after the first occurrence of `stop_token_id` + postprocessed_response = response + if self.stop_token_id is not None: # handle the edge case when stop_token_id exists but is 0 + postprocessed_response = truncate_response( + self.stop_token_id, processing_class.pad_token_id, response + ) + + # Response Processing 2. run reward model on the truncated responses + postprocessed_query_response = torch.cat((query, postprocessed_response), 1) + sequence_length = first_true_indices(postprocessed_response == processing_class.pad_token_id) - 1 + unwrapped_value_model = accelerator.unwrap_model(model).value_model + full_value, _, _ = get_reward( + unwrapped_value_model, query_response, processing_class.pad_token_id, context_length + ) + value = full_value[:, context_length - 1 : -1].squeeze(-1) + _, score, _ = get_reward( + reward_model, postprocessed_query_response, processing_class.pad_token_id, context_length + ) + + responses.append(response) + postprocessed_responses.append(postprocessed_response) + logprobs.append(logprob) + ref_logprobs.append(ref_logprob) + sequence_lengths.append(sequence_length) + scores.append(score) + values.append(value) + responses = torch.cat(responses, 0) + postprocessed_responses = torch.cat(postprocessed_responses, 0) + logprobs = torch.cat(logprobs, 0) + ref_logprobs = torch.cat(ref_logprobs, 0) + sequence_lengths = torch.cat(sequence_lengths, 0) + scores = torch.cat(scores, 0) + values = torch.cat(values, 0) + del (logprob, ref_logprob, full_value, value, score, unwrapped_model) + empty_cache() + gc.collect() + + # Response Processing 3. Filter completion. Ensure that the sample contains stop_token_id + # Completions not passing that filter will receive a lower score. + contain_eos_token = torch.any(postprocessed_responses == self.processing_class.eos_token_id, dim=-1) + if self.args.missing_eos_penalty is not None: + scores[~contain_eos_token] -= self.args.missing_eos_penalty + # accelerator.print(f"{scores=}, {(contain_eos_token.sum() / len(contain_eos_token))=}") + + # be very careful with `padding_mask_p1`; see https://excalidraw.com/#json=LWnzG4w2k5DjF_EOL_xPt,e2w3a-hFJ_gX5vOfeyXGTw + response_idxs = torch.arange(responses.shape[1], device=responses.device).repeat(responses.shape[0], 1) + padding_mask = response_idxs > sequence_lengths.unsqueeze(1) + logprobs = torch.masked_fill(logprobs, padding_mask, INVALID_LOGPROB) + ref_logprobs = torch.masked_fill(ref_logprobs, padding_mask, INVALID_LOGPROB) + sequence_lengths_p1 = sequence_lengths + 1 + padding_mask_p1 = response_idxs > (sequence_lengths_p1.unsqueeze(1)) + values = torch.masked_fill(values, padding_mask_p1, 0) + + # 4. compute rewards + # Formula used by http://joschu.net/blog/kl-approx.html for the k1 and k3 estimators + logr = ref_logprobs - logprobs + kl = -logr if args.kl_estimator == "k1" else (logr.exp() - 1) - logr # Else statement is k3 + non_score_reward = -args.kl_coef * kl + rewards = non_score_reward.clone() + actual_start = torch.arange(rewards.size(0), device=rewards.device) + actual_end = torch.where(sequence_lengths_p1 < rewards.size(1), sequence_lengths_p1, sequence_lengths) + rewards[[actual_start, actual_end]] += scores + + # 5. whiten rewards + if args.whiten_rewards: + rewards = masked_whiten(rewards, mask=~padding_mask_p1, shift_mean=False) + rewards = torch.masked_fill(rewards, padding_mask_p1, 0) + + # 6. compute advantages and returns + lastgaelam = 0 + advantages_reversed = [] + gen_length = responses.shape[1] + for t in reversed(range(gen_length)): + nextvalues = values[:, t + 1] if t < gen_length - 1 else 0.0 + delta = rewards[:, t] + args.gamma * nextvalues - values[:, t] + lastgaelam = delta + args.gamma * args.lam * lastgaelam + advantages_reversed.append(lastgaelam) + advantages = torch.stack(advantages_reversed[::-1], axis=1) + returns = advantages + values + advantages = masked_whiten(advantages, ~padding_mask) + advantages = torch.masked_fill(advantages, padding_mask, 0) + empty_cache() + + # Do multiple epochs of PPO training, with a fresh random shuffle in each epoch + for ppo_epoch_idx in range(args.num_ppo_epochs): + b_inds = np.random.permutation(args.local_batch_size) + minibatch_idx = 0 + for mini_batch_start in range(0, args.local_batch_size, args.local_mini_batch_size): + mini_batch_end = mini_batch_start + args.local_mini_batch_size + mini_batch_inds = b_inds[mini_batch_start:mini_batch_end] + gradient_accumulation_idx = 0 + for micro_batch_start in range(0, args.local_mini_batch_size, args.per_device_train_batch_size): + with accelerator.accumulate(model): + micro_batch_end = micro_batch_start + args.per_device_train_batch_size + micro_batch_inds = mini_batch_inds[micro_batch_start:micro_batch_end] + mb_advantage = advantages[micro_batch_inds] + mb_responses = responses[micro_batch_inds] + mb_query_responses = query_responses[micro_batch_inds] + mb_logprobs = logprobs[micro_batch_inds] + mb_return = returns[micro_batch_inds] + mb_values = values[micro_batch_inds] + + output, vpred_temp = forward(model, mb_query_responses, processing_class.pad_token_id) + logits = output.logits[:, context_length - 1 : -1] + logits /= args.temperature + 1e-7 + new_logprobs = selective_log_softmax(logits, mb_responses) + new_logprobs = torch.masked_fill( + new_logprobs, padding_mask[micro_batch_inds], INVALID_LOGPROB + ) + vpred = vpred_temp[:, context_length - 1 : -1].squeeze(-1) + vpred = torch.masked_fill(vpred, padding_mask_p1[micro_batch_inds], 0) + vpredclipped = torch.clamp( + vpred, + mb_values - args.cliprange_value, + mb_values + args.cliprange_value, + ) + vf_losses1 = torch.square(vpred - mb_return) + vf_losses2 = torch.square(vpredclipped - mb_return) + vf_loss_max = torch.max(vf_losses1, vf_losses2) + vf_loss = 0.5 * masked_mean(vf_loss_max, ~padding_mask_p1[micro_batch_inds]) + vf_clipfrac = masked_mean( + (vf_losses2 > vf_losses1).float(), ~padding_mask_p1[micro_batch_inds] + ) + logprobs_diff = new_logprobs - mb_logprobs + ratio = torch.exp(logprobs_diff) + pg_losses = -mb_advantage * ratio + pg_losses2 = -mb_advantage * torch.clamp(ratio, 1.0 - args.cliprange, 1.0 + args.cliprange) + pg_loss_max = torch.max(pg_losses, pg_losses2) + pg_loss = masked_mean(pg_loss_max, ~padding_mask[micro_batch_inds]) + loss = pg_loss + args.vf_coef * vf_loss + accelerator.backward(loss) + optimizer.step() + optimizer.zero_grad() + with torch.no_grad(): + pg_clipfrac = masked_mean( + (pg_losses2 > pg_losses).float(), ~padding_mask[micro_batch_inds] + ) + prob_dist = torch.nn.functional.softmax(logits, dim=-1) + entropy = torch.logsumexp(logits, dim=-1) - torch.sum(prob_dist * logits, dim=-1) + approxkl = 0.5 * (logprobs_diff**2).mean() + approxkl_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = approxkl + pg_clipfrac_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = ( + pg_clipfrac + ) + pg_loss_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = pg_loss + vf_loss_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = vf_loss + vf_clipfrac_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = ( + vf_clipfrac + ) + entropy_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = entropy.mean() + ratio_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = ratio.mean() + gradient_accumulation_idx += 1 + minibatch_idx += 1 + # del everything and empty cache + # fmt: off + del ( + output, vpred_temp, logits, new_logprobs, vpred, vpredclipped, + vf_losses1, vf_losses2, vf_loss, vf_clipfrac, logprobs_diff, ratio, pg_losses, pg_losses2, pg_loss_max, + pg_loss, loss, pg_clipfrac, prob_dist, entropy, approxkl, mb_return, + mb_advantage, mb_values, mb_responses, mb_query_responses, mb_logprobs, + ) + # fmt: on + empty_cache() + with torch.no_grad(): + mean_kl = kl.sum(1).mean() + mean_entropy = (-logprobs).sum(1).mean() + mean_non_score_reward = non_score_reward.sum(1).mean() + rlhf_reward = mean_non_score_reward + scores.mean() + eps = int(self.state.episode / (time.time() - start_time)) + metrics = {} + metrics["eps"] = eps + metrics["objective/kl"] = self.accelerator.gather_for_metrics(mean_kl).mean().item() + metrics["objective/entropy"] = self.accelerator.gather_for_metrics(mean_entropy).mean().item() + metrics["objective/non_score_reward"] = ( + self.accelerator.gather_for_metrics(mean_non_score_reward).mean().item() + ) + metrics["objective/rlhf_reward"] = self.accelerator.gather_for_metrics(rlhf_reward).mean().item() + metrics["objective/scores"] = self.accelerator.gather_for_metrics(scores.mean()).mean().item() + metrics["policy/approxkl_avg"] = self.accelerator.gather_for_metrics(approxkl_stats).mean().item() + metrics["policy/clipfrac_avg"] = self.accelerator.gather_for_metrics(pg_clipfrac_stats).mean().item() + metrics["loss/policy_avg"] = self.accelerator.gather_for_metrics(pg_loss_stats).mean().item() + metrics["loss/value_avg"] = self.accelerator.gather_for_metrics(vf_loss_stats).mean().item() + metrics["val/clipfrac_avg"] = self.accelerator.gather_for_metrics(vf_clipfrac_stats).mean().item() + metrics["policy/entropy_avg"] = self.accelerator.gather_for_metrics(entropy_stats).mean().item() + metrics["val/ratio"] = self.accelerator.gather_for_metrics(ratio_stats).mean().item() + metrics["val/ratio_var"] = self.accelerator.gather_for_metrics(ratio_stats).var().item() + metrics["val/num_eos_tokens"] = (responses == processing_class.eos_token_id).sum().item() + metrics["lr"] = self.lr_scheduler.get_last_lr()[0] + metrics["episode"] = self.state.episode + self.state.epoch = self.state.episode / self.train_dataset_len # used by self.log + self.state.global_step += 1 + self.log(metrics) + + self.lr_scheduler.step() + self.control = self.callback_handler.on_step_end(args, self.state, self.control) + if self.control.should_save: + self._save_checkpoint(model, trial=None) + self.control = self.callback_handler.on_save(self.args, self.state, self.control) + del kl, mean_kl, mean_entropy, mean_non_score_reward, scores, metrics, non_score_reward + empty_cache() + gc.collect() + + if args.num_sample_generations > 0 and (update - 1) % self.sample_generations_freq == 0: + self.generate_completions(sampling=True) + empty_cache() + del ( + query_responses, + responses, + postprocessed_responses, + logprobs, + ref_logprobs, + values, + sequence_lengths, + contain_eos_token, + sequence_lengths_p1, + response_idxs, + padding_mask, + padding_mask_p1, + rewards, + actual_start, + actual_end, + advantages, + returns, + ) + empty_cache() + + # HF trainer specifics + self.control = self.callback_handler.on_train_end(args, self.state, self.control) + if self.control.should_save: + self._save_checkpoint(model, trial=None, metrics=None) + self.control = self.callback_handler.on_save(self.args, self.state, self.control) + + def generate_completions(self, sampling: bool = False): + args = self.args + processing_class = self.processing_class + generation_config = GenerationConfig( + max_new_tokens=self.args.response_length, + temperature=(0.01 + 1e-7), + top_k=0.0, + top_p=1.0, + do_sample=True, + ) + + table = defaultdict(list) + with unwrap_model_for_generation( + self.model, self.accelerator, gather_deepspeed3_params=self.args.ds3_gather_for_generation + ) as unwrapped_model: + for batch in self.eval_dataloader: + query = batch["input_ids"] + with torch.no_grad(): + context_length = query.shape[1] + query_response, _ = batch_generation( + unwrapped_model.policy, + query, + query.shape[0], + processing_class.pad_token_id, + generation_config, + ) + response = query_response[:, context_length:] + postprocessed_response = response + if self.stop_token_id is not None: # handle the edge case when stop_token_id exists but is 0 + postprocessed_response = truncate_response( + self.stop_token_id, processing_class.pad_token_id, response + ) + table["query"].extend( + gather_object(processing_class.batch_decode(query, skip_special_tokens=True)) + ) + table["model response"].extend( + gather_object(processing_class.batch_decode(postprocessed_response)) + ) + + postprocessed_query_response = torch.cat((query, postprocessed_response), 1) + _, score, _ = get_reward( + self.reward_model, postprocessed_query_response, processing_class.pad_token_id, context_length + ) + table["score"].extend(self.accelerator.gather_for_metrics(score).float().cpu().numpy()) + + if sampling: + break + df = pd.DataFrame(table) + + if self.accelerator.is_main_process: + if is_rich_available(): + print_rich_table(df.iloc[0 : 0 + 5]) + if "wandb" in args.report_to: + import wandb + + if wandb.run is not None: + wandb.log({"completions": wandb.Table(dataframe=df)}) + + if "comet_ml" in args.report_to: + log_table_to_comet_experiment( + name="completions.csv", + table=df, + ) + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) + + 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 + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + citation = textwrap.dedent("""\ + @article{mziegler2019fine-tuning, + title = {{Fine-Tuning Language Models from Human Preferences}}, + author = {Daniel M. Ziegler and Nisan Stiennon and Jeffrey Wu and Tom B. Brown and Alec Radford and Dario Amodei and Paul F. Christiano and Geoffrey Irving}, + year = 2019, + eprint = {arXiv:1909.08593} + }""") + + 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="PPO", + trainer_citation=citation, + paper_title="Fine-Tuning Language Models from Human Preferences", + paper_id="1909.08593", + ) + + model_card.save(os.path.join(self.args.output_dir, "README.md")) +class UnslothPPOTrainer(_UnslothPPOTrainer): + """ + + """ + def __init__( + self, + args, + processing_class, + model, + ref_model, + reward_model, + train_dataset, + value_model, + data_collator = None, + eval_dataset = None, + callbacks = None, + peft_config = None, + **kwargs + ): + if args is None: args = UnslothPPOConfig() + 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 + 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) + 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 + 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('ppo_trainer', other_metrics) + + super().__init__( + args = args, + processing_class = processing_class, + model = model, + ref_model = ref_model, + reward_model = reward_model, + train_dataset = train_dataset, + value_model = value_model, + data_collator = data_collator, + eval_dataset = eval_dataset, + callbacks = callbacks, + 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 diff --git a/unsloth_compiled_cache/UnslothPRMTrainer.py b/unsloth_compiled_cache/UnslothPRMTrainer.py new file mode 100644 index 0000000..d44abf0 --- /dev/null +++ b/unsloth_compiled_cache/UnslothPRMTrainer.py @@ -0,0 +1,826 @@ +""" +2025.7.1 +2025.7.1 +4.53.1 +0.19.1 +__UNSLOTH_VERSIONING__ +""" +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from trl.trainer.prm_trainer import (BaseImageProcessor, Callable, DataCollator, DataCollatorForTokenClassification, Dataset, EvalPrediction, FeatureExtractionMixin, Optional, PRMConfig, PRMTrainer, PartialState, Path, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, Trainer, TrainerCallback, Union, chain, compute_accuracy, disable_dropout_in_model, features, generate_model_card, inspect, is_peft_available, is_wandb_available, nn, os, prepare_model_for_kbit_training, textwrap, torch, warnings, Optional, PeftModel, PreTrainedModel, Trainer, is_peft_available, os, torch) + + +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 UnslothPRMConfig(PRMConfig): + """ + + Configuration class for the [`PRMTrainer`]. + + This class includes only the parameters that are specific to PRM training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may + differ from those in [`~transformers.TrainingArguments`]. + + 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: + max_length (`int` or `None`, *optional*, defaults to `1024`): + Maximum length of the sequences (prompt + completion) used for truncation. + max_prompt_length (`int` or `None`, *optional*, defaults to `512`): + Maximum length of the prompt used for truncation. + max_completion_length (`int` or `None`, *optional*, defaults to `None`): + Maximum length of the completion used for truncation. The completion is the concatenation of the steps. + disable_dropout (`bool`, *optional*, defaults to `True`): + Whether to disable dropout in the model. + step_separator (`str`, *optional*, defaults to `"\n"`): + Separator used to separate each step of the reasoning process. + train_on_last_step_only (`bool`, *optional*, defaults to `False`): + Whether to train only on the last step. + dataset_num_proc (`int`, *optional*, defaults to `None`): + Number of processes to use for processing the dataset. + + """ + 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, + hub_revision = None, + 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, + liger_kernel_config = None, + eval_use_gather_object = False, + average_tokens_across_devices = True, + max_length = 1024, + max_prompt_length = 512, + max_completion_length = None, + disable_dropout = True, + step_separator = '\ +', + train_on_last_step_only = False, + dataset_num_proc = None, + 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, + 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, + max_length = max_length, + max_prompt_length = max_prompt_length, + max_completion_length = max_completion_length, + disable_dropout = disable_dropout, + step_separator = step_separator, + train_on_last_step_only = train_on_last_step_only, + dataset_num_proc = dataset_num_proc,**kwargs) + self.vllm_sampling_params = vllm_sampling_params + self.unsloth_num_chunks = unsloth_num_chunks +pass + +class _UnslothPRMTrainer(Trainer): + """""" + + _tag_names = ["trl", "prm"] + + def __init__( + self, + model: Optional[Union[PreTrainedModel, nn.Module]] = None, + args: Optional[PRMConfig] = 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[[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, + peft_config: Optional[dict] = None, + ): + 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 not isinstance(model, PeftModel): + if getattr(model, "is_loaded_in_8bit", False) or getattr(model, "is_quantized", False): + _supports_gc_kwargs = "gradient_checkpointing_kwargs" in list( + inspect.signature(prepare_model_for_kbit_training).parameters + ) + + prepare_model_kwargs = {"use_gradient_checkpointing": args.gradient_checkpointing} + + if not _supports_gc_kwargs and args.gradient_checkpointing_kwargs is not None: + warnings.warn( + "You passed `gradient_checkpointing_kwargs` in the trainer's kwargs, but your peft version does not support it. " + "please update to the latest version of peft to use `gradient_checkpointing_kwargs`." + ) + elif _supports_gc_kwargs and args.gradient_checkpointing_kwargs is not None: + prepare_model_kwargs["gradient_checkpointing_kwargs"] = args.gradient_checkpointing_kwargs + + model = prepare_model_for_kbit_training(model, **prepare_model_kwargs) + + model = model + + # Disable dropout in the model + if args.disable_dropout: + disable_dropout_in_model(model) + + if compute_metrics is None: + compute_metrics = compute_accuracy + + if data_collator is None: + if processing_class is None: + raise ValueError( + "A processing_class must be specified when using the default DataCollatorForTokenClassification" + ) + data_collator = DataCollatorForTokenClassification(processing_class, max_length=args.max_length) + + if "input_ids" not in train_dataset.column_names: + with PartialState().main_process_first(): + fn_kwargs = { + "tokenizer": processing_class, + "step_separator": args.step_separator, + "max_length": args.max_length, + "max_prompt_length": args.max_prompt_length, + "max_completion_length": args.max_completion_length, + "train_on_last_step_only": args.train_on_last_step_only, + } + train_fn_kwargs = {**fn_kwargs, "is_eval": False} + train_dataset = train_dataset.map( + self.tokenize_row, + fn_kwargs=train_fn_kwargs, + num_proc=args.dataset_num_proc, + remove_columns=train_dataset.features, + desc="Tokenizing train dataset", + features=features.Features( # needed to avoid map to cast labels to bool + { + "labels": features.Sequence(features.Value("int64")), + "input_ids": features.Sequence(features.Value("int64")), + } + ), + ) + + eval_fn_kwargs = {**fn_kwargs, "is_eval": True} + if eval_dataset is not None: + eval_dataset = eval_dataset.map( + self.tokenize_row, + fn_kwargs=eval_fn_kwargs, + num_proc=args.dataset_num_proc, + remove_columns=eval_dataset.features, + desc="Tokenizing eval dataset", + features=features.Features( # needed to avoid map to cast labels to bool + { + "labels": features.Sequence(features.Value("int64")), + "input_ids": features.Sequence(features.Value("int64")), + } + ), + ) + + 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, + ) + + # 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) + + @staticmethod + def tokenize_row( + features, + tokenizer, + step_separator, + max_length, + max_prompt_length, + max_completion_length, + train_on_last_step_only, + is_eval, + ): + r""" + Tokenize a row of the dataset. + + Args: + features (`dict[str, str]`): + Row of the dataset, should contain the keys `"prompt"`, `"completions"`, and `"labels"`. + tokenizer (`PreTrainedTokenizerBase`): + Tokenizer used to process the data. + step_separator (`str`): + Separator between steps in the completion. + max_length (`int` or `None`): + Maximum length of the sequences (prompt + completion). If `None`, the sequences are not truncated. + max_prompt_length (`int` or `None`): + Maximum length of the prompt. If `None`, the prompt is not truncated. + max_completion_length (`int` or `None`): + Maximum length of the completion sequences. If `None`, the completion sequences are not truncated. + train_on_last_step_only (`bool`): + Whether to train only on the last step. If `True`, the labels are `-100` for all tokens except the last + token of the completion. + is_eval (`bool`): + Whether the function is used to tokenize samples from a training or an evaluation dataset. Used only if + `train_on_last_step_only` is set to `True`. + + Returns: + `dict[str, list[int]]`: + Tokenized sequences with the keys `"input_ids"`, and `"labels". + + Example: + ```python + >>> from transformers import AutoTokenizer + + >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B") + >>> features = { + ... "prompt": "Which number is larger, 9.8 or 9.11?", + ... "completions": ["11 is greater than 8.", "Hence, 9.11 > 9.8."], + ... "labels": [True, False], + ... } + >>> PRMTrainer.tokenize_row( + ... features, tokenizer, "\n", max_completion_length=None, train_on_last_step_only=False, is_eval=False + ... ) + {'input_ids': [23085, 1372, 374, 8131, 11, 220, 24, 13, 23, 476, 220, 24, 13, 16, 16, 30, 16, 16, 374, 7046, 1091, 220, 23, 13, 198, 39, 763, 11, 220, 24, 13, 16, 16, 861, 220, 24, 13, 23, 13, 198], + 'labels': [-100, -100, -100, -100, -100, -100, -100, -100, 1, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, 0]} + ``` + """ + # Tokenize the prompt and completions + prompt_ids = tokenizer(features["prompt"], add_special_tokens=False)["input_ids"] + completions_ids = [ + tokenizer(completion, add_special_tokens=False)["input_ids"] for completion in features["completions"] + ] + if train_on_last_step_only and not is_eval: + labels = [-100] * (len(features["labels"]) - 1) + [int(features["labels"][-1])] + else: + labels = [int(label) for label in features["labels"]] + + # Get the ID of the separator token and add it to the completions + separator_ids = tokenizer.encode(step_separator, add_special_tokens=False) + completions_ids = [completion + separator_ids for completion in completions_ids] + + # Create the label + labels = [[-100] * (len(completion) - 1) + [label] for completion, label in zip(completions_ids, labels)] + + # Join the completions and labels steps + completion_ids = list(chain(*completions_ids)) + labels = list(chain(*labels)) + + if tokenizer.bos_token_id is not None: + prompt_ids = [tokenizer.bos_token_id] + prompt_ids + + # Truncate prompt and completion sequences + if max_prompt_length is not None: + prompt_ids = prompt_ids[-max_prompt_length:] + if max_completion_length is not None: + completion_ids = completion_ids[:max_completion_length] + labels = labels[:max_completion_length] + + input_ids = prompt_ids + completion_ids + labels = [-100] * len(prompt_ids) + labels + + if max_length is not None: + input_ids = input_ids[:max_length] + labels = labels[:max_length] + + return {"input_ids": input_ids, "labels": labels} + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) + + 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 + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + citation = textwrap.dedent("""\ + @article{uesato2022solving, + title = {{Solving Math Word Problems With Process- and Outcome-Based Feedback}}, + author = {Uesato, Jonathan and Kushman, Nate and Kumar, Ramana and Song, Francis and Siegel, Noah and Wang, Lisa and Creswell, Antonia and Irving, Geoffrey and Higgins, Irina}, + year = 2022, + journal = {arXiv preprint arXiv:2211.14275} + }""") + + 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, + trainer_name="PRM", + trainer_citation=citation, + paper_title="Solving math word problems with process-and outcome-based feedback", + ) + + model_card.save(os.path.join(self.args.output_dir, "README.md")) +class UnslothPRMTrainer(_UnslothPRMTrainer): + """ + + Initialize PRMTrainer. + + Args: + model (`transformers.PreTrainedModel`): + The model to train, preferably an `AutoModelForTokenClassification`. + args (`PRMConfig`): + The arguments to use for training. + data_collator (`transformers.DataCollator`): + The data collator to use for training. If None is specified, the default data collator + (`DataCollatorForTokenClassification`) 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[[transformers.EvalPrediction], dict]`, *optional* defaults to `compute_accuracy`): + The metrics to use for evaluation. If no metrics are specified, the default metric (`compute_accuracy`) + will be used. + 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, + 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 = UnslothPRMConfig() + 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 + 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) + 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 + 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('prm_trainer', other_metrics) + + 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, + 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 diff --git a/unsloth_compiled_cache/UnslothRLOOTrainer.py b/unsloth_compiled_cache/UnslothRLOOTrainer.py new file mode 100644 index 0000000..e337d2c --- /dev/null +++ b/unsloth_compiled_cache/UnslothRLOOTrainer.py @@ -0,0 +1,1155 @@ +""" +2025.7.1 +2025.7.1 +4.53.1 +0.19.1 +__UNSLOTH_VERSIONING__ +""" +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from trl.trainer.rloo_trainer import (Accelerator, BaseImageProcessor, Callable, CallbackHandler, DEFAULT_CALLBACKS, DEFAULT_PROGRESS_CALLBACK, DataCollatorWithPadding, DataLoader, Dataset, ExportableState, FeatureExtractionMixin, GenerationConfig, INVALID_LOGPROB, OnlineTrainerState, Optional, Path, PreTrainedTokenizerBase, PrinterCallback, ProcessorMixin, RLOOConfig, RLOOTrainer, Trainer, TrainerCallback, TrainerControl, Union, batch_generation, broadcast, defaultdict, disable_dropout_in_model, empty_cache, exact_div, first_true_indices, forward, gather_object, gc, generate_model_card, get_comet_experiment_url, get_reporting_integration_callbacks, get_reward, is_rich_available, is_wandb_available, log_table_to_comet_experiment, math, nn, np, os, pd, prepare_deepspeed, print_rich_table, textwrap, time, torch, truncate_response, unwrap_model_for_generation, Optional, Trainer, os, torch) + + +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 UnslothRLOOConfig(RLOOConfig): + """ + + Configuration class for the [`RLOOTrainer`]. + + This class includes only the parameters that are specific to RLOO training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] and [`OnPolicyConfig`] documentation. Note that default + values in this class may differ from those in [`~transformers.TrainingArguments`]. + + 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: + exp_name (`str`, *optional*, defaults to `os.path.basename(__file__)[: -len(".py")]`): + Name of this experiment. + reward_model_path (`str`, *optional*, defaults to `"EleutherAI/pythia-160m"`): + Path to the reward model. + num_ppo_epochs (`int`, *optional*, defaults to `4`): + Number of epochs to train. + whiten_rewards (`bool`, *optional*, defaults to `False`): + Whether to whiten the rewards. + kl_coef (`float`, *optional*, defaults to `0.05`): + KL coefficient. + cliprange (`float`, *optional*, defaults to `0.2`): + Clip range. + rloo_k (`int`, *optional*, defaults to `2`): + REINFORCE Leave-One-Out (RLOO) number of online samples per prompt. + normalize_reward (`bool`, *optional*, defaults to `False`): + Whether to normalize rewards. + reward_clip_range (`float`, *optional*, defaults to `10.0`): + Clip range for rewards. + normalize_advantage (`bool`, *optional*, defaults to `False`): + Whether to normalize advantages. + token_level_kl (`bool`, *optional*, defaults to `True`): + Whether to use token-level KL penalty or sequence-level KL penalty. + ds3_gather_for_generation (`bool`, *optional*, defaults to `True`): + This setting applies to DeepSpeed ZeRO-3. If enabled, the policy model weights are gathered for generation, + improving generation speed. However, disabling this option allows training models that exceed the VRAM + capacity of a single GPU, albeit at the cost of slower generation. + + """ + 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, + hub_revision = None, + 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, + liger_kernel_config = None, + eval_use_gather_object = False, + average_tokens_across_devices = False, + dataset_num_proc = None, + num_mini_batches = 1, + total_episodes = None, + local_rollout_forward_batch_size = 64, + num_sample_generations = 10, + response_length = 53, + stop_token = None, + stop_token_id = None, + temperature = 0.7, + missing_eos_penalty = None, + sft_model_path = 'EleutherAI/pythia-160m', + world_size = None, + num_total_batches = None, + micro_batch_size = None, + local_batch_size = None, + batch_size = None, + local_mini_batch_size = None, + mini_batch_size = None, + exp_name = 'rloo_config', + reward_model_path = 'EleutherAI/pythia-160m', + num_ppo_epochs = 4, + whiten_rewards = False, + kl_coef = 0.05, + cliprange = 0.2, + rloo_k = 2, + normalize_reward = False, + reward_clip_range = 10.0, + normalize_advantage = False, + token_level_kl = False, + ds3_gather_for_generation = True, + 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() + if temperature <= 0: + raise MathError('Unsloth: Please set a positive non-zero temperature since your results will be wrong.') + elif temperature >= 10: + raise MathError('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, + 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, + 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, + dataset_num_proc = dataset_num_proc, + num_mini_batches = num_mini_batches, + total_episodes = total_episodes, + local_rollout_forward_batch_size = local_rollout_forward_batch_size, + num_sample_generations = num_sample_generations, + response_length = response_length, + stop_token = stop_token, + stop_token_id = stop_token_id, + temperature = temperature, + missing_eos_penalty = missing_eos_penalty, + sft_model_path = sft_model_path, + world_size = world_size, + num_total_batches = num_total_batches, + micro_batch_size = micro_batch_size, + local_batch_size = local_batch_size, + batch_size = batch_size, + local_mini_batch_size = local_mini_batch_size, + mini_batch_size = mini_batch_size, + exp_name = exp_name, + reward_model_path = reward_model_path, + num_ppo_epochs = num_ppo_epochs, + whiten_rewards = whiten_rewards, + kl_coef = kl_coef, + cliprange = cliprange, + rloo_k = rloo_k, + normalize_reward = normalize_reward, + reward_clip_range = reward_clip_range, + normalize_advantage = normalize_advantage, + token_level_kl = token_level_kl, + ds3_gather_for_generation = ds3_gather_for_generation,**kwargs) + self.vllm_sampling_params = vllm_sampling_params + self.unsloth_num_chunks = unsloth_num_chunks +pass + +class _UnslothRLOOTrainer(Trainer): + _tag_names = ["trl", "rloo"] + + def __init__( + self, + config: RLOOConfig, + processing_class: Optional[ + Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin] + ], + policy: nn.Module, + ref_policy: nn.Module, + reward_model: Union[nn.Module, Callable[[list[str]], list[float]]], + train_dataset: Dataset, + data_collator: Optional[DataCollatorWithPadding] = None, + eval_dataset: Optional[Union[Dataset, dict[str, Dataset]]] = None, + # less commonly used + optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), + callbacks: Optional[list[TrainerCallback]] = None, + ) -> None: + if ref_policy is policy: + raise ValueError( + "`policy` and `ref_policy` cannot be the same object. If you want `ref_policy` to be the " + "same as `policy`, you must mass a copy of it, or `None` if you use peft." + ) + + self.args = config + args = config + self.processing_class = processing_class + self.policy = policy + + # Define the collator if not provided + if data_collator is None: + data_collator = DataCollatorWithPadding(self.processing_class) + + self.policy.generation_config.eos_token_id = ( + None # disable `pad_token_id` and `eos_token_id` because we just want to + ) + self.policy.generation_config.pad_token_id = None # generate tokens without truncation / padding + + self.ref_policy = ref_policy + self.reward_model = reward_model + self.train_dataset = train_dataset + self.train_dataset_len = len(train_dataset) + self.data_collator = data_collator + self.eval_dataset = eval_dataset + self.optimizer, self.lr_scheduler = optimizers + self.optimizer_cls_and_kwargs = None # needed for transformers >= 4.47 + + ######### + # calculate various batch sizes + ######### + if args.total_episodes is None: # allow the users to define episodes in terms of epochs. + args.total_episodes = int(args.num_train_epochs * self.train_dataset_len) + accelerator = Accelerator(gradient_accumulation_steps=args.gradient_accumulation_steps) + self.accelerator = accelerator + args.world_size = accelerator.num_processes + args.local_batch_size = ( + args.per_device_train_batch_size * args.gradient_accumulation_steps * args.num_mini_batches + ) + args.micro_batch_size = int(args.per_device_train_batch_size * args.world_size) + args.batch_size = int(args.local_batch_size * args.world_size) + args.mini_batch_size = exact_div( + args.batch_size, args.num_mini_batches, "`batch_size` must be a multiple of `num_mini_batches`" + ) + args.local_mini_batch_size = exact_div( + args.local_batch_size, args.num_mini_batches, "`local_batch_size` must be a multiple of `num_mini_batches`" + ) + args.num_total_batches = math.ceil( + args.total_episodes / args.batch_size + ) # we may train for more than `total_episodes` + time_tensor = torch.tensor(int(time.time()), device=accelerator.device) + time_int = broadcast(time_tensor, 0).item() # avoid different timestamps across processes + args.run_name = f"{args.exp_name}__{args.seed}__{time_int}" + self.local_seed = args.seed + accelerator.process_index * 100003 # Prime + if args.num_sample_generations > 0: + self.sample_generations_freq = max(1, args.num_total_batches // args.num_sample_generations) + self.local_dataloader_batch_size = exact_div( + args.local_batch_size, args.rloo_k, "`local_batch_size` must be a multiple of rloo_k" + ) # RLOO logic: needed because RLOO repeats the same prompt args.rloo_k times + + ######### + # setup model, optimizer, and others + ######### + for module in [policy, ref_policy, reward_model]: + if isinstance(module, nn.Module): + disable_dropout_in_model(module) + if args.stop_token and args.stop_token == "eos": + args.stop_token_id = self.processing_class.eos_token_id + self.model = policy + self.create_optimizer_and_scheduler( + num_training_steps=args.num_total_batches + ) # note that we are calling `self.lr_scheduler.step[]` manually only at the batch level + + ######### + ### trainer specifics + ######### + default_callbacks = DEFAULT_CALLBACKS + get_reporting_integration_callbacks(self.args.report_to) + self.callbacks = default_callbacks if callbacks is None else default_callbacks + callbacks + self.callback_handler = CallbackHandler( + self.callbacks, self.model, self.processing_class, self.optimizer, self.lr_scheduler + ) + self.add_callback(PrinterCallback if self.args.disable_tqdm else DEFAULT_PROGRESS_CALLBACK) + self.control = TrainerControl() + self.state = OnlineTrainerState( + is_local_process_zero=self.is_local_process_zero(), + is_world_process_zero=self.is_world_process_zero(), + stateful_callbacks=[ + cb for cb in self.callback_handler.callbacks + [self.control] if isinstance(cb, ExportableState) + ], + ) + + self.current_flos = 0 + self.hp_search_backend = None + self.is_deepspeed_enabled = getattr(self.accelerator.state, "deepspeed_plugin", None) is not None + self.is_fsdp_enabled = getattr(self.accelerator.state, "fsdp_plugin", None) is not None + # Create distant repo and output directory if needed + self.hub_model_id = None + if self.args.push_to_hub: + self.init_hf_repo() + if self.args.should_save: + os.makedirs(self.args.output_dir, exist_ok=True) + self.backup_model = None + + # 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) + + ######### + ### setup dataloader + ######### + self.dataloader = DataLoader( + self.train_dataset, + batch_size=self.local_dataloader_batch_size, + shuffle=True, + collate_fn=self.data_collator, + drop_last=True, # needed; otherwise the last batch will be of ragged shape + ) + # sync random states for DataLoader[shuffle=True] before `accelerator.prepare` + # see https://gist.github.com/vwxyzjn/2581bff1e48e185e0b85b6dfe1def79c + torch.manual_seed(args.seed) + self.model, self.optimizer, self.dataloader = accelerator.prepare(self.model, self.optimizer, self.dataloader) + torch.manual_seed(self.local_seed) # reset the local seed again + + self.eval_dataloader = DataLoader( + self.eval_dataset, + batch_size=args.per_device_eval_batch_size, + collate_fn=self.data_collator, + drop_last=True, + ) # no need to shuffle eval dataset + self.eval_dataloader = accelerator.prepare(self.eval_dataloader) + + if self.is_deepspeed_enabled: + if isinstance(self.reward_model, nn.Module): + self.reward_model = prepare_deepspeed( + self.reward_model, args.per_device_train_batch_size, args.fp16, args.bf16 + ) + self.ref_policy = prepare_deepspeed( + self.ref_policy, args.per_device_train_batch_size, args.fp16, args.bf16 + ) + self.deepspeed = self.model + else: + self.ref_policy = self.ref_policy.to(self.accelerator.device) + if isinstance(self.reward_model, nn.Module): + self.reward_model = self.reward_model.to(self.accelerator.device) + + def get_train_dataloader(self) -> DataLoader: + return self.dataloader + + def get_eval_dataloader(self) -> DataLoader: + return self.eval_dataloader + + def train(self): + args = self.args + accelerator = self.accelerator + optimizer = self.optimizer + model = self.model + self.model_wrapped = self.model + ref_policy = self.ref_policy + reward_model = self.reward_model + processing_class = self.processing_class + dataloader = self.dataloader + device = accelerator.device + + def repeat_generator(): + while True: + yield from dataloader + + iter_dataloader = iter(repeat_generator()) + generation_config = GenerationConfig( + max_new_tokens=args.response_length, + temperature=(args.temperature + 1e-7), + top_k=0.0, + top_p=1.0, + do_sample=True, + ) + + accelerator.print("===training policy===") + start_time = time.time() + stats_shape = (args.num_ppo_epochs, args.num_mini_batches, args.gradient_accumulation_steps) + approxkl_stats = torch.zeros(stats_shape, device=device) + pg_clipfrac_stats = torch.zeros(stats_shape, device=device) + pg_loss_stats = torch.zeros(stats_shape, device=device) + vf_clipfrac_stats = torch.zeros(stats_shape, device=device) + entropy_stats = torch.zeros(stats_shape, device=device) + ratio_stats = torch.zeros(stats_shape, device=device) + model.train() + + # trainer state initialization + self.state.global_step = 0 + self.state.episode = 0 + self.state.max_steps = (args.num_total_batches * args.num_mini_batches) // 2 + self.state.num_train_epochs = args.total_episodes / self.train_dataset_len + # Compute absolute values for logging, eval, and save if given as ratio + if args.logging_steps is not None: + if args.logging_steps < 1: + self.state.logging_steps = math.ceil(self.state.max_steps * args.logging_steps) + else: + self.state.logging_steps = args.logging_steps + if args.eval_steps is not None: + if args.eval_steps < 1: + self.state.eval_steps = math.ceil(self.state.max_steps * args.eval_steps) + else: + self.state.eval_steps = args.eval_steps + if args.save_steps is not None: + if args.save_steps < 1: + self.state.save_steps = math.ceil(self.state.max_steps * args.save_steps) + else: + self.state.save_steps = args.save_steps + self.control = self.callback_handler.on_train_begin(args, self.state, self.control) + + for update in range(1, args.num_total_batches + 1): + self.state.episode += 1 * args.batch_size + data = next(iter_dataloader) + with torch.no_grad(): + queries = data["input_ids"].to(device) + queries = queries.repeat(args.rloo_k, 1) + context_length = queries.shape[1] + responses = [] + postprocessed_responses = [] + logprobs = [] + ref_logprobs = [] + scores = [] + sequence_lengths = [] + + # Generate responses and compute logprobs + with unwrap_model_for_generation( + self.model, self.accelerator, gather_deepspeed3_params=self.args.ds3_gather_for_generation + ) as unwrapped_model: + query_responses, logitss = batch_generation( + unwrapped_model, + queries, + args.local_rollout_forward_batch_size, + processing_class.pad_token_id, + generation_config, + ) + + # Process responses in batches + for i in range(0, queries.shape[0], args.local_rollout_forward_batch_size): + query = queries[i : i + args.local_rollout_forward_batch_size] + query_response = query_responses[i : i + args.local_rollout_forward_batch_size] + response = query_response[:, context_length:] + logits = logitss[i : i + args.local_rollout_forward_batch_size] + logprob = selective_log_softmax(logits, response) + del logits + empty_cache() + + ref_output = forward(ref_policy, query_response, processing_class.pad_token_id) + ref_logits = ref_output.logits[:, context_length - 1 : -1] + ref_logits /= args.temperature + 1e-7 + ref_logprob = selective_log_softmax(ref_logits, response) + del ref_output, ref_logits + empty_cache() + + # Response Processing 1. truncate response after the first occurrence of `stop_token_id` + postprocessed_response = response + if args.stop_token_id is not None: # handle the edge case when stop_token_id exists but is 0 + postprocessed_response = truncate_response( + args.stop_token_id, processing_class.pad_token_id, response + ) + + # Response Processing 2. run reward model on the truncated responses + postprocessed_query_response = torch.cat((query, postprocessed_response), 1) + sequence_length = first_true_indices(postprocessed_response == processing_class.pad_token_id) - 1 + + if isinstance(reward_model, nn.Module): + _, score, _ = get_reward( + reward_model, postprocessed_query_response, processing_class.pad_token_id, context_length + ) + else: + score = torch.tensor( + reward_model( + processing_class.batch_decode(postprocessed_query_response, skip_special_tokens=True) + ), + dtype=torch.float, + ).to(device) + + # Store batch results + responses.append(response) + postprocessed_responses.append(postprocessed_response) + logprobs.append(logprob) + ref_logprobs.append(ref_logprob) + sequence_lengths.append(sequence_length) + scores.append(score) + + # Concatenate all batched results + responses = torch.cat(responses, 0) + postprocessed_responses = torch.cat(postprocessed_responses, 0) + logprobs = torch.cat(logprobs, 0) + ref_logprobs = torch.cat(ref_logprobs, 0) + sequence_lengths = torch.cat(sequence_lengths, 0) + scores = torch.cat(scores, 0) + del (logprob, ref_logprob, score) + empty_cache() + gc.collect() + + # Response Processing 3. filter response. Ensure that the sample contains stop_token_id + # responses not passing that filter will receive a low (fixed) score + # only query humans on responses that pass that filter + contain_eos_token = torch.any(postprocessed_responses == processing_class.eos_token_id, dim=-1) + if args.missing_eos_penalty is not None: + scores[~contain_eos_token] -= self.args.missing_eos_penalty + # accelerator.print(f"{scores=}, {(contain_eos_token.sum() / len(contain_eos_token))=}") + + # be very careful with `padding_mask_p1`; see https://excalidraw.com/#json=LWnzG4w2k5DjF_EOL_xPt,e2w3a-hFJ_gX5vOfeyXGTw + response_idxs = torch.arange(responses.shape[1], device=responses.device).repeat(responses.shape[0], 1) + padding_mask = response_idxs > sequence_lengths.unsqueeze(1) + logprobs = torch.masked_fill(logprobs, padding_mask, INVALID_LOGPROB) + ref_logprobs = torch.masked_fill(ref_logprobs, padding_mask, INVALID_LOGPROB) + + # 4. compute rewards + # Compute KL divergence + kl = logprobs - ref_logprobs + + # Normalize rewards + if args.normalize_reward: + scores = (scores - scores.mean()) / (scores.std() + 1e-8) + scores = torch.clamp(scores, -args.reward_clip_range, args.reward_clip_range) + + # Compute total reward with KL penalty + if args.token_level_kl: + # Token-level KL penalty: apply KL penalty per token + kl_reward = -args.kl_coef * kl + + # Get the index of the last non-padded token for each sequence + eos_indices = padding_mask.size(1) - 1 - padding_mask.long().fliplr().argmax(dim=1, keepdim=True) + last_reward = torch.zeros_like(kl) + # Ensure scores has correct shape and type + scores_shaped = scores.reshape(-1, 1).to(kl.dtype) + last_reward.scatter_(dim=1, index=eos_indices, src=scores_shaped) + + # Combine KL reward and last reward + non_score_reward = kl_reward.sum(1) # Keep this for logging + reward = last_reward + kl_reward + rlhf_reward = reward.sum(1) # Sum across sequence length + else: + # Sequence-level KL penalty: sum KL across tokens first + sequence_kl = kl.sum(1) + non_score_reward = -args.kl_coef * sequence_kl + rlhf_reward = non_score_reward + scores + + # vectorized RLOO advantages implementation + rlhf_reward = rlhf_reward.reshape(args.rloo_k, -1) + baseline = (rlhf_reward.sum(0) - rlhf_reward) / (args.rloo_k - 1) + advantages = rlhf_reward - baseline + advantages = advantages.flatten() + + # Normalize advantages + if args.normalize_advantage: + advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8) + + empty_cache() + + # Do multiple epochs of PPO training, with a fresh random shuffle in each epoch + for ppo_epoch_idx in range(args.num_ppo_epochs): + b_inds = np.random.permutation(args.local_batch_size) + minibatch_idx = 0 + for mini_batch_start in range(0, args.local_batch_size, args.local_mini_batch_size): + mini_batch_end = mini_batch_start + args.local_mini_batch_size + mini_batch_inds = b_inds[mini_batch_start:mini_batch_end] + gradient_accumulation_idx = 0 + for micro_batch_start in range(0, args.local_mini_batch_size, args.per_device_train_batch_size): + with accelerator.accumulate(model): + micro_batch_end = micro_batch_start + args.per_device_train_batch_size + micro_batch_inds = mini_batch_inds[micro_batch_start:micro_batch_end] + + # Get batch data + mb_advantage = advantages[micro_batch_inds] + mb_responses = responses[micro_batch_inds] + mb_query_responses = query_responses[micro_batch_inds] + mb_logprobs = logprobs[micro_batch_inds] + + # Forward pass + output = forward(model, mb_query_responses, processing_class.pad_token_id) + logits = output.logits[:, context_length - 1 : -1] + logits /= args.temperature + 1e-7 + + # Compute new logprobs + new_logprobs = selective_log_softmax(logits, mb_responses) + new_logprobs = torch.masked_fill( + new_logprobs, padding_mask[micro_batch_inds], INVALID_LOGPROB + ) + + # Compute probability ratios + new_ratio = (new_logprobs - mb_logprobs).exp() + new_logprobs = new_logprobs.sum(1) + mb_logprobs = mb_logprobs.sum(1) + logprobs_diff = new_logprobs - mb_logprobs + ratio = torch.exp(logprobs_diff) + + # PPO clipped loss + pg_losses = -mb_advantage * ratio + pg_losses2 = -mb_advantage * torch.clamp(ratio, 1.0 - args.cliprange, 1.0 + args.cliprange) + pg_loss_max = torch.max(pg_losses, pg_losses2) + pg_loss = pg_loss_max.mean() + + # Final loss + loss = pg_loss + + # Optimization step + accelerator.backward(loss) + optimizer.step() + optimizer.zero_grad() + + with torch.no_grad(): + pg_clipfrac = (pg_losses2 > pg_losses).float().mean() + prob_dist = torch.nn.functional.softmax(logits, dim=-1) + entropy = torch.logsumexp(logits, dim=-1) - torch.sum(prob_dist * logits, dim=-1) + approxkl = 0.5 * (logprobs_diff**2).mean() + approxkl_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = approxkl + pg_clipfrac_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = ( + pg_clipfrac + ) + pg_loss_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = pg_loss + entropy_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = entropy.mean() + ratio_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = new_ratio.mean() + gradient_accumulation_idx += 1 + minibatch_idx += 1 + + # del everything and empty cache + # fmt: off + del ( + output, logits, new_logprobs, logprobs_diff, ratio, pg_losses, + pg_losses2, pg_loss, loss, pg_clipfrac, prob_dist, entropy, approxkl, + mb_advantage, mb_responses, mb_query_responses, mb_logprobs, + ) + # fmt: on + empty_cache() + + # Compute metrics + with torch.no_grad(): + mean_kl = kl.sum(1).mean() + mean_entropy = (-logprobs).sum(1).mean() + mean_non_score_reward = non_score_reward.mean() + eps = int(self.state.episode / (time.time() - start_time)) + metrics = {} + metrics["eps"] = eps + metrics["objective/kl"] = self.accelerator.gather_for_metrics(mean_kl).mean().item() + metrics["objective/entropy"] = self.accelerator.gather_for_metrics(mean_entropy).mean().item() + metrics["objective/non_score_reward"] = ( + self.accelerator.gather_for_metrics(mean_non_score_reward).mean().item() + ) + metrics["objective/rlhf_reward"] = self.accelerator.gather_for_metrics(rlhf_reward).mean().item() + metrics["objective/scores"] = self.accelerator.gather_for_metrics(scores.mean()).mean().item() + metrics["policy/approxkl_avg"] = self.accelerator.gather_for_metrics(approxkl_stats).mean().item() + metrics["policy/clipfrac_avg"] = self.accelerator.gather_for_metrics(pg_clipfrac_stats).mean().item() + metrics["loss/policy_avg"] = self.accelerator.gather_for_metrics(pg_loss_stats).mean().item() + metrics["val/clipfrac_avg"] = self.accelerator.gather_for_metrics(vf_clipfrac_stats).mean().item() + metrics["policy/entropy_avg"] = self.accelerator.gather_for_metrics(entropy_stats).mean().item() + metrics["val/ratio"] = self.accelerator.gather_for_metrics(ratio_stats).mean().item() + metrics["val/ratio_var"] = self.accelerator.gather_for_metrics(ratio_stats).var().item() + metrics["val/num_eos_tokens"] = (responses == processing_class.eos_token_id).sum().item() + metrics["lr"] = self.lr_scheduler.get_last_lr()[0] + metrics["episode"] = self.state.episode + self.state.epoch = self.state.episode / (args.rloo_k * self.train_dataset_len) # used by self.log + self.log(metrics) + del kl, mean_kl, mean_entropy, scores + + self.lr_scheduler.step() + self.state.global_step += 1 + self.control = self.callback_handler.on_step_end(args, self.state, self.control) + if self.control.should_save: + self._save_checkpoint(model, trial=None) + self.control = self.callback_handler.on_save(self.args, self.state, self.control) + empty_cache() + gc.collect() + + if args.num_sample_generations > 0 and (update - 1) % self.sample_generations_freq == 0: + self.generate_completions(sampling=True) + + # HF trainer specifics + self.control = self.callback_handler.on_train_end(args, self.state, self.control) + if self.control.should_save: + self._save_checkpoint(model, trial=None, metrics=None) + self.control = self.callback_handler.on_save(self.args, self.state, self.control) + + def generate_completions(self, sampling: bool = False): + args = self.args + processing_class = self.processing_class + generation_config = GenerationConfig( + max_new_tokens=self.args.response_length, + temperature=(0.01 + 1e-7), + top_k=0.0, + top_p=1.0, + do_sample=True, + ) + + table = defaultdict(list) + with unwrap_model_for_generation( + self.model, self.accelerator, gather_deepspeed3_params=self.args.ds3_gather_for_generation + ) as unwrapped_model: + for batch in self.eval_dataloader: + query = batch["input_ids"] + with torch.no_grad(): + context_length = query.shape[1] + query_response, _ = batch_generation( + unwrapped_model, + query, + query.shape[0], + processing_class.pad_token_id, + generation_config, + ) + response = query_response[:, context_length:] + postprocessed_response = response + if args.stop_token_id is not None: # handle the edge case when stop_token_id exists but is 0 + postprocessed_response = truncate_response( + args.stop_token_id, processing_class.pad_token_id, response + ) + table["query"].extend( + gather_object(processing_class.batch_decode(query, skip_special_tokens=True)) + ) + table["model response"].extend( + gather_object(processing_class.batch_decode(postprocessed_response)) + ) + + postprocessed_query_response = torch.cat((query, postprocessed_response), 1) + + if isinstance(self.reward_model, nn.Module): + _, score, _ = get_reward( + self.reward_model, + postprocessed_query_response, + processing_class.pad_token_id, + context_length, + ) + else: + score = torch.tensor( + self.reward_model( + processing_class.batch_decode(postprocessed_query_response, skip_special_tokens=True) + ), + dtype=torch.float, + ).to(postprocessed_query_response.device) + table["score"].extend(self.accelerator.gather_for_metrics(score).float().cpu().numpy()) + + if sampling: + break + df = pd.DataFrame(table) + + if self.accelerator.is_main_process: + if is_rich_available(): + print_rich_table(df.iloc[0 : 0 + 5]) + if "wandb" in args.report_to: + import wandb + + if wandb.run is not None: + wandb.log({"completions": wandb.Table(dataframe=df)}) + + if "comet_ml" in args.report_to: + log_table_to_comet_experiment( + name="completions.csv", + table=df, + ) + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) + + 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 + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + citation = textwrap.dedent("""\ + @inproceedings{ahmadian2024back, + title = {{Back to Basics: Revisiting REINFORCE-Style Optimization for Learning from Human Feedback in LLMs}}, + author = {Arash Ahmadian and Chris Cremer and Matthias Gall{\'{e}} and Marzieh Fadaee and Julia Kreutzer and Olivier Pietquin and Ahmet {\"{U}}st{\"{u}}n and Sara Hooker}, + year = 2024, + booktitle = {Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), {ACL} 2024, Bangkok, Thailand, August 11-16, 2024}, + publisher = {Association for Computational Linguistics}, + pages = {12248--12267}, + editor = {Lun{-}Wei Ku and Andre Martins and Vivek Srikumar}, + }""") + + 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="RLOO", + trainer_citation=citation, + paper_title="Back to Basics: Revisiting REINFORCE-Style Optimization for Learning from Human Feedback in LLMs", + paper_id="2402.14740", + ) + + model_card.save(os.path.join(self.args.output_dir, "README.md")) +class UnslothRLOOTrainer(_UnslothRLOOTrainer): + """ + + """ + def __init__( + self, + config, + processing_class, + policy, + ref_policy, + reward_model, + train_dataset, + data_collator = None, + eval_dataset = None, + callbacks = None, + **kwargs + ): + if args is None: args = UnslothRLOOConfig() + _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('rloo_trainer', other_metrics) + + super().__init__( + config = config, + processing_class = processing_class, + policy = policy, + ref_policy = ref_policy, + reward_model = reward_model, + train_dataset = train_dataset, + data_collator = data_collator, + eval_dataset = eval_dataset, + callbacks = callbacks,**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 diff --git a/unsloth_compiled_cache/UnslothRewardTrainer.py b/unsloth_compiled_cache/UnslothRewardTrainer.py new file mode 100644 index 0000000..5fbbb18 --- /dev/null +++ b/unsloth_compiled_cache/UnslothRewardTrainer.py @@ -0,0 +1,844 @@ +""" +2025.7.1 +2025.7.1 +4.53.1 +0.19.1 +__UNSLOTH_VERSIONING__ +""" +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from trl.trainer.reward_trainer import (Any, BaseImageProcessor, Callable, DataCollator, Dataset, EvalPrediction, FeatureExtractionMixin, FrozenInstanceError, Optional, PartialState, Path, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, RewardConfig, RewardDataCollatorWithPadding, RewardTrainer, Trainer, TrainerCallback, Union, _tokenize, compute_accuracy, decode_and_strip_padding, defaultdict, disable_dropout_in_model, gather_object, generate_model_card, get_comet_experiment_url, inspect, is_peft_available, is_rich_available, is_wandb_available, log_table_to_comet_experiment, maybe_apply_chat_template, nested_detach, nn, os, pd, prepare_model_for_kbit_training, print_rich_table, replace, torch, warnings, Optional, PeftModel, PreTrainedModel, Trainer, is_peft_available, os, torch) + + +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 UnslothRewardConfig(RewardConfig): + """ + + Configuration class for the [`RewardTrainer`]. + + This class includes only the parameters that are specific to Reward training. For a full list of training + arguments, please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this + class may differ from those in [`~transformers.TrainingArguments`]. + + 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: + max_length (`int` or `None`, *optional*, defaults to `1024`): + Maximum length of the sequences (prompt + completion) in the batch, filters out entries that exceed the + limit. This argument is required if you want to use the default data collator. + disable_dropout (`bool`, *optional*, defaults to `True`): + Whether to disable dropout in the model. + dataset_num_proc (`int`, *optional*, defaults to `None`): + Number of processes to use for processing the dataset. + center_rewards_coefficient (`float`, *optional*, defaults to `None`): + Coefficient to incentivize the reward model to output mean-zero rewards (proposed by + https://huggingface.co/papers/2312.09244, Eq. 2). Recommended value: `0.01`. + remove_unused_columns (`bool`, *optional*, defaults to `False`): + Whether to remove the columns that are not used by the model's forward pass. Can be `True` only if the + dataset is pretokenized. + + """ + 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 = False, + 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, + hub_revision = None, + 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, + liger_kernel_config = None, + eval_use_gather_object = False, + average_tokens_across_devices = True, + max_length = 1024, + disable_dropout = True, + dataset_num_proc = None, + center_rewards_coefficient = None, + 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, + 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, + max_length = max_length, + disable_dropout = disable_dropout, + dataset_num_proc = dataset_num_proc, + center_rewards_coefficient = center_rewards_coefficient,**kwargs) + self.vllm_sampling_params = vllm_sampling_params + self.unsloth_num_chunks = unsloth_num_chunks +pass + +class _UnslothRewardTrainer(Trainer): + _tag_names = ["trl", "reward-trainer"] + + def __init__( + self, + model: Optional[Union[PreTrainedModel, nn.Module]] = None, + args: Optional[RewardConfig] = 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[[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, + peft_config: Optional[dict] = None, + ): + """ + Initialize RewardTrainer. + + Args: + model (`transformers.PreTrainedModel`): + The model to train, preferably an `AutoModelForSequenceClassification`. + args (`RewardConfig`): + The arguments to use for training. + data_collator (`transformers.DataCollator`): + The data collator to use for training. If None is specified, the default data collator + (`RewardDataCollatorWithPadding`) 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[[transformers.EvalPrediction], dict]`, *optional* defaults to `compute_accuracy`): + The metrics to use for evaluation. If no metrics are specified, the default metric (`compute_accuracy`) + will be used. + 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. + """ + 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 not isinstance(model, PeftModel): + if getattr(model, "is_loaded_in_8bit", False) or getattr(model, "is_quantized", False): + _supports_gc_kwargs = "gradient_checkpointing_kwargs" in list( + inspect.signature(prepare_model_for_kbit_training).parameters + ) + + prepare_model_kwargs = {"use_gradient_checkpointing": args.gradient_checkpointing} + + if not _supports_gc_kwargs and args.gradient_checkpointing_kwargs is not None: + warnings.warn( + "You passed `gradient_checkpointing_kwargs` in the trainer's kwargs, but your peft version does not support it. " + "please update to the latest version of peft to use `gradient_checkpointing_kwargs`.", + UserWarning, + ) + elif _supports_gc_kwargs and args.gradient_checkpointing_kwargs is not None: + prepare_model_kwargs["gradient_checkpointing_kwargs"] = args.gradient_checkpointing_kwargs + + model = prepare_model_for_kbit_training(model, **prepare_model_kwargs) + + model = model + + # Disable dropout in the model + if args.disable_dropout: + disable_dropout_in_model(model) + + if compute_metrics is None: + compute_metrics = compute_accuracy + + if data_collator is None: + if processing_class is None: + raise ValueError( + "A processing_class must be specified when using the default RewardDataCollatorWithPadding" + ) + + max_length = args.max_length + + data_collator = RewardDataCollatorWithPadding(processing_class) + + if args.remove_unused_columns: + try: # for bc before https://github.com/huggingface/transformers/pull/25435 + args.remove_unused_columns = False + except FrozenInstanceError: + args = replace(args, remove_unused_columns=False) + # warn users + warnings.warn( + "When using RewardDataCollatorWithPadding, you should set `remove_unused_columns=False` in your RewardConfig" + " we have set it for you, but you should do it yourself in the future.", + UserWarning, + ) + + self.use_reward_data_collator = True + else: + self.use_reward_data_collator = False + + # 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 Reward, the sampled data does not include the + # "input_ids" key. Instead, the available keys are "input_ids_chosen" and "input_ids_rejected". 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 + + if "input_ids_chosen" not in train_dataset.column_names: + with PartialState().main_process_first(): + fn_kwargs = {"tokenizer": processing_class} + train_dataset = train_dataset.map(maybe_apply_chat_template, fn_kwargs={"tokenizer": processing_class}) + train_dataset = train_dataset.map( + _tokenize, + batched=True, + fn_kwargs=fn_kwargs, + num_proc=args.dataset_num_proc, + ) + # This filter is important because otherwise you get samples that exceed the model's context length and + # get truncated => noisy signal the chosen/rejected label gets lost. The downside is that the + # user might get surprised if N samples are missing from training. + train_dataset = train_dataset.filter( + lambda x: len(x["input_ids_chosen"]) <= max_length and len(x["input_ids_rejected"]) <= max_length, + num_proc=args.dataset_num_proc, + ) + if eval_dataset is not None: + eval_dataset = eval_dataset.map( + maybe_apply_chat_template, fn_kwargs={"tokenizer": processing_class} + ) + eval_dataset = eval_dataset.map( + _tokenize, + fn_kwargs=fn_kwargs, + batched=True, + num_proc=args.dataset_num_proc, + ) + # This filter is important because otherwise you get samples that exceed the model's context length and + # get truncated => noisy signal the chosen/rejected label gets lost. The downside is that the + # user might get surprised if N samples are missing from training. + eval_dataset = eval_dataset.filter( + lambda x: len(x["input_ids_chosen"]) <= max_length + and len(x["input_ids_rejected"]) <= max_length, + num_proc=args.dataset_num_proc, + ) + + 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, + ) + + # 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) + + 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]]]: + rewards_chosen = model( + input_ids=inputs["input_ids_chosen"], + attention_mask=inputs["attention_mask_chosen"], + return_dict=True, + )["logits"] + rewards_rejected = model( + input_ids=inputs["input_ids_rejected"], + attention_mask=inputs["attention_mask_rejected"], + return_dict=True, + )["logits"] + # calculate loss, optionally modulate with margin + if "margin" in inputs: + loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected - inputs["margin"]).mean() + else: + loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected).mean() + + if self.args.center_rewards_coefficient is not None: + loss += self.args.center_rewards_coefficient * torch.mean((rewards_chosen + rewards_rejected) ** 2) + + if return_outputs: + return loss, { + "rewards_chosen": rewards_chosen, + "rewards_rejected": rewards_rejected, + } + return loss + + 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, + ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]: + inputs = self._prepare_inputs(inputs) + if ignore_keys is None: + if hasattr(self.model, "config"): + ignore_keys = getattr(self.model.config, "keys_to_ignore_at_inference", []) + else: + ignore_keys = [] + + with torch.no_grad(): + loss, logits_dict = self.compute_loss(model, inputs, return_outputs=True) + + if prediction_loss_only: + return (loss, None, None) + + loss = loss.detach() + logits = tuple(v for k, v in logits_dict.items() if k not in ignore_keys) + logits = nested_detach(logits) + # Stack accepted against rejected, mean over logits + # and softmax to get preferences between accepted and rejected to sum to 1 + logits = torch.stack(logits).mean(dim=2).softmax(dim=0).T + + labels = torch.zeros(logits.shape[0]) + labels = self._prepare_inputs(labels) + + return loss, logits, labels + + def evaluate(self, *args, **kwargs): + num_print_samples = kwargs.pop("num_print_samples", 4) + self.visualize_samples(num_print_samples) + return super().evaluate(*args, **kwargs) + + def visualize_samples(self, num_print_samples: int): + """ + Visualize the reward model logits prediction + + Args: + num_print_samples (`int`, defaults to `4`): + The number of samples to print. Set to `-1` to print all samples. + """ + eval_dataloader = self.get_eval_dataloader() + table = defaultdict(list) + for _, inputs in enumerate(eval_dataloader): + _, logits, _ = self.prediction_step(self.model, inputs, prediction_loss_only=False) + chosen_text = decode_and_strip_padding(inputs["input_ids_chosen"], self.processing_class) + rejected_text = decode_and_strip_padding(inputs["input_ids_rejected"], self.processing_class) + table["chosen_text"].extend(gather_object(chosen_text)) + table["rejected_text"].extend(gather_object(rejected_text)) + table["logits"].extend( + gather_object([[round(inner_item, 4) for inner_item in item] for item in logits.tolist()]) + ) + if num_print_samples >= 0 and len(table["chosen_text"]) >= num_print_samples: + break + df = pd.DataFrame(table) + if self.accelerator.process_index == 0: + if is_rich_available(): + print_rich_table(df[:num_print_samples]) + if "wandb" in self.args.report_to: + import wandb + + if wandb.run is not None: + wandb.log({"completions": wandb.Table(dataframe=df)}) + + if "comet_ml" in self.args.report_to: + log_table_to_comet_experiment( + name="completions.csv", + table=df, + ) + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) + + 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 + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + 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="Reward", + ) + + model_card.save(os.path.join(self.args.output_dir, "README.md")) +class UnslothRewardTrainer(_UnslothRewardTrainer): + """ + + """ + def __init__( + self, + 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 = UnslothRewardConfig() + 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 + 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) + 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 + 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('reward_trainer', other_metrics) + + 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, + 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 diff --git a/unsloth_compiled_cache/UnslothSFTTrainer.py b/unsloth_compiled_cache/UnslothSFTTrainer.py new file mode 100644 index 0000000..f24c9ca --- /dev/null +++ b/unsloth_compiled_cache/UnslothSFTTrainer.py @@ -0,0 +1,1206 @@ +""" +2025.7.1 +2025.7.1 +4.53.1 +0.19.1 +__UNSLOTH_VERSIONING__ +""" +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from trl.trainer.sft_trainer import (Any, AutoModelForCausalLM, AutoTokenizer, BaseImageProcessor, Callable, ConstantLengthDataset, DataCollator, DataCollatorForLanguageModeling, Dataset, EvalPrediction, FeatureExtractionMixin, IterableDataset, Optional, Path, PeftConfig, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, SFTConfig, SFTTrainer, Trainer, TrainerCallback, TrainingArguments, Union, clone_chat_template, contextlib, dataclass, dataclasses, defaultdict, generate_model_card, get_act_offloading_ctx_manager, get_comet_experiment_url, get_peft_model, is_conversational, is_peft_available, is_wandb_available, nn, os, pad, peft, peft_module_casting_to_bf16, prepare_model_for_kbit_training, torch, version, warnings, Callable, ConstantLengthDataset, DataCollator, DataCollatorForLanguageModeling, Dataset, IterableDataset, Optional, Union, os, pad, Optional, PeftModel, PreTrainedModel, Trainer, is_peft_available, os, peft, torch, os) + + +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 UnslothSFTConfig(SFTConfig): + """ + + Configuration class for the [`SFTTrainer`]. + + This class includes only the parameters that are specific to SFT training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may + differ from those in [`~transformers.TrainingArguments`]. + + 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 + + model_init_kwargs (`dict[str, Any]` or `None`, *optional*, defaults to `None`): + Keyword arguments for [`~transformers.AutoModelForCausalLM.from_pretrained`], used when the `model` + argument of the [`SFTTrainer`] is provided as a string. + chat_template_path (`str` or `None`, *optional*, defaults to `None`): + If specified, sets the model's chat template. This can either be the path to a tokenizer (local directory + or Hugging Face Hub model) or a direct path to a Jinja template file. When using a Jinja file, you must + ensure that any special tokens referenced in the template are added to the tokenizer and that the model's + embedding layer is resized accordingly. + + > Parameters that control the data preprocessing + + dataset_text_field (`str`, *optional*, defaults to `"text"`): + Name of the column that contains text data in the dataset. + dataset_kwargs (`dict[str, Any]` or `None`, *optional*, defaults to `None`): + Dictionary of optional keyword arguments for the dataset preparation. The only supported key is + `skip_prepare_dataset`. + dataset_num_proc (`int` or `None`, *optional*, defaults to `None`): + Number of processes to use for processing the dataset. + eos_token (`str` or `None`, *optional*, defaults to `None`): + Token used to indicate the end of a turn or sequence. If `None`, it defaults to + `processing_class.eos_token`. + pad_token (`int` or `None`, *optional*, defaults to `None`): + Token used for padding. If `None`, it defaults to `processing_class.pad_token`, or if that is also `None`, + it falls back to `processing_class.eos_token`. + max_length (`int` or `None`, *optional*, defaults to `1024`): + Maximum length of the tokenized sequence. Sequences longer than `max_length` are truncated from the right. + If `None`, no truncation is applied. When packing is enabled, this value sets the sequence length. + packing (`bool`, *optional*, defaults to `False`): + Whether to group multiple sequences into fixed-length blocks to improve computational efficiency and reduce + padding. Uses `max_length` to define sequence length. + packing_strategy (`str`, *optional*, defaults to `"ffd"`): + Strategy for packing sequences. Can be either `"ffd"` (first-fit decreasing, default), or `"wrapped"`. + 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. When packing is enabled with strategy `"ffd"`, padding-free is enabled, regardless of the + value of this parameter. + pad_to_multiple_of (`int` or `None`, *optional*, defaults to `None`): + If set, the sequences will be padded to a multiple of this value. + eval_packing (`bool` or `None`, *optional*, defaults to `None`): + Whether to pack the eval dataset. If `None`, uses the same value as `packing`. + + > Parameters that control the training + + completion_only_loss (`bool` or `None`, *optional*, defaults to `None`): + Whether to compute loss only on the completion part of the sequence. If set to `True`, loss is computed + only on the completion, which is supported only for [prompt-completion](#prompt-completion) datasets. If + `False`, loss is computed on the entire sequence. If `None` (default), the behavior depends on the dataset: + loss is computed on the completion for [prompt-completion](#prompt-completion) datasets, and on the full + sequence for [language modeling](#language-modeling) datasets. + assistant_only_loss (`bool`, *optional*, defaults to `False`): + Whether to compute loss only on the assistant part of the sequence. If set to `True`, loss is computed + only on the assistant responses, which is supported only for [conversational](#conversational) datasets. If `False`, + loss is computed on the entire sequence. + activation_offloading (`bool`, *optional*, defaults to `False`): + Whether to offload the activations to the CPU. + + """ + 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, + hub_revision = None, + 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, + liger_kernel_config = None, + eval_use_gather_object = False, + average_tokens_across_devices = True, + model_init_kwargs = None, + chat_template_path = None, + dataset_text_field = 'text', + dataset_kwargs = None, + dataset_num_proc = None, + eos_token = None, + pad_token = None, + max_length = 1024, + packing = False, + packing_strategy = 'ffd', + padding_free = False, + pad_to_multiple_of = None, + eval_packing = None, + completion_only_loss = None, + assistant_only_loss = False, + activation_offloading = False, + max_seq_length = None, + 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, + 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, + model_init_kwargs = model_init_kwargs, + chat_template_path = chat_template_path, + dataset_text_field = dataset_text_field, + dataset_kwargs = dataset_kwargs, + dataset_num_proc = dataset_num_proc, + eos_token = eos_token, + pad_token = pad_token, + max_length = max_length, + packing = packing, + packing_strategy = packing_strategy, + padding_free = padding_free, + pad_to_multiple_of = pad_to_multiple_of, + eval_packing = eval_packing, + completion_only_loss = completion_only_loss, + assistant_only_loss = assistant_only_loss, + activation_offloading = activation_offloading, + max_seq_length = max_seq_length,**kwargs) + self.vllm_sampling_params = vllm_sampling_params + self.unsloth_num_chunks = unsloth_num_chunks +pass + +class _UnslothSFTTrainer(Trainer): + """""" + + _tag_names = ["trl", "sft"] + + def __init__( + self, + model: Union[str, nn.Module, PreTrainedModel], + args: Optional[Union[SFTConfig, TrainingArguments]] = None, + data_collator: Optional[DataCollator] = None, # type: ignore + 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, + compute_loss_func: Optional[Callable] = None, + compute_metrics: Optional[Callable[[EvalPrediction], dict]] = None, + callbacks: Optional[list[TrainerCallback]] = None, + optimizers: tuple[Optional[torch.optim.Optimizer], Optional[torch.optim.lr_scheduler.LambdaLR]] = (None, None), + optimizer_cls_and_kwargs: Optional[tuple[type[torch.optim.Optimizer], dict[str, Any]]] = None, + preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None, + peft_config: Optional["PeftConfig"] = None, + formatting_func: Optional[Callable[[dict], str]] = None, + ): + # Args + model_id = model if isinstance(model, str) else model.config._name_or_path + if args is None: + model_name = model_id.split("/")[-1] + args = SFTConfig(f"{model_name}-SFT") + elif isinstance(args, TrainingArguments) and not isinstance(args, SFTConfig): + dict_args = args.to_dict() + dict_args["hub_token"] = args.hub_token # to_dict hides the hub_token + dict_args.pop("push_to_hub_token") + args = SFTConfig(**dict_args) + + # Handle the tokenizer + if processing_class is None: + processing_class = AutoTokenizer.from_pretrained(model_id) + + if args.eos_token is not None: + eos_token = args.eos_token + eos_token_id = processing_class.convert_tokens_to_ids(eos_token) + if eos_token_id is None: + raise ValueError( + f"The specified `eos_token` ('{eos_token}') is not found in the vocabulary of the given " + f"`processing_class` ({processing_class.__class__.__name__}). Ensure that the `eos_token` exists " + "in the vocabulary before using it as an EOS token." + ) + processing_class.eos_token_id = eos_token_id + + # Model + if args.model_init_kwargs is not None and not isinstance(model, str): + warnings.warn( + "You passed model_init_kwargs to the `SFTConfig`, but your model is already instantiated. " + "The `model_init_kwargs` will be ignored." + ) + if isinstance(model, str): + model = self._create_model_from_path(model, args) + + if args.chat_template_path is not None: + if os.path.isfile(args.chat_template_path) and args.chat_template_path.endswith((".jinja", ".j2")): + with open(args.chat_template_path, encoding="utf-8") as chat_template_file: + processing_class.chat_template = chat_template_file.read() + else: + model, processing_class = clone_chat_template(model, processing_class, args.chat_template_path) + + # PEFT configuration and model wrapping + if False: + model = self._prepare_peft_model(model, peft_config, args) + + # Data collator + # FFD packing requires padding-free mode; otherwise, the collator outputs padded attention masks, causing + # FlashAttention to ignore position_ids and recompute them incorrectly from the padded attention mask. + self.padding_free = args.padding_free or (args.packing and args.packing_strategy == "ffd") + if self.padding_free: + if data_collator is not None: + raise ValueError("Passing a custom data collator is not supported when using padding-free.") + if args.packing and args.packing_strategy == "wrapped": + warnings.warn( + "You are passing `padding_free=True` with the 'wrapped' packing strategy, which is not " + "recommended. Please refer to the documentation to understand why this is not recommended." + ) + 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." + ) + if args.per_device_train_batch_size == 1 and not args.packing: + warnings.warn( + "You are using a per_device_train_batch_size of 1 with padding-free training. Using a batch size " + "of 1 anihilate the benefits of padding-free training. Please consider increasing the batch size " + "to at least 2." + ) + + if args.completion_only_loss is None: + first_example = next(iter(train_dataset)) + self.completion_only_loss = "prompt" in first_example + else: + self.completion_only_loss = args.completion_only_loss + + if data_collator is None: + # Get the pad token: if not provided, use the one from the processing class or the eos token + # if the processing class does not have a pad token. + pad_token = args.pad_token or processing_class.pad_token or processing_class.eos_token + pad_token_id = processing_class.convert_tokens_to_ids(pad_token) + if pad_token_id is None: + raise ValueError( + f"The specified `pad_token` ('{pad_token}') is not found in the vocabulary of the given " + f"`processing_class` ({processing_class.__class__.__name__}). Ensure that the `pad_token` exists " + "in the vocabulary before using it as a padding token." + ) + data_collator = DataCollatorForLanguageModeling( + pad_token_id=pad_token_id, + completion_only_loss=self.completion_only_loss, + padding_free=self.padding_free, + # Using position_ids without flash_attn hurts the training + return_position_ids=model.config._attn_implementation == "flash_attention_2", + pad_to_multiple_of=args.pad_to_multiple_of, + ) + + if ( + args.packing + and args.packing_strategy == "ffd" + and model.config._attn_implementation != "flash_attention_2" + ): + warnings.warn( + "You are using packing, but the attention implementation is not set to 'flash_attention_2'. Packing " + "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 cross-contamination " + "between batches. To avoid this, either disable packing by setting `packing=False`, or set " + "`attn_implementation='flash_attention_2'` in the model configuration." + ) + if args.assistant_only_loss and not is_conversational(train_dataset[0]): + raise ValueError( + "You set `assistant_only_loss=True`, but the dataset is not conversational. This option is only " + "supported for conversational datasets." + ) + + # Dataset + preprocess_dataset = args.dataset_kwargs is None or not args.dataset_kwargs.get("skip_prepare_dataset", False) + if preprocess_dataset: + if self.completion_only_loss and formatting_func: + raise ValueError( + "A formatting function was provided while `completion_only_loss=True`, which is incompatible. " + "Using a formatter converts the dataset to a language modeling type, conflicting with " + "completion-only loss. To resolve this, apply your formatting function before passing the " + "dataset, or disable `completion_only_loss` in `SFTConfig`." + ) + train_dataset = self._prepare_dataset( + train_dataset, processing_class, args, args.packing, formatting_func, "train" + ) + if eval_dataset is not None: + packing = args.packing if args.eval_packing is None else args.eval_packing + if isinstance(eval_dataset, dict): + eval_dataset = { + key: self._prepare_dataset(dataset, processing_class, args, packing, formatting_func, key) + for key, dataset in eval_dataset.items() + } + else: + eval_dataset = self._prepare_dataset( + eval_dataset, processing_class, args, packing, formatting_func, "eval" + ) + + # Initialize the metrics + self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)} + self._total_train_tokens = 0 + + # Initialize the Trainer. Parent class will handle: + # - DeepSpeed configuration [through create_accelerator_and_postprocess] + # - FSDP setup + # - Distributed training setup + # - Optimizer and scheduler creation + + super().__init__( + model=model, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + compute_loss_func=compute_loss_func, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + optimizer_cls_and_kwargs=optimizer_cls_and_kwargs, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + ) + + # Initialize activation offloading context + if self.args.activation_offloading: + self.maybe_activation_offload_context = get_act_offloading_ctx_manager(model=self.model) + else: + self.maybe_activation_offload_context = contextlib.nullcontext() + + # 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) + + def _create_model_from_path(self, model_path: str, args: SFTConfig) -> PreTrainedModel: + """Creates a model from a path or model identifier.""" + model_init_kwargs = args.model_init_kwargs or {} + # Handle torch dtype + torch_dtype = model_init_kwargs.get("torch_dtype") + if isinstance(torch_dtype, torch.dtype) or torch_dtype == "auto" or torch_dtype is None: + pass # torch_dtype is already a torch.dtype or "auto" or None + elif isinstance(torch_dtype, str): # it's a str, but not "auto" + torch_dtype = getattr(torch, torch_dtype) + model_init_kwargs["torch_dtype"] = torch_dtype + else: + raise ValueError( + "Invalid `torch_dtype` passed to `SFTConfig`. Expected either 'auto' or a string representing " + f"a `torch.dtype` (e.g., 'float32'), but got {torch_dtype}." + ) + # Disable caching if gradient checkpointing is enabled (not supported) + # if args.gradient_checkpointing: + # model_init_kwargs["use_cache"] = False + + # Create model + model = AutoModelForCausalLM.from_pretrained(model_path, **model_init_kwargs) + return model + + def _prepare_peft_model(self, model: PreTrainedModel, peft_config: Any, args: SFTConfig) -> PreTrainedModel: + """Prepares a model for PEFT training.""" + if not is_peft_available(): + raise ImportError("To use PeftModel, you need to install the `peft` library.") + + if not isinstance(peft_config, PeftConfig): + raise ValueError( + f"Expected PeftConfig object but got {type(peft_config)}. If you want to use the PeftModel, you need " + "to pass a PeftConfig object to the SFTTrainer." + ) + + if isinstance(model, PeftModel): + return model + + # Handle quantized models (QLoRA) + is_qlora = getattr(model, "is_loaded_in_4bit", False) or getattr(model, "is_loaded_in_8bit", False) + + is_sharded_qlora = False + if getattr(model, "is_loaded_in_4bit", False): + # Check if model is sharded (FSDP/DS-Zero3) + for _, param in model.named_parameters(): + if param.__class__.__name__ == "Params4bit": + is_sharded_qlora = param.data.device.type in {"cpu", "meta"} + break + + # Prepare model for kbit training if needed + if is_qlora and not is_sharded_qlora: + model = self._prepare_model_for_kbit_training(model, args) + # Disable gradient checkpointing as it's handled by prepare_model_for_kbit_training + args = dataclasses.replace(args, gradient_checkpointing=False) + elif args.gradient_checkpointing: + model = self._enable_gradient_checkpointing(model, args) + + # Create PEFT model + if ( + version.parse(peft.__version__) >= version.parse("0.12") # autocast_adapter_dtype introduced in 0.12 + and getattr(model, "is_loaded_in_4bit", False) + and is_sharded_qlora + ): + model = get_peft_model(model, peft_config, autocast_adapter_dtype=False) + else: + model = get_peft_model(model, peft_config) + + # Handle bf16 casting for 4-bit models + if args.bf16 and getattr(model, "is_loaded_in_4bit", False) and not is_sharded_qlora: + peft_module_casting_to_bf16(model) + + return model + + def _prepare_model_for_kbit_training(self, model: PreTrainedModel, args: SFTConfig) -> PreTrainedModel: + """Prepares a quantized model for kbit training.""" + prepare_model_kwargs = { + "use_gradient_checkpointing": args.gradient_checkpointing, + "gradient_checkpointing_kwargs": args.gradient_checkpointing_kwargs or {}, + } + + return prepare_model_for_kbit_training(model, **prepare_model_kwargs) + + def _enable_gradient_checkpointing(self, model: PreTrainedModel, args: SFTConfig) -> PreTrainedModel: + """Enables gradient checkpointing for the model.""" + gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs or {} + use_reentrant = ( + "use_reentrant" not in gradient_checkpointing_kwargs or gradient_checkpointing_kwargs["use_reentrant"] + ) + + if use_reentrant: + 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) + + return model + + def _prepare_dataset( + self, + dataset: Union[Dataset, IterableDataset], + processing_class, + args, + packing: bool, + formatting_func: Optional[Callable[[dict], str]], + dataset_name: str, + ) -> Union[Dataset, IterableDataset]: + # All Unsloth Zoo code licensed under LGPLv3 + if isinstance(dataset, ConstantLengthDataset): return dataset + + map_kwargs = {} + use_desc = isinstance(dataset, Dataset) + is_vlm = hasattr(processing_class, "tokenizer") + tokenizer = processing_class + if is_vlm: tokenizer = processing_class.tokenizer + + # Get max length + max_seq_length = getattr(args, "max_length", 0) + if max_seq_length == 0: max_seq_length = getattr(args, "max_seq_length", 0) + if max_seq_length == 0: max_seq_length = getattr(self, "max_seq_length", 0) + if max_seq_length == 0: max_seq_length = getattr(self, "max_seq", 0) + if max_seq_length == 0: raise RuntimeError("Unsloth: max_seq_length is 0! Please specify one!") + dataset_text_field = getattr(args, "dataset_text_field", "text") + do_truncation = max_seq_length != 0 + do_formatting_func = False + do_tokenize = True + + # Get correct column names + column_names = set(next(iter(dataset)).keys()) + used_column_names = ["input_ids"] + if "attention_mask" in column_names: + used_column_names.append("attention_mask") + + # Check if already tokenized so skip + from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling + if "labels" in column_names: + # Most likely forgot data collator! + if is_vlm and not hasattr(tokenizer, "pad"): + # Check if processing_class has a .pad, if not, use tokenizer.tokenizer + raise RuntimeError(f"Unsloth: {processing_class.__class__} does not have .pad!") + self.data_collator = DataCollatorForSeq2Seq(tokenizer) + used_column_names.append("labels") + do_tokenize = False + elif "input_ids" in column_names: + # Skip dataset prep, and set data collator + if is_vlm and not hasattr(tokenizer, "pad"): + # Check if processing_class has a .pad, if not, use tokenizer.tokenizer + raise RuntimeError(f"Unsloth: {processing_class.__class__} does not have .pad!") + self.data_collator = DataCollatorForLanguageModeling(tokenizer, mlm = False) + do_tokenize = False + elif dataset_text_field not in column_names: + do_formatting_func = True + if formatting_func is None: + raise RuntimeError("Unsloth: You must specify a `formatting_func`") + pass + + if do_tokenize: + # Check double BOS tokens + if do_formatting_func: + test_text = formatting_func(next(iter(dataset))) + if not isinstance(test_text, list): + raise ValueError( + "Unsloth: The `formatting_func` should return a list of processed strings." + ) + test_text = test_text[0] + else: + test_text = next(iter(dataset))[dataset_text_field][0] + + # Get chat template + chat_template = getattr(processing_class, 'chat_template', '') + if chat_template == '' and is_vlm: + chat_template = getattr(tokenizer, 'chat_template', '') + if chat_template is None: + chat_template = '' + + # Get bos_token + add_special_tokens = True + bos_token_1 = getattr(processing_class, 'bos_token', None) + bos_token_2 = getattr(tokenizer, 'bos_token', None) + bos_token = bos_token_1 or bos_token_2 + + if bos_token is not None: + if test_text.startswith(bos_token) or bos_token in chat_template: + add_special_tokens = False + print("Unsloth: We found double BOS tokens - we shall remove one automatically.") + pass + + # Create tokenize function + def _tokenize(example): + return tokenizer( + example[dataset_text_field] if not do_formatting_func else formatting_func(example), + truncation = do_truncation, + max_length = max_seq_length, + return_token_type_ids = False, + add_special_tokens = add_special_tokens, + ) + pass + + if not isinstance(dataset, IterableDataset): + map_kwargs["num_proc"] = getattr(args, "dataset_num_proc", 2) + else: + map_kwargs["batch_size"] = dataset._ex_iterable.batch_size + + if use_desc: map_kwargs["desc"] = f'Unsloth: Tokenizing ["{dataset_text_field}"]' + dataset = dataset.map(_tokenize, batched = True, **map_kwargs) + + # If VLM, switch data collator since .pad is needed! + if is_vlm and not hasattr(processing_class, "pad"): + data_collator = DataCollatorForLanguageModeling(tokenizer, mlm = False) + self.data_collator = data_collator + pass + pass + if packing: + print("Unsloth: Hugging Face's packing is currently buggy - we're disabling it for now!") + return dataset + + if max_seq_length == 0: + raise ValueError("When packing is enabled, `max_seq_length` can't be `None`.") + + if use_desc: map_kwargs["desc"] = f"Unsloth: Packing {dataset_name} dataset" + dataset = dataset.select_columns(used_column_names).map( + pack_examples, + batched = True, + fn_kwargs = {"seq_length": max_seq_length,}, + **map_kwargs, + ) + pass + return dataset + + 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 (usually, "input_ids" + # and "attention_mask"). When using `train_on_completion_only` we add a "completion_mask" column to the + # dataset. So we need to override the default signature columns to include "completion_mask" as well. + if self._signature_columns is None: + self._signature_columns = [ + "input_ids", + "labels", + "position_ids", + "completion_mask", + "assistant_masks", + ] + + def compute_loss(self, model, inputs, return_outputs = False, num_items_in_batch = None): + outputs = super().compute_loss( + model, + inputs, + return_outputs = return_outputs, + num_items_in_batch = num_items_in_batch, + ) + return outputs + + # Override training step to add activation offloading context. + def training_step(self, *args, **kwargs): + with self.maybe_activation_offload_context: + return super().training_step(*args, **kwargs) + + def log(self, logs: dict[str, float], start_time: Optional[float] = None) -> None: + mode = "train" if self.model.training else "eval" + metrics = {key: sum(val) / len(val) for key, val in self._metrics[mode].items()} # average the metrics + + # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs` + # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format. + if mode == "eval": + metrics = {f"eval_{key}": val for key, val in metrics.items()} + + logs = {**logs, **metrics} + super().log(logs, start_time) + self._metrics[mode].clear() + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) + + 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 + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + 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=list(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="SFT", + ) + + model_card.save(os.path.join(self.args.output_dir, "README.md")) +class UnslothSFTTrainer(_UnslothSFTTrainer): + """ + + Trainer for Supervised Fine-Tuning (SFT) method. + + This class is a wrapper around the [`transformers.Trainer`] class and inherits all of its attributes and methods. + + Example: + + ```python + from datasets import load_dataset + from trl import SFTTrainer + + dataset = load_dataset("roneneldan/TinyStories", split="train[:1%]") + + trainer = SFTTrainer(model="Qwen/Qwen2-0.5B-Instruct", train_dataset=dataset) + trainer.train() + ``` + + Args: + model (`Union[str, PreTrainedModel]`): + Model to be trained. Can be either: + + - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or a + path to a *directory* containing model weights saved using + [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded + using [`~transformers.AutoModelForCausalLM.from_pretrained`] with the keyword arguments in + `args.model_init_kwargs`. + - A [`~transformers.PreTrainedModel`] object. Only causal language models are supported. + args ([`SFTConfig`], *optional*, defaults to `None`): + Configuration for this trainer. If `None`, a default configuration is used. + data_collator (`DataCollator`, *optional*): + Function to use to form a batch from a list of elements of the processed `train_dataset` or `eval_dataset`. + Will default to a custom [`DataCollatorForLanguageModeling`]. + train_dataset ([`~datasets.Dataset`] or [`~datasets.IterableDataset`]): + Dataset to use for training. SFT supports both [language modeling](#language-modeling) type and + [prompt-completion](#prompt-completion) type. The format of the samples can be either: + + - [Standard](dataset_formats#standard): Each sample contains plain text. + - [Conversational](dataset_formats#conversational): Each sample contains structured messages (e.g., role + and content). + + The trainer also supports processed datasets (tokenized) as long as they contain an `input_ids` field. + eval_dataset ([`~datasets.Dataset`], [`~datasets.IterableDataset`] or `dict[str, Union[Dataset, IterableDataset]]`): + Dataset to use for evaluation. It must meet the same requirements as `train_dataset`. + processing_class ([`~transformers.PreTrainedTokenizerBase`], *optional*, defaults to `None`): + Processing class used to process the data. If `None`, the processing class is loaded from the model's name + with [`~transformers.AutoTokenizer.from_pretrained`]. + callbacks (list of [`~transformers.TrainerCallback`], *optional*, defaults to `None`): + List of callbacks to customize the training loop. Will add those to the list of default callbacks detailed + in [here](https://huggingface.co/docs/transformers/main_classes/callback). + + If you want to remove one of the default callbacks used, use the [`~transformers.Trainer.remove_callback`] + method. + optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`, *optional*, defaults to `(None, None)`): + A tuple containing the optimizer and the scheduler to use. Will default to an instance of [`AdamW`] on your + model and a scheduler given by [`get_linear_schedule_with_warmup`] controlled by `args`. + optimizer_cls_and_kwargs (`Tuple[Type[torch.optim.Optimizer], Dict[str, Any]]`, *optional*, defaults to `None`): + A tuple containing the optimizer class and keyword arguments to use. Overrides `optim` and `optim_args` in + `args`. Incompatible with the `optimizers` argument. + + Unlike `optimizers`, this argument avoids the need to place model parameters on the correct devices before + initializing the Trainer. + preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`, *optional*, defaults to `None`): + A function that preprocess the logits right before caching them at each evaluation step. Must take two + tensors, the logits and the labels, and return the logits once processed as desired. The modifications made + by this function will be reflected in the predictions received by `compute_metrics`. + + Note that the labels (second parameter) will be `None` if the dataset does not have them. + peft_config ([`~peft.PeftConfig`], *optional*, defaults to `None`): + PEFT configuration used to wrap the model. If `None`, the model is not wrapped. + formatting_func (`Optional[Callable]`): + Formatting function applied to the dataset before tokenization. Applying the formatting function explicitly + converts the dataset into a [language modeling](#language-modeling) type. + + """ + def __init__( + self, + model, + args = None, + data_collator = None, + train_dataset = None, + eval_dataset = None, + processing_class = None, + compute_loss_func = None, + compute_metrics = None, + callbacks = None, + optimizer_cls_and_kwargs = None, + preprocess_logits_for_metrics = None, + peft_config = None, + formatting_func = None, + **kwargs + ): + if args is None: args = UnslothSFTConfig() + 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 + 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) + 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 + if 'max_length' not in locals() and not hasattr(args, 'max_length'): + pass + else: + if hasattr(args, 'max_seq_length') and args.max_seq_length is not None and args.max_seq_length > 0: + if hasattr(args, 'max_length'): + args.max_length = args.max_seq_length + max_length = args.max_length + else: + model_max_length = getattr(model, 'max_seq_length', None) + # print(model_max_length, 'mml1') + if model_max_length is None: model_max_length = getattr(model, 'max_length', None) + # print(model_max_length, 'mml2') + if model_max_length is not None: + args.max_length = model_max_length + max_length = args.max_length + elif hasattr(args, 'max_length') and args.max_length is not None: + max_length = args.max_length + # if we are here, then we are in a weird case where max_length is set but max_seq_length is not set + setattr(model, 'max_seq_length', max_length) + else: + print('Unsloth: We did not find `max_seq_length` or `max_length` in the model or args. We will set it to 1024.') + args.max_length = 1024 + 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('sft_trainer', other_metrics) + IGNORED_TOKENIZER_NAMES = os.environ.get('UNSLOTH_IGNORED_TOKENIZER_NAMES', '').split('\n') + from unsloth_zoo.tokenizer_utils import fix_untrained_tokens + from unsloth_zoo.training_utils import fix_zero_training_loss + if 'tokenizer' not in locals(): tokenizer = processing_class + fix_untrained_tokens(model, tokenizer, train_dataset, IGNORED_TOKENIZER_NAMES, eps = 1e-16) + fix_zero_training_loss(model, tokenizer, train_dataset) + + super().__init__( + model = model, + args = args, + data_collator = data_collator, + train_dataset = train_dataset, + eval_dataset = eval_dataset, + processing_class = processing_class, + compute_loss_func = compute_loss_func, + compute_metrics = compute_metrics, + callbacks = callbacks, + optimizer_cls_and_kwargs = optimizer_cls_and_kwargs, + preprocess_logits_for_metrics = preprocess_logits_for_metrics, + peft_config = peft_config, + formatting_func = formatting_func,**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 diff --git a/unsloth_compiled_cache/UnslothXPOTrainer.py b/unsloth_compiled_cache/UnslothXPOTrainer.py new file mode 100644 index 0000000..a7f37f2 --- /dev/null +++ b/unsloth_compiled_cache/UnslothXPOTrainer.py @@ -0,0 +1,1036 @@ +""" +2025.7.1 +2025.7.1 +4.53.1 +0.19.1 +__UNSLOTH_VERSIONING__ +""" +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from trl.trainer.xpo_trainer import (Any, BaseImageProcessor, BasePairwiseJudge, Callable, Dataset, EvalPrediction, F, FeatureExtractionMixin, IterableDataset, OnlineDPOTrainer, OptimizerNames, Optional, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, SIMPLE_CHAT_TEMPLATE, TrainerCallback, Union, XPOConfig, XPOTrainer, empty_cache, generate_model_card, get_comet_experiment_url, get_reward, is_conversational, is_peft_available, is_wandb_available, jinja2, maybe_apply_chat_template, nn, os, 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 +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 UnslothXPOConfig(XPOConfig): + """ + + Configuration class for the [`XPOTrainer`]. + + Subclass of [`OnlineDPOConfig`] we can use all its arguments and add the following: + + Parameters: + alpha (`float` or `list[float]`, *optional*, defaults to `1e-5`): + Weight of the XPO loss term. If a list of floats is provided then the alpha is selected for each new epoch + and the last alpha 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.'}, + ) + 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, + hub_revision = None, + 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, + liger_kernel_config = None, + eval_use_gather_object = False, + average_tokens_across_devices = False, + reward_model_path = None, + judge = None, + max_new_tokens = 64, + max_length = 512, + temperature = 0.9, + missing_eos_penalty = None, + loss_type = 'sigmoid', + dataset_num_proc = None, + disable_dropout = True, + use_vllm = False, + gpu_memory_utilization = 0.55, + ds3_gather_for_generation = True, + 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() + if temperature <= 0: + raise MathError('Unsloth: Please set a positive non-zero temperature since your results will be wrong.') + elif temperature >= 10: + raise MathError('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, + 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, + 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, + missing_eos_penalty = missing_eos_penalty, + loss_type = loss_type, + dataset_num_proc = dataset_num_proc, + disable_dropout = disable_dropout, + use_vllm = use_vllm, + gpu_memory_utilization = gpu_memory_utilization, + ds3_gather_for_generation = ds3_gather_for_generation,**kwargs) + self.vllm_sampling_params = vllm_sampling_params + self.unsloth_num_chunks = unsloth_num_chunks +pass + +class _UnslothXPOTrainer(OnlineDPOTrainer): + r"""""" + + _tag_names = ["trl", "xpo"] + + def __init__( + self, + model: Union[PreTrainedModel, nn.Module] = None, + ref_model: Union[PreTrainedModel, nn.Module] = None, + reward_model: Optional[nn.Module] = None, + judge: Optional[BasePairwiseJudge] = None, + args: Optional[XPOConfig] = 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, + ) -> None: + super().__init__( + model=model, + ref_model=ref_model, + judge=judge, + reward_model=reward_model, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + reward_processing_class=processing_class, # for now, XPOTrainer can't use any reward model + peft_config=peft_config, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + ) + + self._alpha = self.args.alpha + + # Overwrite the stats dictionary to include XPO specific statistics + self.stats = { + # Remove "non_score_reward", "rlhf_reward", "scores" + # Add "loss/dpo", "loss/xpo" + "loss/dpo": [], + "loss/xpo": [], + "objective/kl": [], + "objective/entropy": [], + "rewards/chosen": [], + "rewards/rejected": [], + "rewards/accuracies": [], + "rewards/margins": [], + "logps/chosen": [], + "logps/rejected": [], + # Replace "contain_eos_token" by "model_contain_eos_token" and "ref_contain_eos_token" + "val/model_contain_eos_token": [], + "val/ref_contain_eos_token": [], + "alpha": [], + "beta": [], + } + if self.reward_model is not None: + # Replace "scores" by "model_scores" and "ref_scores" + self.stats["objective/model_scores"] = [] + self.stats["objective/ref_scores"] = [] + self.stats["objective/scores_margin"] = [] + + @property + def alpha(self): + if isinstance(self._alpha, list): + epoch = self.state.epoch + return self._alpha[epoch] if epoch < len(self._alpha) else self._alpha[-1] + else: + return self._alpha + + def _generate_completions(self, prompts, model): + with unwrap_model_for_generation(model, self.accelerator) as unwrapped_policy_model_for_gen: + model_output = unwrapped_policy_model_for_gen.generate( + input_ids=prompts["input_ids"], + attention_mask=prompts["attention_mask"], + generation_config=self.generation_config, + ) + + actual_model_for_ref_generation: torch.nn.Module + if self.ref_model is None: + unwrapped_main_model_for_ref_logic = self.accelerator.unwrap_model(model) + + if is_peft_available() and isinstance(unwrapped_main_model_for_ref_logic, PeftModel): + actual_model_for_ref_generation = unwrapped_main_model_for_ref_logic.get_base_model() + else: + actual_model_for_ref_generation = unwrapped_main_model_for_ref_logic + else: + actual_model_for_ref_generation = self.accelerator.unwrap_model(self.ref_model) + + with unwrap_model_for_generation(actual_model_for_ref_generation, self.accelerator) as final_ref_model_for_gen: + ref_output = final_ref_model_for_gen.generate( + input_ids=prompts["input_ids"], + attention_mask=prompts["attention_mask"], + generation_config=self.generation_config, + ) + + return model_output, ref_output + + def _process_completions(self, model_output, ref_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 + ref_completion_ids = ref_output[:, context_length:] + ref_completion_ids, ref_completion_mask = truncate_right( + ref_completion_ids, self.processing_class.eos_token_id, self.processing_class.pad_token_id + ) + ref_data = { + "input_ids": torch.cat((prompts["input_ids"], ref_completion_ids), dim=1), + "attention_mask": torch.cat((prompts["attention_mask"], ref_completion_mask), dim=1), + "raw": prompts["raw"], + } + + return model_data, ref_data + + def _compute_rewards(self, model_data, ref_data, context_length): + with torch.no_grad(): + _, model_scores, _ = get_reward( + self.reward_model, model_data["input_ids"], self.processing_class.pad_token_id, context_length + ) + _, ref_scores, _ = get_reward( + self.reward_model, ref_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) + ref_contain_eos = torch.any(ref_data["input_ids"] == self.processing_class.eos_token_id, dim=-1) + model_scores[~model_contain_eos] -= self.args.missing_eos_penalty + ref_scores[~ref_contain_eos] -= self.args.missing_eos_penalty + + return model_scores, ref_scores + + def _compute_judge(self, model_data, ref_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] + + ref_data_completions = self.processing_class.batch_decode( + ref_data["input_ids"][:, context_length:], skip_special_tokens=True + ) + ref_data_completions = [completion.strip() for completion in ref_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] + + ref_data_completions = [ + [{"role": "assistant", "content": completion}] for completion in ref_data_completions + ] + ref_data_completions = [template.render(messages=completion) for completion in ref_data_completions] + + ranks_of_first_completion = self.judge.judge( + prompts, + list(zip(model_data_completions, ref_data_completions)), + ) + # convert ranks to a True/False mask: + # when rank == 0, it means the first completion is the best + # when rank == 1, it means the second completion is the best + return torch.tensor([rank == 0 for rank in ranks_of_first_completion], device=model_data["input_ids"].device) + + def _compute_logprobs(self, model, model_data, ref_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 + model_logprobs_model_data = compute_logprobs_for_data(model, model_data) + # Compute logprobs for model on reference completions (for XPO loss) + model_logprobs_ref_data = compute_logprobs_for_data(model, ref_data) + + # Compute logprobs for reference model completions + 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) + ref_logprobs_ref_data = compute_logprobs_for_data(model, ref_data) + else: + ref_logprobs_model_data = compute_logprobs_for_data(self.ref_model, model_data) + ref_logprobs_ref_data = compute_logprobs_for_data(self.ref_model, ref_data) + + # Mask padding tokens + model_padding_mask = model_data["attention_mask"][:, context_length:] == 0 + ref_padding_mask = ref_data["attention_mask"][:, context_length:] == 0 + model_logprobs_model_data = model_logprobs_model_data.masked_fill(model_padding_mask, 0.0) + model_logprobs_ref_data = model_logprobs_ref_data.masked_fill(ref_padding_mask, 0.0) + ref_logprobs_ref_data = ref_logprobs_ref_data.masked_fill(ref_padding_mask, 0.0) + ref_logprobs_model_data = ref_logprobs_model_data.masked_fill(model_padding_mask, 0.0) + + return model_logprobs_model_data, model_logprobs_ref_data, ref_logprobs_ref_data, ref_logprobs_model_data + + def _compute_losses( + self, + model_logprobs_model_data, + model_logprobs_ref_data, + ref_logprobs_ref_data, + ref_logprobs_model_data, + chosen_mask, + ): + # Compute log probs + model_logprobs_model_data_sum = model_logprobs_model_data.sum(1) + model_logprobs_ref_data_sum = model_logprobs_ref_data.sum(1) + ref_logprobs_ref_data_sum = ref_logprobs_ref_data.sum(1) + ref_logprobs_model_data_sum = ref_logprobs_model_data.sum(1) + + chosen_model_logprobs = torch.where(chosen_mask, model_logprobs_model_data_sum, model_logprobs_ref_data_sum) + chosen_ref_logprobs = torch.where(chosen_mask, ref_logprobs_model_data_sum, ref_logprobs_ref_data_sum) + chosen_log_ratios = chosen_model_logprobs - chosen_ref_logprobs + + rejected_model_logprobs = torch.where(~chosen_mask, model_logprobs_model_data_sum, model_logprobs_ref_data_sum) + rejected_ref_logprobs = torch.where(~chosen_mask, ref_logprobs_model_data_sum, ref_logprobs_ref_data_sum) + rejected_log_ratios = rejected_model_logprobs - rejected_ref_logprobs + + # Compute logits as the difference between chosen and rejected log ratios + logits = chosen_log_ratios - rejected_log_ratios + + if self.args.loss_type == "sigmoid": + dpo_losses = -F.logsigmoid(self.beta * logits) + elif self.args.loss_type == "ipo": + dpo_losses = (logits - 1 / (2 * self.beta)) ** 2 + else: + raise NotImplementedError(f"invalid loss type {self.args.loss_type}") + + # Compute XPO specific loss + xpo_losses = self.alpha * model_logprobs_ref_data_sum + + # Total loss + loss = (dpo_losses + xpo_losses).mean() + + return loss, dpo_losses, xpo_losses + + def _log_statistics( + self, + model_data, + ref_data, + model_logprobs_model_data, + model_logprobs_ref_data, + ref_logprobs_ref_data, + ref_logprobs_model_data, + chosen_mask, + dpo_losses, + xpo_losses, + context_length, + model_scores=None, + ref_scores=None, + ): + # Helper function to gather and compute mean + def gather_mean(tensor): + return self.accelerator.gather_for_metrics(tensor).mean().item() + + # Log losses + self.stats["loss/dpo"].append(gather_mean(dpo_losses)) + self.stats["loss/xpo"].append(gather_mean(xpo_losses)) + + # Log scores + if self.reward_model is not None: + self.stats["objective/model_scores"].append(gather_mean(model_scores)) + self.stats["objective/ref_scores"].append(gather_mean(ref_scores)) + self.stats["objective/scores_margin"].append(gather_mean(model_scores - ref_scores)) + + # Log logprobs + model_logprobs_model_data_sum = model_logprobs_model_data.sum(1) + model_logprobs_ref_data_sum = model_logprobs_ref_data.sum(1) + ref_logprobs_ref_data_sum = ref_logprobs_ref_data.sum(1) + ref_logprobs_model_data_sum = ref_logprobs_model_data.sum(1) + + chosen_model_logprobs = torch.where(chosen_mask, model_logprobs_model_data_sum, model_logprobs_ref_data_sum) + chosen_ref_logprobs = torch.where(chosen_mask, ref_logprobs_model_data_sum, ref_logprobs_ref_data_sum) + chosen_log_ratios = chosen_model_logprobs - chosen_ref_logprobs + + rejected_model_logprobs = torch.where(~chosen_mask, model_logprobs_model_data_sum, model_logprobs_ref_data_sum) + rejected_ref_logprobs = torch.where(~chosen_mask, ref_logprobs_model_data_sum, ref_logprobs_ref_data_sum) + rejected_log_ratios = rejected_model_logprobs - rejected_ref_logprobs + + self.stats["logps/chosen"].append(gather_mean(chosen_model_logprobs.mean() + chosen_ref_logprobs.mean())) + self.stats["logps/rejected"].append(gather_mean(rejected_model_logprobs.mean() + rejected_ref_logprobs.mean())) + + # Log rewards + # Compute various statistics + chosen_rewards = chosen_log_ratios * self.beta + rejected_rewards = rejected_log_ratios * self.beta + self.stats["rewards/chosen"].append(gather_mean(chosen_rewards.mean())) + self.stats["rewards/rejected"].append(gather_mean(rejected_rewards.mean())) + + # Calculate KL divergence for model and ref data + kl_model_data = model_logprobs_model_data - ref_logprobs_model_data + kl_ref_data = model_logprobs_ref_data - ref_logprobs_ref_data + mean_kl = (kl_model_data.sum(1) + kl_ref_data.sum(1)).mean() / 2 + self.stats["objective/kl"].append(gather_mean(mean_kl)) + + # Calculate entropy for model and ref data + entropy_model_data = -model_logprobs_model_data.sum(1) + entropy_ref_data = -model_logprobs_ref_data.sum(1) + mean_entropy = (entropy_model_data.mean() + entropy_ref_data.mean()) / 2 + self.stats["objective/entropy"].append(gather_mean(mean_entropy)) + + # Calculate margins + margin = chosen_rewards - rejected_rewards + self.stats["rewards/margins"].append(gather_mean(margin.mean())) + + # Calculate accuracy + accuracy = (margin > 0).float() + self.stats["rewards/accuracies"].append(gather_mean(accuracy.mean())) + + # Log EOS token statistics + model_eos = (model_data["input_ids"][:, context_length:] == self.processing_class.eos_token_id).any(dim=1) + ref_eos = (ref_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(ref_eos.float())) + + # Log alpha and beta + self.stats["alpha"].append(self.alpha) + self.stats["beta"].append(self.beta) + + 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, ref_output = self._generate_completions(prompts, model) + + # Process model completions + model_data, ref_data = self._process_completions(model_output, ref_output, prompts) + + # Compute rewards + if self.reward_model is not None: + model_scores, ref_scores = self._compute_rewards(model_data, ref_data, context_length) + chosen_mask = model_scores >= ref_scores + else: + model_scores, ref_scores = None, None + chosen_mask = self._compute_judge(model_data, ref_data, context_length) + + # Compute logprobs + model_logprobs_model_data, model_logprobs_ref_data, ref_logprobs_ref_data, ref_logprobs_model_data = ( + self._compute_logprobs(model, model_data, ref_data, context_length) + ) + + # Compute loss + loss, dpo_losses, xpo_losses = self._compute_losses( + model_logprobs_model_data, + model_logprobs_ref_data, + ref_logprobs_ref_data, + ref_logprobs_model_data, + chosen_mask, + ) + + # Log everything + self._log_statistics( + model_data, + ref_data, + model_logprobs_model_data.detach(), + model_logprobs_ref_data.detach(), + ref_logprobs_ref_data, + ref_logprobs_model_data, + chosen_mask, + dpo_losses.detach(), + xpo_losses.detach(), + context_length, + model_scores, + ref_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 + + if self.use_apex: + with amp.scale_loss(loss, self.optimizer) as scaled_loss: + scaled_loss.backward() + else: + self.accelerator.backward(loss, **kwargs) + + return loss.detach() / self.args.gradient_accumulation_steps + + 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 + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + citation = textwrap.dedent("""\ + @article{jung2024binary, + title = {{Exploratory Preference Optimization: Harnessing Implicit Q*-Approximation for Sample-Efficient RLHF}}, + author = {Tengyang Xie and Dylan J. Foster and Akshay Krishnamurthy and Corby Rosset and Ahmed Awadallah and Alexander Rakhlin}, + year = 2024, + eprint = {arXiv:2405.21046} + }""") + + 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="XPO", + trainer_citation=citation, + paper_title="Exploratory Preference Optimization: Harnessing Implicit Q*-Approximation for Sample-Efficient RLHF", + paper_id="2405.21046", + ) + + model_card.save(os.path.join(self.args.output_dir, "README.md")) +class UnslothXPOTrainer(_UnslothXPOTrainer): + """ + + Initialize XPOTrainer as a subclass of [`OnlineDPOConfig`]. + + 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_model (`transformers.PreTrainedModel`): + The reward model to score completions with, preferably an `AutoModelForSequenceClassification`. + judge (`BasePairwiseJudge`): + The judge to use for pairwise comparison of model completions. + args (`XPOConfig`): + The XPO 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 (`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. + 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. + + """ + def __init__( + self, + model = None, + ref_model = None, + reward_model = 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, + **kwargs + ): + if args is None: args = UnslothXPOConfig() + 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 + 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) + 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 + 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('xpo_trainer', other_metrics) + + super().__init__( + model = model, + ref_model = ref_model, + reward_model = reward_model, + 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,**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