Spaces:
Running on CPU Upgrade
Running on CPU Upgrade
| """ | |
| 2026.3.2 | |
| 2026.3.4 | |
| 5.3.0 | |
| 0.24.0 | |
| __UNSLOTH_VERSIONING__ | |
| """ | |
| # Unsloth auto generated code | |
| # Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. | |
| # | |
| # This program is free software: you can redistribute it and/or modify | |
| # it under the terms of the GNU Lesser General Public License as published by | |
| # the Free Software Foundation, either version 3 of the License, or | |
| # (at your option) any later version. | |
| # | |
| # This program is distributed in the hope that it will be useful, | |
| # but WITHOUT ANY WARRANTY; without even the implied warranty of | |
| # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the | |
| # GNU General Public License for more details. | |
| # | |
| # You should have received a copy of the GNU Lesser General Public License | |
| # along with this program. If not, see <https://www.gnu.org/licenses/>. | |
| from torch import Tensor | |
| import torch | |
| import torch.nn as nn | |
| from torch.nn import functional as F | |
| from unsloth_zoo.temporary_patches.common import torch_compile | |
| from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable | |
| from trl.trainer.reward_trainer import (Any, AutoModelForSequenceClassification, AutoTokenizer, BaseTrainer, Callable, DataCollator, DataCollatorForPreference, Dataset, EvalPrediction, IterableDataset, Optional, PartialState, Path, PeftConfig, PreTrainedModel, PreTrainedTokenizerBase, RewardConfig, RewardTrainer, TrainerCallback, Union, clone_chat_template, contextlib, dataclass, defaultdict, disable_dropout_in_model, get_act_offloading_ctx_manager, is_conversational, logger, logging, nn, os, pad, re, remove_none_values, suppress_from_pretrained_warning, torch, transformers, Any, AutoModelForSequenceClassification, AutoTokenizer, Callable, DataCollator, DataCollatorForPreference, Dataset, EvalPrediction, IterableDataset, Optional, PeftConfig, PreTrainedModel, PreTrainedTokenizerBase, RewardConfig, TrainerCallback, Union, clone_chat_template, contextlib, defaultdict, disable_dropout_in_model, get_act_offloading_ctx_manager, logger, os, pad, re, suppress_from_pretrained_warning, torch, transformers, PreTrainedModel, logger, os, re, torch) | |
| import os | |
| import math | |
| import logging | |
| from typing import * | |
| from dataclasses import dataclass, field | |
| from packaging.version import Version | |
| import torch | |
| import numpy as np | |
| from contextlib import nullcontext | |
| from torch.nn import functional as F | |
| import inspect | |
| from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling as TransformersDataCollatorForLanguageModeling | |
| from transformers.training_args import ParallelMode | |
| from unsloth_zoo.device_type import DEVICE_TYPE, device_synchronize | |
| # Wrap trainer with padding to right and enable training mode | |
| # Also patches W&B since multiple runs must use wandb.finish() | |
| import functools | |
| from types import MethodType | |
| try: | |
| from unsloth_zoo.gradient_checkpointing import reset_unsloth_gradient_checkpointing_buffers | |
| except: | |
| def reset_unsloth_gradient_checkpointing_buffers(): pass | |
| def prepare_for_training_mode(f): | |
| def wrapper(self, *args, **kwargs): | |
| # Enable training mode | |
| _was_training = None | |
| # Get gradient checkpointing setting from training arguments | |
| use_gc = getattr(self.args, 'gradient_checkpointing', True) | |
| if hasattr(self, 'model') and hasattr(self.model, "training"): | |
| _was_training = self.model.training | |
| if hasattr(self, 'model') and hasattr(self.model, "for_training"): | |
| self.model.for_training(use_gradient_checkpointing=use_gc) | |
| output = f(self, *args, **kwargs) | |
| # Restore previous mode when possible | |
| if hasattr(self, 'model') and hasattr(self.model, "for_inference"): | |
| if _was_training is False: | |
| self.model.for_inference() | |
| elif _was_training is True and hasattr(self.model, "for_training"): | |
| self.model.for_training(use_gradient_checkpointing=use_gc) | |
| # Reset gradient checkpointing buffers to free memory while staying ready for next run | |
| try: | |
| reset_unsloth_gradient_checkpointing_buffers() | |
| except: | |
| pass | |
| # Patch W&B to enable logging on future runs, otherwise it'll overwrite the first run | |
| try: | |
| import wandb | |
| wandb.finish() | |
| except: | |
| pass | |
| return output | |
| return wrapper | |
| pass | |
| torch_compile_options = { | |
| "epilogue_fusion" : True, | |
| "max_autotune" : False, | |
| "shape_padding" : True, | |
| "trace.enabled" : False, | |
| "triton.cudagraphs" : False, | |
| } | |
| def chunked_hidden_states_selective_log_softmax( | |
| hidden_states: torch.Tensor, | |
| lm_head: torch.Tensor, | |
| index: torch.Tensor, | |
| chunks: int = 4, | |
| logit_scale_multiply: float = 0.0, | |
| logit_scale_divide: float = 0.0, | |
| logit_softcapping: float = 0.0, | |
| temperature: float = 1.0, | |
| ) -> torch.Tensor: | |
| # All Unsloth Zoo code licensed under AGPL3 | |
| flat_hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1]) | |
| flat_index = index.reshape(-1) | |
| chunked_hidden_states = torch.chunk(flat_hidden_states, chunks=chunks, dim=0) | |
| chunked_index = torch.chunk(flat_index, chunks=chunks, dim=0) | |
| all_per_token_logps = [] | |
| for chunk_hidden_states, chunk_index in zip(chunked_hidden_states, chunked_index): | |
| chunk_logits = chunk_hidden_states.to(lm_head.dtype) @ lm_head.t() | |
| if logit_scale_multiply != 0.0: | |
| chunk_logits = chunk_logits * logit_scale_multiply | |
| if logit_scale_divide != 0.0: | |
| chunk_logits = chunk_logits / logit_scale_divide | |
| if logit_softcapping != 0.0: | |
| chunk_logits = chunk_logits * torch.tanh(chunk_logits / logit_softcapping) | |
| chunk_logits = chunk_logits.to(torch.float32) | |
| if temperature != 1.0: | |
| chunk_logits = chunk_logits / temperature | |
| selected_logits = torch.gather(chunk_logits, dim=-1, index=chunk_index.unsqueeze(-1)).squeeze(-1) | |
| logsumexp_values = torch.logsumexp(chunk_logits, dim=-1) | |
| per_token_logps = selected_logits - logsumexp_values | |
| all_per_token_logps.append(per_token_logps) | |
| all_per_token_logps = torch.concat(all_per_token_logps) | |
| all_per_token_logps = all_per_token_logps.reshape((hidden_states.shape[0], hidden_states.shape[1])) | |
| return all_per_token_logps | |
| def chunked_selective_log_softmax(logits, index): | |
| # Split into 4 chunks only | |
| chunked_logits = torch.chunk(logits.reshape(-1, logits.shape[-1]), chunks = 4, dim = 0) | |
| chunked_index = torch.chunk(index.reshape(-1), chunks = 4, dim = 0) | |
| all_per_token_logps = [] | |
| # Below loop does the same as selective_log_softmax(chunk_logits, chunk_index) | |
| for chunk_logits, chunk_index in zip(chunked_logits, chunked_index): | |
| chunk_logits = chunk_logits.to(torch.float32) | |
| selected_logits = torch.gather(chunk_logits, dim = -1, index = chunk_index.unsqueeze(-1)).squeeze(-1) | |
| logsumexp_values = torch.logsumexp(chunk_logits, dim = -1) | |
| per_token_logps = selected_logits - logsumexp_values | |
| all_per_token_logps.append(per_token_logps) | |
| pass | |
| all_per_token_logps = torch.concat(all_per_token_logps) | |
| all_per_token_logps = all_per_token_logps.reshape((logits.shape[0], logits.shape[1])) | |
| return all_per_token_logps | |
| def calculate_pad_tokens_in_prompt( | |
| input_ids: torch.Tensor, | |
| logits_to_keep: int, | |
| pad_token_id: int | |
| ) -> torch.Tensor: | |
| """ | |
| Given prompt tensor, it returns all the left padded tokens in that sequence. so [pad, pad, pad, cat] = 3 tokens | |
| """ | |
| if logits_to_keep >= input_ids.shape[1]: | |
| raise ValueError("logits_to_keep must be smaller than the sequence length.") | |
| prompt_section = input_ids[:, :-logits_to_keep] | |
| padding_mask = (prompt_section == pad_token_id) | |
| pad_token_counts = padding_mask.sum(dim=1) | |
| return pad_token_counts | |
| def create_completion_attention_mask( | |
| completion_input_ids: torch.Tensor, | |
| left_pad_tokens_per_prompt: torch.Tensor, | |
| max_left_pad: int, | |
| pad_token_id: int | |
| ) -> torch.Tensor: | |
| """ | |
| Given that we have a sequence, [p,p,p,c,c,c,pad,pad,pad] | |
| Where p are extra prompt tokens we got from slicing the torch tensor, c is completion tokens | |
| and pad are pad tokens, this function would make a completion mask that would 0 out the pad | |
| and p tokens. so in this example [0,0,0,1,1,1,0,0,0] | |
| """ | |
| batch_size, completion_len = completion_input_ids.shape | |
| device = completion_input_ids.device | |
| num_tokens_to_mask = max_left_pad - left_pad_tokens_per_prompt | |
| indices = torch.arange(completion_len, device=device).unsqueeze(0) | |
| shift_mask = indices >= num_tokens_to_mask.unsqueeze(1) | |
| non_padding_mask = (completion_input_ids != pad_token_id) | |
| final_mask = shift_mask & non_padding_mask | |
| return final_mask | |
| def left_pack_padding(tensor: torch.Tensor, pad_id: int) -> torch.Tensor: | |
| """ | |
| Moves all padding tokens in each sequence of a batch to the right. | |
| """ | |
| mask = (tensor != pad_id) | |
| # Must do stable=True since binary mark is unordered | |
| sorted_indices = torch.argsort(mask, dim=1, descending=True, stable=True) | |
| packed_tensor = torch.gather(tensor, 1, sorted_indices) | |
| return packed_tensor | |
| def align_logprobs_with_mask( | |
| logprob_tensor: torch.Tensor, | |
| attention_mask: torch.Tensor, | |
| pad_value: float = 0.0 | |
| ) -> torch.Tensor: | |
| """ | |
| Aligns a log probability tensor with a given attention mask. | |
| """ | |
| device = logprob_tensor.device | |
| batch_size, logprob_seq_len = logprob_tensor.shape | |
| mask_seq_len = attention_mask.shape[1] | |
| padded_logprobs = torch.full( | |
| attention_mask.shape, | |
| fill_value=pad_value, | |
| dtype=logprob_tensor.dtype, | |
| device=device | |
| ) | |
| left_pad_counts = torch.argmax(attention_mask, dim=1) | |
| cols = torch.arange(logprob_seq_len, device=device) | |
| dest_indices = left_pad_counts.unsqueeze(1) + cols | |
| # Create destination row indices | |
| # Shape: [batch_size, logprob_seq_len] | |
| row_indices = torch.arange(batch_size, device=device).unsqueeze(1).expand_as(dest_indices) | |
| # --- 4. Filter out-of-bounds indices and perform assignment --- | |
| # Create a mask to identify only the indices that are within the bounds | |
| # of the target tensor's sequence length. | |
| valid_mask = dest_indices < mask_seq_len | |
| # Use this mask to select only the valid row indices, column indices, | |
| # and the corresponding values from the logprob tensor. | |
| # This flattens the selected elements into 1D tensors. | |
| valid_rows = row_indices[valid_mask] | |
| valid_cols = dest_indices[valid_mask] | |
| valid_vals = logprob_tensor[valid_mask] | |
| # Place the valid values into their correct positions in the padded tensor | |
| # using a single, efficient advanced indexing operation. | |
| padded_logprobs[valid_rows, valid_cols] = valid_vals | |
| return padded_logprobs | |
| def autotune_batch_and_chunks( | |
| total_input_rows, | |
| seq_len, | |
| hidden_size, | |
| vocab_size, | |
| dtype_bytes=16, | |
| multiplier=None | |
| ): | |
| if multiplier is None: | |
| final_m = max(4, seq_len // 4096) | |
| else: | |
| final_m = multiplier | |
| if torch.cuda.is_available(): | |
| free_bytes, _ = torch.cuda.mem_get_info() | |
| limit_gb = (free_bytes / (1024**3))*.80 | |
| elif hasattr(torch, "xpu") and torch.xpu.is_available(): | |
| # For XPU: estimate free memory from total - reserved | |
| total_mem = torch.xpu.get_device_properties(0).total_memory | |
| reserved_mem = torch.xpu.memory_reserved() | |
| free_bytes = total_mem - reserved_mem | |
| limit_gb = (free_bytes / (1024**3)) * 0.80 | |
| else: | |
| # Fallback: assume 8GB available | |
| limit_gb = 8.0 | |
| bytes_to_gb = 1024**3 | |
| b_vals = torch.arange(total_input_rows, 0, -1, device='cpu', dtype=torch.float32) | |
| hidden_gb = (b_vals * seq_len * hidden_size * dtype_bytes) / bytes_to_gb | |
| base_logits = ((b_vals/total_input_rows) * b_vals * seq_len * vocab_size * dtype_bytes) / bytes_to_gb | |
| logits_gb = base_logits / final_m | |
| total_mem_gb = hidden_gb + logits_gb | |
| valid_mask = total_mem_gb <= limit_gb | |
| valid_indices = torch.nonzero(valid_mask, as_tuple=False) | |
| if valid_indices.shape[0] == 0: | |
| #This means your GPU will OOM | |
| return 4, final_m | |
| best_idx = valid_indices[0].item() | |
| final_b = int(b_vals[best_idx].item()) | |
| return final_b, final_m | |
| def sanitize_logprob(logprob): | |
| """Local port of trl.scripts.vllm_serve.sanitize_logprob. | |
| Filters NaN logprobs from vLLM outputs.""" | |
| value = logprob.logprob | |
| if math.isnan(value): | |
| logging.getLogger(__name__).warning( | |
| f"Generated NaN logprob, token logprob '{logprob}' will be ignored" | |
| ) | |
| return None | |
| return value | |
| 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: | |
| > Parameters that control the model | |
| model_init_kwargs (`dict[str, Any]`, *optional*): | |
| Keyword arguments for [`~transformers.AutoModelForCausalLM.from_pretrained`], used when the `model` | |
| argument of the [`RewardTrainer`] is provided as a string. If you're training a MoE architecture and want | |
| to include the load balancing/auxilliary loss as a part of the final loss, remember to set | |
| `output_router_logits=True` in this dictionary. | |
| chat_template_path (`str`, *optional*): | |
| 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. | |
| disable_dropout (`bool`, *optional*, defaults to `True`): | |
| Whether to disable dropout in the model. | |
| > Parameters that control the data preprocessing | |
| dataset_num_proc (`int`, *optional*): | |
| Number of processes to use for processing the dataset. | |
| eos_token (`str`, *optional*): | |
| Token used to indicate the end of a turn or sequence. If `None`, it defaults to | |
| `processing_class.eos_token`. | |
| pad_token (`str`, *optional*): | |
| 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. Samples are filtered out if either chosen or rejected sequence | |
| exceeds this value. If `None`, no filtering is applied. | |
| pad_to_multiple_of (`int`, *optional*): | |
| If set, the sequences will be padded to a multiple of this value. | |
| > Parameters that control the training | |
| center_rewards_coefficient (`float`, *optional*): | |
| 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`. | |
| 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.'}, | |
| ) | |
| unsloth_logit_chunk_multiplier : Optional[int] = field( | |
| default = None, | |
| metadata = {'help': 'Multiplier for chunked logit computations.'}, | |
| ) | |
| unsloth_grpo_mini_batch : Optional[int] = field( | |
| default = None, | |
| metadata = {'help': 'Mini batch size for GRPO hidden state accumulation. Default is None unless user defines it.'}, | |
| ) | |
| max_seq_length : Optional[int] = field( | |
| default = None, | |
| metadata = {'help': 'Maximum sequence length to truncate to.'}, | |
| ) | |
| def __init__( | |
| self, | |
| output_dir = None, | |
| per_device_train_batch_size = 4, | |
| num_train_epochs = 3.0, | |
| max_steps = -1, | |
| learning_rate = 5e-05, | |
| lr_scheduler_type = 'linear', | |
| lr_scheduler_kwargs = None, | |
| warmup_steps = 0.1, | |
| optim = 'adamw_8bit', | |
| optim_args = None, | |
| weight_decay = 0.01, | |
| adam_beta1 = 0.9, | |
| adam_beta2 = 0.999, | |
| adam_epsilon = 1e-08, | |
| optim_target_modules = None, | |
| gradient_accumulation_steps = 2, | |
| average_tokens_across_devices = True, | |
| max_grad_norm = 1.0, | |
| label_smoothing_factor = 0.0, | |
| bf16 = False, | |
| fp16 = False, | |
| bf16_full_eval = False, | |
| fp16_full_eval = False, | |
| tf32 = None, | |
| gradient_checkpointing = True, | |
| gradient_checkpointing_kwargs = None, | |
| torch_compile = False, | |
| torch_compile_backend = None, | |
| torch_compile_mode = None, | |
| use_liger_kernel = False, | |
| liger_kernel_config = None, | |
| use_cache = False, | |
| neftune_noise_alpha = None, | |
| torch_empty_cache_steps = 250, | |
| auto_find_batch_size = False, | |
| logging_strategy = 'steps', | |
| logging_steps = 1, | |
| logging_first_step = False, | |
| log_on_each_node = True, | |
| logging_nan_inf_filter = False, | |
| include_num_input_tokens_seen = False, | |
| log_level = 'passive', | |
| log_level_replica = 'warning', | |
| disable_tqdm = None, | |
| report_to = 'none', | |
| run_name = None, | |
| project = 'huggingface', | |
| trackio_space_id = 'trackio', | |
| eval_strategy = 'no', | |
| eval_steps = None, | |
| eval_delay = 0, | |
| per_device_eval_batch_size = 4, | |
| prediction_loss_only = False, | |
| eval_on_start = False, | |
| eval_do_concat_batches = True, | |
| eval_use_gather_object = False, | |
| eval_accumulation_steps = 2, | |
| batch_eval_metrics = False, | |
| save_only_model = False, | |
| save_strategy = 'steps', | |
| save_steps = 500, | |
| save_on_each_node = False, | |
| save_total_limit = None, | |
| enable_jit_checkpoint = False, | |
| push_to_hub = False, | |
| hub_token = None, | |
| hub_private_repo = None, | |
| hub_model_id = None, | |
| hub_strategy = 'every_save', | |
| hub_always_push = False, | |
| hub_revision = None, | |
| load_best_model_at_end = False, | |
| metric_for_best_model = None, | |
| greater_is_better = None, | |
| ignore_data_skip = False, | |
| restore_callback_states_from_checkpoint = False, | |
| full_determinism = False, | |
| seed = 3407, | |
| data_seed = 3407, | |
| use_cpu = False, | |
| accelerator_config = None, | |
| parallelism_config = None, | |
| dataloader_drop_last = False, | |
| dataloader_num_workers = 0, | |
| dataloader_pin_memory = True, | |
| dataloader_persistent_workers = False, | |
| dataloader_prefetch_factor = None, | |
| remove_unused_columns = True, | |
| label_names = None, | |
| train_sampling_strategy = 'random', | |
| length_column_name = 'length', | |
| ddp_find_unused_parameters = None, | |
| ddp_bucket_cap_mb = None, | |
| ddp_broadcast_buffers = None, | |
| ddp_backend = None, | |
| ddp_timeout = 1800, | |
| fsdp = None, | |
| fsdp_config = None, | |
| deepspeed = None, | |
| debug = '', | |
| skip_memory_metrics = True, | |
| do_train = False, | |
| do_eval = False, | |
| do_predict = False, | |
| resume_from_checkpoint = None, | |
| warmup_ratio = None, | |
| logging_dir = None, | |
| local_rank = -1, | |
| model_init_kwargs = None, | |
| chat_template_path = None, | |
| disable_dropout = True, | |
| dataset_num_proc = None, | |
| eos_token = None, | |
| pad_token = None, | |
| max_length = 1024, | |
| pad_to_multiple_of = None, | |
| center_rewards_coefficient = None, | |
| activation_offloading = False, | |
| vllm_sampling_params = None, | |
| unsloth_num_chunks = -1, | |
| unsloth_logit_chunk_multiplier = None, | |
| unsloth_grpo_mini_batch = None, | |
| max_seq_length = None, | |
| **kwargs, | |
| ): | |
| if learning_rate < 1e-7: print(f'Unsloth: Your learning rate of `{learning_rate}` is too small and less than 1e-7! Consider increasing it, otherwise gradient updates will be close to 0!') | |
| if learning_rate > 1: print(f'Unsloth: Your learning rate of `{learning_rate}` is way too larger > 1! Consider decreasing it to 1e-1, otherwise gradient updates will explode!') | |
| if num_train_epochs is None: | |
| num_train_epochs = 3.0 # Default to 3 epochs if None, max_steps will override | |
| if output_dir is None and save_strategy == 'steps' and save_steps == 500: | |
| output_dir = 'unsloth_training_checkpoints' | |
| save_strategy = 'no' | |
| import multiprocessing as _mp | |
| if _mp.get_start_method() != 'fork': | |
| dataset_num_proc = None | |
| elif dataset_num_proc is None: | |
| import psutil | |
| dataset_num_proc = min(max((psutil.cpu_count() or 1)+4, 2), 64) | |
| memory_gb_left = psutil.virtual_memory().available / (1024**3) | |
| if memory_gb_left <= 2: dataset_num_proc = 1 | |
| else: dataset_num_proc = min(dataset_num_proc, int(memory_gb_left)) | |
| if os.environ.get('UNSLOTH_ENABLE_FLEX_ATTENTION', '0') == '1': | |
| from unsloth_zoo.flex_attention import HAS_FLEX_ATTENTION | |
| if HAS_FLEX_ATTENTION and pad_to_multiple_of is None: | |
| from unsloth_zoo.flex_attention import FLEX_ATTENTION_BLOCK_SIZE | |
| pad_to_multiple_of = FLEX_ATTENTION_BLOCK_SIZE | |
| super().__init__( | |
| output_dir = output_dir, | |
| per_device_train_batch_size = per_device_train_batch_size, | |
| num_train_epochs = num_train_epochs, | |
| max_steps = max_steps, | |
| learning_rate = learning_rate, | |
| lr_scheduler_type = lr_scheduler_type, | |
| lr_scheduler_kwargs = lr_scheduler_kwargs, | |
| warmup_steps = warmup_steps, | |
| optim = optim, | |
| optim_args = optim_args, | |
| weight_decay = weight_decay, | |
| adam_beta1 = adam_beta1, | |
| adam_beta2 = adam_beta2, | |
| adam_epsilon = adam_epsilon, | |
| optim_target_modules = optim_target_modules, | |
| gradient_accumulation_steps = gradient_accumulation_steps, | |
| average_tokens_across_devices = average_tokens_across_devices, | |
| max_grad_norm = max_grad_norm, | |
| label_smoothing_factor = label_smoothing_factor, | |
| bf16 = bf16, | |
| fp16 = fp16, | |
| bf16_full_eval = bf16_full_eval, | |
| fp16_full_eval = fp16_full_eval, | |
| tf32 = tf32, | |
| gradient_checkpointing = gradient_checkpointing, | |
| gradient_checkpointing_kwargs = gradient_checkpointing_kwargs, | |
| torch_compile = torch_compile, | |
| torch_compile_backend = torch_compile_backend, | |
| torch_compile_mode = torch_compile_mode, | |
| use_liger_kernel = use_liger_kernel, | |
| liger_kernel_config = liger_kernel_config, | |
| use_cache = use_cache, | |
| neftune_noise_alpha = neftune_noise_alpha, | |
| torch_empty_cache_steps = torch_empty_cache_steps, | |
| auto_find_batch_size = auto_find_batch_size, | |
| logging_strategy = logging_strategy, | |
| logging_steps = logging_steps, | |
| logging_first_step = logging_first_step, | |
| log_on_each_node = log_on_each_node, | |
| logging_nan_inf_filter = logging_nan_inf_filter, | |
| include_num_input_tokens_seen = include_num_input_tokens_seen, | |
| log_level = log_level, | |
| log_level_replica = log_level_replica, | |
| disable_tqdm = disable_tqdm, | |
| report_to = report_to, | |
| run_name = run_name, | |
| project = project, | |
| trackio_space_id = trackio_space_id, | |
| eval_strategy = eval_strategy, | |
| eval_steps = eval_steps, | |
| eval_delay = eval_delay, | |
| per_device_eval_batch_size = per_device_eval_batch_size, | |
| prediction_loss_only = prediction_loss_only, | |
| eval_on_start = eval_on_start, | |
| eval_do_concat_batches = eval_do_concat_batches, | |
| eval_use_gather_object = eval_use_gather_object, | |
| eval_accumulation_steps = eval_accumulation_steps, | |
| batch_eval_metrics = batch_eval_metrics, | |
| save_only_model = save_only_model, | |
| save_strategy = save_strategy, | |
| save_steps = save_steps, | |
| save_on_each_node = save_on_each_node, | |
| save_total_limit = save_total_limit, | |
| enable_jit_checkpoint = enable_jit_checkpoint, | |
| push_to_hub = push_to_hub, | |
| hub_token = hub_token, | |
| hub_private_repo = hub_private_repo, | |
| hub_model_id = hub_model_id, | |
| hub_strategy = hub_strategy, | |
| hub_always_push = hub_always_push, | |
| hub_revision = hub_revision, | |
| 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, | |
| restore_callback_states_from_checkpoint = restore_callback_states_from_checkpoint, | |
| full_determinism = full_determinism, | |
| seed = seed, | |
| data_seed = data_seed, | |
| use_cpu = use_cpu, | |
| accelerator_config = accelerator_config, | |
| parallelism_config = parallelism_config, | |
| dataloader_drop_last = dataloader_drop_last, | |
| dataloader_num_workers = dataloader_num_workers, | |
| dataloader_pin_memory = dataloader_pin_memory, | |
| dataloader_persistent_workers = dataloader_persistent_workers, | |
| dataloader_prefetch_factor = dataloader_prefetch_factor, | |
| remove_unused_columns = remove_unused_columns, | |
| label_names = label_names, | |
| train_sampling_strategy = train_sampling_strategy, | |
| length_column_name = length_column_name, | |
| ddp_find_unused_parameters = ddp_find_unused_parameters, | |
| ddp_bucket_cap_mb = ddp_bucket_cap_mb, | |
| ddp_broadcast_buffers = ddp_broadcast_buffers, | |
| ddp_backend = ddp_backend, | |
| ddp_timeout = ddp_timeout, | |
| fsdp = fsdp, | |
| fsdp_config = fsdp_config, | |
| deepspeed = deepspeed, | |
| debug = debug, | |
| skip_memory_metrics = skip_memory_metrics, | |
| do_train = do_train, | |
| do_eval = do_eval, | |
| do_predict = do_predict, | |
| resume_from_checkpoint = resume_from_checkpoint, | |
| warmup_ratio = warmup_ratio, | |
| logging_dir = logging_dir, | |
| local_rank = local_rank, | |
| model_init_kwargs = model_init_kwargs, | |
| chat_template_path = chat_template_path, | |
| disable_dropout = disable_dropout, | |
| dataset_num_proc = dataset_num_proc, | |
| eos_token = eos_token, | |
| pad_token = pad_token, | |
| max_length = max_length, | |
| pad_to_multiple_of = pad_to_multiple_of, | |
| center_rewards_coefficient = center_rewards_coefficient, | |
| activation_offloading = activation_offloading,**kwargs) | |
| self.vllm_sampling_params = vllm_sampling_params | |
| self.unsloth_num_chunks = unsloth_num_chunks | |
| if unsloth_grpo_mini_batch is not None: | |
| if self.generation_batch_size >= unsloth_grpo_mini_batch: | |
| self.unsloth_grpo_mini_batch = unsloth_grpo_mini_batch | |
| else: | |
| raise ValueError( | |
| f"Unsloth GRPO mini batch size needs to be less than or equal to the effective generation batch size, " | |
| f"which is self.per_device_train_batch_size * gradient_accumulation_steps." | |
| ) | |
| self.unsloth_logit_chunk_multiplier = unsloth_logit_chunk_multiplier | |
| self.max_seq_length = max_seq_length | |
| pass | |
| class _UnslothRewardTrainer(BaseTrainer): | |
| """""" | |
| _tag_names = ["trl", "reward-trainer"] | |
| _name = "Reward" | |
| _template_file = "rm_model_card.md" | |
| def __init__( | |
| self, | |
| model: Union[str, PreTrainedModel], | |
| args: Optional[RewardConfig] = None, | |
| data_collator: Optional[DataCollator] = None, | |
| train_dataset: Optional[Union[Dataset, IterableDataset]] = None, | |
| eval_dataset: Optional[Union[Dataset, dict[str, Dataset]]] = None, | |
| processing_class: Optional[PreTrainedTokenizerBase] = 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, | |
| ): | |
| # 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 = RewardConfig(f"{model_name}-Reward") | |
| # Model | |
| model_init_kwargs = args.model_init_kwargs or {} | |
| if isinstance(model, str): | |
| model_id = model | |
| dtype = model_init_kwargs.get("dtype") | |
| if isinstance(dtype, torch.dtype) or dtype == "auto" or dtype is None: | |
| pass # dtype is already a torch.dtype or "auto" or None | |
| elif isinstance(dtype, str) and dtype in ["bfloat16", "float16", "float32"]: | |
| model_init_kwargs["dtype"] = getattr(torch, dtype) | |
| else: | |
| raise ValueError( | |
| "Invalid `dtype` passed to `RewardConfig`. Expected either 'auto' or a string representing " | |
| f"a valid `torch.dtype` (e.g., 'float32'), but got {dtype}." | |
| ) | |
| with suppress_from_pretrained_warning(transformers.modeling_utils.logger): | |
| model = AutoModelForSequenceClassification.from_pretrained(model_id, num_labels=1, **model_init_kwargs) | |
| else: | |
| model_id = model.config._name_or_path | |
| if args.model_init_kwargs is not None: | |
| logger.warning( | |
| "You passed `model_init_kwargs` to the `RewardConfig`, but your model is already instantiated. " | |
| "The `model_init_kwargs` will be ignored." | |
| ) | |
| # Processing class | |
| if processing_class is None: | |
| processing_class = AutoTokenizer.from_pretrained(model_id) | |
| # Handle pad token for processors or tokenizers | |
| 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 | |
| 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() | |
| added_tokens = [] | |
| else: | |
| model, processing_class, added_tokens = clone_chat_template( | |
| model, processing_class, args.chat_template_path | |
| ) | |
| else: | |
| added_tokens = [] | |
| # PEFT configuration and model wrapping | |
| if False: | |
| if added_tokens: | |
| # Ensure that the added tokens are trainable | |
| if peft_config.trainable_token_indices is None: | |
| peft_config.trainable_token_indices = {"embed_tokens": added_tokens} | |
| elif "embed_tokens" not in peft_config.trainable_token_indices: | |
| peft_config.trainable_token_indices["embed_tokens"] = added_tokens | |
| else: | |
| peft_config.trainable_token_indices["embed_tokens"].extend(added_tokens) | |
| # Ensure that the lm_head is trainable | |
| if peft_config.modules_to_save is None or "lm_head" not in peft_config.modules_to_save: | |
| logger.warning( | |
| "Cloning chat template added new tokens to the tokenizer, but 'lm_head' is not in PEFT's " | |
| "`modules_to_save`. As a result, the model may not learn to generate outputs with these new " | |
| "tokens, leading to degraded generation quality. To fix this, add " | |
| "`modules_to_save=['lm_head']` to your PEFT configuration." | |
| ) | |
| if peft_config.modules_to_save is None: | |
| peft_config.modules_to_save = ["lm_head"] | |
| else: | |
| peft_config.modules_to_save.append("lm_head") | |
| if False: | |
| pass | |
| # Disable dropout in the model | |
| if args.disable_dropout: | |
| disable_dropout_in_model(model) | |
| # Pad token [needed for SequenceClassification models] | |
| # 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." | |
| ) | |
| model.config.pad_token_id = pad_token_id | |
| processing_class.pad_token_id = pad_token_id | |
| # Data collator | |
| if data_collator is None: | |
| data_collator = DataCollatorForPreference( | |
| pad_token_id=pad_token_id, | |
| pad_to_multiple_of=args.pad_to_multiple_of, | |
| ) | |
| # Dataset | |
| 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") | |
| # 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_metrics=compute_metrics, | |
| callbacks=callbacks, | |
| optimizers=optimizers, | |
| optimizer_cls_and_kwargs=optimizer_cls_and_kwargs, | |
| preprocess_logits_for_metrics=preprocess_logits_for_metrics, | |
| ) | |
| # During evaluation, Trainer calls compute_loss[] only if can_return_loss is True and label_names is empty. | |
| self.can_return_loss = True | |
| self.label_names = [] | |
| # 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) | |
| self.aux_loss_enabled = getattr(model.config, "output_router_logits", False) | |
| def _prepare_dataset( | |
| self, | |
| dataset: Union[Dataset, IterableDataset], | |
| processing_class: PreTrainedTokenizerBase, | |
| args: RewardConfig, | |
| dataset_name: str, | |
| ) -> Union[Dataset, IterableDataset]: | |
| # Tabular backends like Arrow/Parquet insert `None` for mismatched keys in nested structures. Clean them from | |
| # sampled data. | |
| if isinstance(dataset, Dataset): # IterableDataset does not support `with_transform` | |
| dataset = dataset.with_transform(remove_none_values) | |
| # If the dataset is already preprocessed (tokenized), skip the processing steps. | |
| column_names = list(next(iter(dataset)).keys()) | |
| is_processed = "chosen_input_ids" in column_names and "rejected_input_ids" in column_names | |
| # Build the kwargs for the `map` function | |
| map_kwargs = {} | |
| if isinstance(dataset, Dataset): # IterableDataset does not support num_proc | |
| map_kwargs["num_proc"] = args.dataset_num_proc | |
| with PartialState().main_process_first(): | |
| if not is_processed: | |
| # Add EOS token to the end of the sequences if needed | |
| first_example = next(iter(dataset)) | |
| if not is_conversational(first_example): | |
| if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` | |
| map_kwargs["desc"] = f"Adding EOS to {dataset_name} dataset" | |
| def add_eos(example, eos_token): | |
| if not example["chosen"].endswith(eos_token): | |
| example["chosen"] = example["chosen"] + eos_token | |
| if "rejected" in example and not example["rejected"].endswith(eos_token): | |
| example["rejected"] = example["rejected"] + eos_token | |
| return example | |
| dataset = dataset.map( | |
| add_eos, | |
| fn_kwargs={"eos_token": processing_class.eos_token}, | |
| **map_kwargs, | |
| ) | |
| # Tokenize the dataset | |
| if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` | |
| map_kwargs["desc"] = f"Tokenizing {dataset_name} dataset" | |
| def tokenize_fn(example, processing_class): | |
| if "prompt" in example: # explicit prompt case | |
| example["chosen"] = example["prompt"] + example["chosen"] | |
| example["rejected"] = example["prompt"] + example["rejected"] | |
| if is_conversational(example): | |
| chosen_input_ids = processing_class.apply_chat_template( | |
| example["chosen"], | |
| tools=example.get("tools"), | |
| **example.get("chat_template_kwargs", {}), | |
| ) | |
| rejected_input_ids = processing_class.apply_chat_template( | |
| example["rejected"], | |
| tools=example.get("tools"), | |
| **example.get("chat_template_kwargs", {}), | |
| ) | |
| output = {"chosen_input_ids": chosen_input_ids, "rejected_input_ids": rejected_input_ids} | |
| else: | |
| output = { | |
| "chosen_input_ids": processing_class(text=example["chosen"])["input_ids"], | |
| "rejected_input_ids": processing_class(text=example["rejected"])["input_ids"], | |
| } | |
| return output | |
| dataset = dataset.map(tokenize_fn, fn_kwargs={"processing_class": processing_class}, **map_kwargs) | |
| # Filter samples that are longer than `max_length` | |
| if args.max_length is not None: | |
| if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` | |
| map_kwargs["desc"] = f"Filtering {dataset_name} >{args.max_length} tokens" | |
| dataset = dataset.filter( | |
| lambda example: len(example["chosen_input_ids"]) <= args.max_length | |
| and len(example["rejected_input_ids"]) <= args.max_length, | |
| **map_kwargs, | |
| ) | |
| 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"). | |
| if self._signature_columns is None: | |
| self._signature_columns = ["chosen_input_ids", "rejected_input_ids", "margin"] | |
| def compute_loss( | |
| self, | |
| model: nn.Module, | |
| inputs: dict[str, Union[torch.Tensor, Any]], | |
| return_outputs: bool = False, | |
| num_items_in_batch: Optional[torch.Tensor] = None, | |
| ): | |
| """ | |
| Compute training loss and additionally compute token accuracies | |
| """ | |
| mode = "train" if self.model.training else "eval" | |
| # If not set, defaults from model config and may warn since cache isn't compatible with gradient checkpointing | |
| inputs["use_cache"] = False | |
| outputs = model(**inputs) | |
| # Split the rewards into chosen and rejected | |
| rewards_chosen, rewards_rejected = torch.chunk(outputs.logits.squeeze(-1), chunks=2) | |
| # 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 mode == "train": | |
| num_tokens_in_batch = self.accelerator.gather_for_metrics(inputs["attention_mask"].sum()).sum().item() | |
| self._total_train_tokens += num_tokens_in_batch | |
| self._metrics[mode]["num_tokens"] = [self._total_train_tokens] | |
| # Compute min, mean, max, accuracy and margin | |
| with torch.no_grad(): | |
| all_rewards = self.accelerator.gather(outputs.logits) | |
| self._metrics[mode]["min_reward"].append(all_rewards.min().item()) | |
| self._metrics[mode]["mean_reward"].append(all_rewards.mean().item()) | |
| self._metrics[mode]["max_reward"].append(all_rewards.max().item()) | |
| mean_accuracy = (rewards_chosen > rewards_rejected).float().mean() | |
| mean_accuracy = self.accelerator.gather_for_metrics(mean_accuracy).mean().item() | |
| self._metrics[mode]["accuracy"].append(mean_accuracy) | |
| mean_margin = (rewards_chosen - rewards_rejected).mean() | |
| mean_margin = self.accelerator.gather_for_metrics(mean_margin).mean() | |
| self._metrics[mode]["margin"].append(mean_margin.item()) | |
| return (loss, outputs) if return_outputs else loss | |
| # 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.update(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) | |
| class UnslothRewardTrainer(_UnslothRewardTrainer): | |
| """ | |
| Trainer for Outcome-supervised Reward Models (ORM). | |
| This class is a wrapper around the [`~transformers.Trainer`] class and inherits all of its attributes and methods. | |
| Example: | |
| ```python | |
| from trl import RewardTrainer | |
| from datasets import load_dataset | |
| dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train") | |
| trainer = RewardTrainer(model="Qwen/Qwen2.5-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 `AutoModelForSequenceClassification.from_pretrained` with the keyword arguments in | |
| `args.model_init_kwargs`. | |
| - A sequence classification [`~transformers.PreTrainedModel`] object. | |
| args ([`RewardConfig`], *optional*): | |
| Configuration for this trainer. If `None`, a default configuration is used. | |
| data_collator ([`~transformers.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 [`~trainer.reward_trainer.DataCollatorForPreference`]. | |
| train_dataset ([`~datasets.Dataset`] or [`~datasets.IterableDataset`]): | |
| Dataset to use for training. This trainer supports [preference](#preference) type (both implicit and | |
| explicit prompt). 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 `chosen_input_ids` and | |
| `rejected_input_ids` fields. | |
| 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*): | |
| Tokenizer used to process the data. If `None`, the tokenizer 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. | |
| compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): | |
| The function that will be used to compute metrics at evaluation. Must take a | |
| [`~transformers.EvalPrediction`] and return a dictionary string to metric values. When passing | |
| [`RewardConfig`] 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*): | |
| 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[Optional[torch.optim.Optimizer], Optional[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 [`~transformers.get_linear_schedule_with_warmup`] controlled by `args`. | |
| optimizer_cls_and_kwargs (`tuple[Type[torch.optim.Optimizer], Dict[str, Any]]`, *optional*): | |
| 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*): | |
| 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*): | |
| PEFT configuration used to wrap the model. If `None`, the model is not wrapped. Note that if the loaded | |
| model is a causal LM, it's highly recommended to set `modules_to_save=["score"]` in the PEFT configuration | |
| to ensure that the reward head is properly trained. | |
| """ | |
| def __init__( | |
| self, | |
| model, | |
| 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 = 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 | |
| full_finetuning = os.environ.get('UNSLOTH_ENABLE_FULL_FINETUNING', '0') == '1' | |
| if not full_finetuning and (os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1'): | |
| print('Unsloth: Switching to float32 training since model cannot work with float16') | |
| force_float32 = True | |
| mixed_precision_dtype = os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') | |
| dtype = getattr(model.config, 'dtype', None) or getattr(model.config, 'torch_dtype', None) | |
| if dtype is None: dtype = model.get_input_embeddings().weight.dtype | |
| from unsloth_zoo.utils import _get_dtype | |
| dtype = _get_dtype(dtype) | |
| float16 = dtype == torch.float16 | |
| if not force_float32 and (float16 and use_bf16): raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`') | |
| if not force_float32 and (not float16 and use_fp16): raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`') | |
| if force_float32: | |
| # Forced float32 training | |
| args.fp16 = False | |
| args.bf16 = False | |
| os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' | |
| if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' | |
| # args.mixed_precision is a new argument which needs to be set now | |
| elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32': | |
| # Mixed precision training | |
| args.fp16 = float16 | |
| args.bf16 = not float16 | |
| os.environ['ACCELERATE_MIXED_PRECISION'] = 'fp16' if float16 else 'bf16' | |
| if hasattr(args, 'mixed_precision'): args.mixed_precision = 'fp16' if float16 else 'bf16' | |
| # args.mixed_precision is a new argument which needs to be set now | |
| elif mixed_precision_dtype == 'bfloat16': | |
| # Both False since bfloat16 full finetuning doesn't do any autocasting. | |
| args.fp16 = False | |
| args.bf16 = False | |
| os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' | |
| if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' | |
| # args.mixed_precision is a new argument which needs to be set now | |
| if getattr(args, 'eval_dataset', None) is not None and getattr(args, 'eval_strategy', 'no') == 'no': | |
| args.eval_strategy = 'steps' | |
| if getattr(args, 'eval_steps', None) is None: args.eval_steps = 0.1 | |
| ga_steps = getattr(args, 'gradient_accumulation_steps', None) | |
| if ga_steps is not None and ga_steps > 1: | |
| from transformers import __version__ as transformers_version | |
| if Version(transformers_version) <= Version('4.45.2'): | |
| print('**** Unsloth: Please use our fixed gradient_accumulation_steps by updating transformers, TRL and Unsloth!\n' | |
| '`pip install --upgrade --no-cache-dir --force-reinstall --no-deps unsloth transformers trl unsloth_zoo`') | |
| if getattr(args, 'eval_strategy', 'no') != 'no': | |
| eval_bsz = getattr(args, 'per_device_eval_batch_size', 8) | |
| if eval_bsz == 8 and args.per_device_train_batch_size < eval_bsz: args.per_device_eval_batch_size = args.per_device_train_batch_size | |
| if getattr(args, 'eval_accumulation_steps', None) is None and ga_steps is not None: args.eval_accumulation_steps = ga_steps | |
| fp16_full_eval = getattr(args, 'fp16_full_eval', False) | |
| if type(fp16_full_eval) is not bool: fp16_full_eval = False | |
| bf16_full_eval = getattr(args, 'bf16_full_eval', False) | |
| if type(bf16_full_eval) is not bool: bf16_full_eval = False | |
| if args.fp16 and bf16_full_eval: args.bf16_full_eval = False; args.fp16_full_eval = True | |
| if args.bf16 and fp16_full_eval: args.bf16_full_eval = True; args.fp16_full_eval = False | |
| if force_float32: | |
| args.bf16_full_eval = False | |
| args.fp16_full_eval = False | |
| elif os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') == 'bfloat16': | |
| args.bf16_full_eval = True | |
| args.fp16_full_eval = False | |
| elif not bf16_full_eval and not fp16_full_eval: | |
| args.bf16_full_eval = args.bf16 | |
| args.fp16_full_eval = args.fp16 | |
| _output_logits = False | |
| if locals().get('compute_metrics', None) is not None: _output_logits = True | |
| if locals().get('preprocess_logits_for_metrics', None) is not None: _output_logits = True | |
| if _output_logits: | |
| os.environ['UNSLOTH_RETURN_LOGITS'] = '1' | |
| if model is not None: | |
| _warnings_issued = getattr(model, 'warnings_issued', None) | |
| if _warnings_issued is None: | |
| model.warnings_issued = {} | |
| elif not isinstance(_warnings_issued, dict): | |
| try: | |
| model.warnings_issued = dict(_warnings_issued) | |
| except Exception: | |
| model.warnings_issued = {} | |
| if 'max_seq_length' not in locals() and not hasattr(args, 'max_seq_length'): | |
| pass | |
| else: | |
| model_max_seq_length = getattr(model, 'max_seq_length', None) | |
| args_max_seq_length = getattr(args, 'max_seq_length', None) | |
| if args_max_seq_length is None and model_max_seq_length is not None: | |
| max_seq_length = model.max_seq_length | |
| if hasattr(args, 'max_seq_length'): args.max_seq_length = max_seq_length | |
| elif args_max_seq_length is not None and model_max_seq_length is not None: | |
| if args_max_seq_length > model_max_seq_length: | |
| print('Unsloth: You set `max_seq_length` as ' + str(args_max_seq_length) + ' but ' | |
| 'the maximum the model supports is ' + str(model_max_seq_length) + '. We shall reduce it.') | |
| args.max_seq_length = model_max_seq_length | |
| if model is not None and hasattr(model, 'for_training'): | |
| model.for_training(use_gradient_checkpointing=getattr(args, 'gradient_checkpointing', True)) | |
| if 'tokenizer' in locals() and hasattr(tokenizer, 'padding_side'): tokenizer.padding_side = 'right' | |
| if 'processing_class' in locals(): | |
| if hasattr(processing_class, 'padding_side'): processing_class.padding_side = 'right' | |
| if hasattr(processing_class, 'tokenizer') and hasattr(processing_class.tokenizer, 'padding_side'): processing_class.tokenizer.padding_side = 'right' | |
| __tokenizer = processing_class if 'processing_class' in locals() else tokenizer | |
| from unsloth_zoo.vision_utils import UnslothVisionDataCollator | |
| if not isinstance(data_collator, UnslothVisionDataCollator): | |
| if isinstance(data_collator, DataCollatorForSeq2Seq) and 'labels' not in train_dataset.column_names: | |
| data_collator = TransformersDataCollatorForLanguageModeling( | |
| __tokenizer, | |
| mlm = False, | |
| mlm_probability = 0.0, | |
| pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), | |
| ) | |
| elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling) and 'labels' in train_dataset.column_names: | |
| data_collator = DataCollatorForSeq2Seq( | |
| __tokenizer, | |
| pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), | |
| ) | |
| else: | |
| if hasattr(args, 'remove_unused_columns'): args.remove_unused_columns = False | |
| if hasattr(args, 'dataset_text_field'): args.dataset_text_field = '' | |
| if hasattr(args, 'dataset_kwargs'): args.dataset_kwargs = {'skip_prepare_dataset': True} | |
| if not isinstance(data_collator, UnslothVisionDataCollator): | |
| if not hasattr(__tokenizer, 'pad') and hasattr(__tokenizer, 'tokenizer'): | |
| if isinstance(data_collator, DataCollatorForSeq2Seq): | |
| data_collator = DataCollatorForSeq2Seq( | |
| __tokenizer.tokenizer, | |
| pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), | |
| ) | |
| else: | |
| data_collator = TransformersDataCollatorForLanguageModeling( | |
| __tokenizer.tokenizer, | |
| mlm = False, | |
| mlm_probability = 0.0, | |
| pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), | |
| ) | |
| other_metrics = [] | |
| from unsloth_zoo.logging_utils import PatchRLStatistics | |
| PatchRLStatistics('reward_trainer', other_metrics) | |
| # [TODO] Fix up DataParallel multiplying batch sizes | |
| # [TODO] DDP works, but DP seems to not work? [TODO] | |
| if getattr(args, "parallel_mode", None) == ParallelMode.NOT_DISTRIBUTED and args.n_gpu > 1: | |
| if getattr(args, "_n_gpu", 1) != 1: | |
| args._n_gpu = 1 | |
| if "model" in locals() and hasattr(model, "for_training"): | |
| model.for_training(use_gradient_checkpointing=getattr(args, 'gradient_checkpointing', True)) | |
| super().__init__( | |
| model = model, | |
| 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 "model" in locals() and hasattr(model, "for_inference"): | |
| model.for_inference() | |
| if hasattr(self, 'neftune_hook_handle'): | |
| self.neftune_hook_handle.remove() | |
| if hasattr(self, 'neftune_hook_handle'): del self.neftune_hook_handle | |
| if getattr(args, 'neftune_noise_alpha', None) is not None: | |
| model.get_input_embeddings().neftune_noise_alpha = self.neftune_noise_alpha | |
| pass | |
| if hasattr(self, 'accelerator'): | |
| scaler = self.accelerator.scaler | |
| current_model = model | |
| while hasattr(current_model, 'model'): | |
| current_model.accelerator_scaler = scaler | |
| current_model = current_model.model | |
| current_model.accelerator_scaler = scaler | |
| pass | |
| if hasattr(self, 'train'): | |
| self.train = MethodType(prepare_for_training_mode(self.__class__.train), self) | |
| pass | |
| if hasattr(self, 'llm') and self.llm is not None and hasattr(self.llm, 'get_tokenizer'): | |
| _vllm_tok = self.llm.get_tokenizer() | |
| _pc = getattr(self, 'processing_class', None) or getattr(self, 'tokenizer', None) | |
| if _vllm_tok is not None and _pc is not None and getattr(_pc, 'chat_template', None) is not None and getattr(_vllm_tok, 'chat_template', None) is None: | |
| _vllm_tok.chat_template = _pc.chat_template | |
| pass | |
| pass | |
| if hasattr(logger, "addFilter"): | |
| import logging | |
| class HideLoggingMessage(logging.Filter): | |
| def __init__(self, text): self.text = text | |
| def filter(self, x): return not (self.text in x.getMessage()) | |
| pass | |
| logger.addFilter(HideLoggingMessage("`use_cache=True`")) | |