s123hree/green-code-optimizer-a100
0
1"""22026.4.932026.4.844.57.250.23.06__UNSLOTH_VERSIONING__7"""8 9# Unsloth auto generated code10# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.11#12# This program is free software: you can redistribute it and/or modify13# it under the terms of the GNU Lesser General Public License as published by14# the Free Software Foundation, either version 3 of the License, or15# (at your option) any later version.16#17# This program is distributed in the hope that it will be useful,18# but WITHOUT ANY WARRANTY; without even the implied warranty of19# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the20# GNU General Public License for more details.21#22# You should have received a copy of the GNU Lesser General Public License23# along with this program. If not, see <https://www.gnu.org/licenses/>.24 25from torch import Tensor26import torch27import torch.nn as nn28from torch.nn import functional as F29from unsloth_zoo.temporary_patches.common import torch_compile30from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable31from 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, logger, logging, 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, selective_log_softmax, textwrap, torch, AutoModelForCausalLM, BaseImageProcessor, CPOConfig, CPOTrainer, Callable, DPODataCollatorWithPadding, DataCollator, Dataset, EvalLoopOutput, F, FeatureExtractionMixin, Optional, PartialState, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, Trainer, TrainerCallback, Union, autocast, defaultdict, disable_dropout_in_model, inspect, is_comet_available, is_peft_available, is_wandb_available, logger, maybe_apply_chat_template, maybe_extract_prompt, nn, np, os, peft_module_casting_to_bf16, prepare_model_for_kbit_training, torch, F, Optional, PeftModel, PreTrainedModel, Trainer, is_peft_available, logger, os, torch)32 33 34import os35import math36import logging37from typing import *38from dataclasses import dataclass, field39from packaging.version import Version40import torch41import numpy as np42from contextlib import nullcontext43from torch.nn import functional as F44import inspect45from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling as TransformersDataCollatorForLanguageModeling46from transformers.training_args import ParallelMode47from unsloth_zoo.device_type import DEVICE_TYPE, device_synchronize48 49# Wrap trainer with padding to right and enable training mode50import functools51from types import MethodType52try:53 from unsloth_zoo.gradient_checkpointing import reset_unsloth_gradient_checkpointing_buffers54except:55 def reset_unsloth_gradient_checkpointing_buffers(): pass56def prepare_for_training_mode(f):57 @functools.wraps(f)58 def wrapper(self, *args, **kwargs):59 # Finish the previous W&B run if this is a subsequent train() call.60 # We do this at the START of train() (not the end) so that61 # evaluate() / log() still work after train() completes.62 # HF's WandbCallback.setup() will call wandb.init() for the new run.63 # See: https://github.com/unslothai/unsloth/issues/395464 if getattr(self, '_unsloth_training_completed', False):65 try:66 import wandb67 if wandb.run is not None:68 wandb.finish()69 # Reset HF's WandbCallback so it calls wandb.init() for the new run70 for cb in self.callback_handler.callbacks:71 if type(cb).__name__ == 'WandbCallback':72 cb._initialized = False73 break74 except:75 pass76 # Enable training mode77 _was_training = None78 # Get gradient checkpointing setting from training arguments79 use_gc = getattr(self.args, 'gradient_checkpointing', True)80 if hasattr(self, 'model') and hasattr(self.model, "training"):81 _was_training = self.model.training82 if hasattr(self, 'model') and hasattr(self.model, "for_training"):83 self.model.for_training(use_gradient_checkpointing=use_gc)84 output = f(self, *args, **kwargs)85 # Restore previous mode when possible86 if hasattr(self, 'model') and hasattr(self.model, "for_inference"):87 if _was_training is False:88 self.model.for_inference()89 elif _was_training is True and hasattr(self.model, "for_training"):90 self.model.for_training(use_gradient_checkpointing=use_gc)91 # Reset gradient checkpointing buffers to free memory while staying ready for next run92 try:93 reset_unsloth_gradient_checkpointing_buffers()94 except:95 pass96 # Mark that training completed so the next train() call can97 # finish this W&B run before starting a new one98 self._unsloth_training_completed = True99 return output100 return wrapper101pass102 103torch_compile_options = {104 "epilogue_fusion" : True,105 "max_autotune" : False,106 "shape_padding" : True,107 "trace.enabled" : False,108 "triton.cudagraphs" : False,109}110 111@torch.compile(dynamic = True, fullgraph = True, options = torch_compile_options,)112def chunked_hidden_states_selective_log_softmax(113 hidden_states: torch.Tensor,114 lm_head: torch.Tensor,115 index: torch.Tensor,116 chunks: int = 4,117 logit_scale_multiply: float = 0.0,118 logit_scale_divide: float = 0.0,119 logit_softcapping: float = 0.0,120 temperature: float = 1.0,121) -> torch.Tensor:122 # All Unsloth Zoo code licensed under AGPL3123 flat_hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1])124 flat_index = index.reshape(-1)125 126 chunked_hidden_states = torch.chunk(flat_hidden_states, chunks=chunks, dim=0)127 chunked_index = torch.chunk(flat_index, chunks=chunks, dim=0)128 129 all_per_token_logps = []130 131 for chunk_hidden_states, chunk_index in zip(chunked_hidden_states, chunked_index):132 chunk_logits = chunk_hidden_states.to(lm_head.dtype) @ lm_head.t()133 134 if logit_scale_multiply != 0.0:135 chunk_logits = chunk_logits * logit_scale_multiply136 if logit_scale_divide != 0.0:137 chunk_logits = chunk_logits / logit_scale_divide138 if logit_softcapping != 0.0:139 chunk_logits = logit_softcapping * torch.tanh(chunk_logits / logit_softcapping)140 141 chunk_logits = chunk_logits.to(torch.float32)142 143 if temperature != 1.0:144 chunk_logits = chunk_logits / temperature145 146 selected_logits = torch.gather(chunk_logits, dim=-1, index=chunk_index.unsqueeze(-1)).squeeze(-1)147 logsumexp_values = torch.logsumexp(chunk_logits, dim=-1)148 per_token_logps = selected_logits - logsumexp_values149 all_per_token_logps.append(per_token_logps)150 151 all_per_token_logps = torch.concat(all_per_token_logps)152 153 all_per_token_logps = all_per_token_logps.reshape((hidden_states.shape[0], hidden_states.shape[1]))154 return all_per_token_logps155 156@torch.compile(dynamic = True, fullgraph = True, options = torch_compile_options,)157def chunked_selective_log_softmax(logits, index, temperature: float = 1.0):158 # Split into 4 chunks only159 chunked_logits = torch.chunk(logits.reshape(-1, logits.shape[-1]), chunks = 4, dim = 0)160 chunked_index = torch.chunk(index.reshape(-1), chunks = 4, dim = 0)161 all_per_token_logps = []162 # Below loop does the same as selective_log_softmax(chunk_logits, chunk_index)163 for chunk_logits, chunk_index in zip(chunked_logits, chunked_index):164 chunk_logits = chunk_logits.to(torch.float32)165 if temperature != 1.0:166 chunk_logits = chunk_logits / temperature167 selected_logits = torch.gather(chunk_logits, dim = -1, index = chunk_index.unsqueeze(-1)).squeeze(-1)168 logsumexp_values = torch.logsumexp(chunk_logits, dim = -1)169 per_token_logps = selected_logits - logsumexp_values170 all_per_token_logps.append(per_token_logps)171 pass172 all_per_token_logps = torch.concat(all_per_token_logps)173 all_per_token_logps = all_per_token_logps.reshape((logits.shape[0], logits.shape[1]))174 return all_per_token_logps175 176def calculate_pad_tokens_in_prompt(177 input_ids: torch.Tensor,178 logits_to_keep: int,179 pad_token_id: int180) -> torch.Tensor:181 """182 Given prompt tensor, it returns all the left padded tokens in that sequence. so [pad, pad, pad, cat] = 3 tokens183 """184 if logits_to_keep >= input_ids.shape[1]:185 raise ValueError("logits_to_keep must be smaller than the sequence length.")186 187 prompt_section = input_ids[:, :-logits_to_keep]188 189 padding_mask = (prompt_section == pad_token_id)190 191 pad_token_counts = padding_mask.sum(dim=1)192 193 return pad_token_counts194 195def create_completion_attention_mask(196 completion_input_ids: torch.Tensor,197 left_pad_tokens_per_prompt: torch.Tensor,198 max_left_pad: int,199 pad_token_id: int200) -> torch.Tensor:201 """202 Given that we have a sequence, [p,p,p,c,c,c,pad,pad,pad]203 204 Where p are extra prompt tokens we got from slicing the torch tensor, c is completion tokens205 and pad are pad tokens, this function would make a completion mask that would 0 out the pad206 and p tokens. so in this example [0,0,0,1,1,1,0,0,0]207 """208 batch_size, completion_len = completion_input_ids.shape209 device = completion_input_ids.device210 211 num_tokens_to_mask = max_left_pad - left_pad_tokens_per_prompt212 213 indices = torch.arange(completion_len, device=device).unsqueeze(0)214 shift_mask = indices >= num_tokens_to_mask.unsqueeze(1)215 216 non_padding_mask = (completion_input_ids != pad_token_id)217 218 final_mask = shift_mask & non_padding_mask219 220 return final_mask221 222def left_pack_padding(tensor: torch.Tensor, pad_id: int) -> torch.Tensor:223 """224 Moves all padding tokens in each sequence of a batch to the right.225 """226 mask = (tensor != pad_id)227 # Must do stable=True since binary mark is unordered228 sorted_indices = torch.argsort(mask, dim=1, descending=True, stable=True)229 packed_tensor = torch.gather(tensor, 1, sorted_indices)230 return packed_tensor231 232def align_logprobs_with_mask(233 logprob_tensor: torch.Tensor,234 attention_mask: torch.Tensor,235 pad_value: float = 0.0236) -> torch.Tensor:237 """238 Aligns a log probability tensor with a given attention mask.239 """240 241 device = logprob_tensor.device242 batch_size, logprob_seq_len = logprob_tensor.shape243 mask_seq_len = attention_mask.shape[1]244 245 padded_logprobs = torch.full(246 attention_mask.shape,247 fill_value=pad_value,248 dtype=logprob_tensor.dtype,249 device=device250 )251 252 left_pad_counts = torch.argmax(attention_mask, dim=1)253 254 cols = torch.arange(logprob_seq_len, device=device)255 dest_indices = left_pad_counts.unsqueeze(1) + cols256 257 # Create destination row indices258 # Shape: [batch_size, logprob_seq_len]259 row_indices = torch.arange(batch_size, device=device).unsqueeze(1).expand_as(dest_indices)260 261 # --- 4. Filter out-of-bounds indices and perform assignment ---262 # Create a mask to identify only the indices that are within the bounds263 # of the target tensor's sequence length.264 valid_mask = dest_indices < mask_seq_len265 266 # Use this mask to select only the valid row indices, column indices,267 # and the corresponding values from the logprob tensor.268 # This flattens the selected elements into 1D tensors.269 valid_rows = row_indices[valid_mask]270 valid_cols = dest_indices[valid_mask]271 valid_vals = logprob_tensor[valid_mask]272 273 # Place the valid values into their correct positions in the padded tensor274 # using a single, efficient advanced indexing operation.275 padded_logprobs[valid_rows, valid_cols] = valid_vals276 277 return padded_logprobs278 279def autotune_batch_and_chunks(280 total_input_rows,281 seq_len,282 hidden_size,283 vocab_size,284 dtype_bytes=16,285 multiplier=None286):287 if multiplier is None:288 final_m = max(4, seq_len // 4096)289 else:290 final_m = multiplier291 292 if torch.cuda.is_available():293 free_bytes, _ = torch.cuda.mem_get_info()294 limit_gb = (free_bytes / (1024**3))*.80295 elif hasattr(torch, "xpu") and torch.xpu.is_available():296 # For XPU: estimate free memory from total - reserved297 total_mem = torch.xpu.get_device_properties(0).total_memory298 reserved_mem = torch.xpu.memory_reserved()299 free_bytes = total_mem - reserved_mem300 limit_gb = (free_bytes / (1024**3)) * 0.80301 else:302 # Fallback: assume 8GB available303 limit_gb = 8.0304 305 bytes_to_gb = 1024**3306 307 b_vals = torch.arange(total_input_rows, 0, -1, device='cpu', dtype=torch.float32)308 309 hidden_gb = (b_vals * seq_len * hidden_size * dtype_bytes) / bytes_to_gb310 311 base_logits = ((b_vals/total_input_rows) * b_vals * seq_len * vocab_size * dtype_bytes) / bytes_to_gb312 logits_gb = base_logits / final_m313 314 total_mem_gb = hidden_gb + logits_gb315 316 valid_mask = total_mem_gb <= limit_gb317 valid_indices = torch.nonzero(valid_mask, as_tuple=False)318 319 if valid_indices.shape[0] == 0:320 #This means your GPU will OOM321 return 4, final_m322 323 best_idx = valid_indices[0].item()324 final_b = int(b_vals[best_idx].item())325 326 return final_b, final_m327 328def sanitize_logprob(logprob):329 """Local port of trl.scripts.vllm_serve.sanitize_logprob.330 Filters NaN logprobs from vLLM outputs."""331 value = logprob.logprob332 if math.isnan(value):333 logging.getLogger(__name__).warning(334 f"Generated NaN logprob, token logprob '{logprob}' will be ignored"335 )336 return None337 return value338@dataclass339class UnslothCPOConfig(CPOConfig):340 """341 342 Configuration class for the [`CPOTrainer`].343 344 This class includes only the parameters that are specific to CPO training. For a full list of training arguments,345 please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may346 differ from those in [`~transformers.TrainingArguments`].347 348 Using [`~transformers.HfArgumentParser`] we can turn this class into349 [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the350 command line.351 352 Parameters:353 max_length (`int` or `None`, *optional*, defaults to `1024`):354 Maximum length of the sequences (prompt + completion) in the batch. This argument is required if you want355 to use the default data collator.356 max_prompt_length (`int` or `None`, *optional*, defaults to `512`):357 Maximum length of the prompt. This argument is required if you want to use the default data collator.358 max_completion_length (`int` or `None`, *optional*, defaults to `None`):359 Maximum length of the completion. This argument is required if you want to use the default data collator360 and your model is an encoder-decoder.361 beta (`float`, *optional*, defaults to `0.1`):362 Parameter controlling the deviation from the reference model. Higher β means less deviation from the363 reference model. For the IPO loss (`loss_type="ipo"`), β is the regularization parameter denoted by τ in364 the [paper](https://huggingface.co/papers/2310.12036).365 label_smoothing (`float`, *optional*, defaults to `0.0`):366 Label smoothing factor. This argument is required if you want to use the default data collator.367 loss_type (`str`, *optional*, defaults to `"sigmoid"`):368 Type of loss to use. Possible values are:369 370 - `"sigmoid"`: sigmoid loss from the original [DPO](https://huggingface.co/papers/2305.18290) paper.371 - `"hinge"`: hinge loss on the normalized likelihood from the372 [SLiC](https://huggingface.co/papers/2305.10425) paper.373 - `"ipo"`: IPO loss from the [IPO](https://huggingface.co/papers/2310.12036) paper.374 - `"simpo"`: SimPO loss from the [SimPO](https://huggingface.co/papers/2405.14734) paper.375 - `"alphapo"`: AlphaPO loss from the [AlphaPO](https://huggingface.co/papers/2501.03884) paper. This376 automatically sets `loss_type="simpo"` and `cpo_alpha=0.0`.377 378 disable_dropout (`bool`, *optional*, defaults to `True`):379 Whether to disable dropout in the model.380 cpo_alpha (`float`, *optional*, defaults to `1.0`):381 Weight of the BC regularizer in CPO training.382 simpo_gamma (`float`, *optional*, defaults to `0.5`):383 Target reward margin for the SimPO loss, used only when the `loss_type="simpo"`.384 alpha (`float`, *optional*, defaults to `0.0`):385 Alpha parameter that controls reward function shape across all loss types. When alpha=0 (default), uses386 standard log probability rewards. When `alpha != 0`, applies AlphaPO transformation: `r = (1 - p^(-alpha))387 / alpha` from the [AlphaPO paper](https://huggingface.co/papers/2501.03884). This parameter works with all388 loss types.389 label_pad_token_id (`int`, *optional*, defaults to `-100`):390 Label pad token id. This argument is required if you want to use the default data collator.391 padding_value (`int` or `None`, *optional*, defaults to `None`):392 Padding value to use. If `None`, the padding value of the tokenizer is used.393 truncation_mode (`str`,*optional*, defaults to `"keep_end"`):394 Truncation mode to use when the prompt is too long. Possible values are `"keep_end"` or `"keep_start"`.395 This argument is required if you want to use the default data collator.396 generate_during_eval (`bool`, *optional*, defaults to `False`):397 If `True`, generates and logs completions from the model to W&B or Comet during evaluation.398 is_encoder_decoder (`bool` or `None`, *optional*, defaults to `None`):399 When using the `model_init` argument (callable) to instantiate the model instead of the `model` argument,400 you need to specify if the model returned by the callable is an encoder-decoder model.401 model_init_kwargs (`dict[str, Any]` or `None`, *optional*, defaults to `None`):402 Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model from a403 string.404 dataset_num_proc (`int` or `None`, *optional*, defaults to `None`):405 Number of processes to use for processing the dataset.406 407 """408 vllm_sampling_params: Optional[Any] = field(409 default = None,410 metadata = {'help': 'vLLM SamplingParams'},411 )412 unsloth_num_chunks : Optional[int] = field(413 default = -1,414 metadata = {'help': 'Chunk size to reduce memory usage. -1 is most efficient.'},415 )416 unsloth_logit_chunk_multiplier : Optional[int] = field(417 default = None,418 metadata = {'help': 'Multiplier for chunked logit computations.'},419 )420 unsloth_grpo_mini_batch : Optional[int] = field(421 default = None,422 metadata = {'help': 'Mini batch size for GRPO hidden state accumulation. Default is None unless user defines it.'},423 )424 max_seq_length : Optional[int] = field(425 default = None,426 metadata = {'help': 'Maximum sequence length to truncate to.'},427 )428 def __init__(429 self,430 output_dir = None,431 overwrite_output_dir = None,432 do_train = False,433 do_eval = False,434 do_predict = False,435 eval_strategy = 'no',436 prediction_loss_only = False,437 per_device_train_batch_size = 4,438 per_device_eval_batch_size = 4,439 per_gpu_train_batch_size = None,440 per_gpu_eval_batch_size = None,441 gradient_accumulation_steps = 2,442 eval_accumulation_steps = 2,443 eval_delay = 0,444 torch_empty_cache_steps = 250,445 learning_rate = 5e-05,446 weight_decay = 0.01,447 adam_beta1 = 0.9,448 adam_beta2 = 0.999,449 adam_epsilon = 1e-08,450 max_grad_norm = 1.0,451 num_train_epochs = 3.0,452 max_steps = -1,453 lr_scheduler_type = 'linear',454 warmup_ratio = 0.1,455 warmup_steps = 0,456 log_level = 'passive',457 log_level_replica = 'warning',458 log_on_each_node = True,459 logging_dir = None,460 logging_strategy = 'steps',461 logging_first_step = False,462 logging_steps = 1,463 logging_nan_inf_filter = False,464 save_strategy = 'steps',465 save_steps = 500,466 save_total_limit = None,467 save_safetensors = True,468 save_on_each_node = False,469 save_only_model = False,470 restore_callback_states_from_checkpoint = False,471 no_cuda = False,472 use_cpu = False,473 use_mps_device = False,474 seed = 3407,475 data_seed = 3407,476 jit_mode_eval = False,477 bf16 = False,478 fp16 = False,479 fp16_opt_level = 'O1',480 half_precision_backend = 'auto',481 bf16_full_eval = False,482 fp16_full_eval = False,483 tf32 = None,484 local_rank = -1,485 ddp_backend = None,486 tpu_num_cores = None,487 tpu_metrics_debug = False,488 debug = '',489 dataloader_drop_last = False,490 eval_steps = None,491 dataloader_num_workers = 0,492 dataloader_prefetch_factor = None,493 past_index = -1,494 run_name = None,495 disable_tqdm = None,496 remove_unused_columns = True,497 label_names = None,498 load_best_model_at_end = False,499 metric_for_best_model = None,500 greater_is_better = None,501 ignore_data_skip = False,502 fsdp = None,503 fsdp_min_num_params = 0,504 fsdp_config = None,505 fsdp_transformer_layer_cls_to_wrap = None,506 accelerator_config = None,507 parallelism_config = None,508 deepspeed = None,509 label_smoothing_factor = 0.0,510 optim = 'adamw_8bit',511 optim_args = None,512 adafactor = False,513 group_by_length = False,514 length_column_name = 'length',515 report_to = 'none',516 project = 'huggingface',517 trackio_space_id = 'trackio',518 ddp_find_unused_parameters = None,519 ddp_bucket_cap_mb = None,520 ddp_broadcast_buffers = None,521 dataloader_pin_memory = True,522 dataloader_persistent_workers = False,523 skip_memory_metrics = True,524 use_legacy_prediction_loop = False,525 push_to_hub = False,526 resume_from_checkpoint = None,527 hub_model_id = None,528 hub_strategy = 'every_save',529 hub_token = None,530 hub_private_repo = None,531 hub_always_push = False,532 hub_revision = None,533 gradient_checkpointing = True,534 gradient_checkpointing_kwargs = None,535 include_inputs_for_metrics = False,536 eval_do_concat_batches = True,537 fp16_backend = 'auto',538 push_to_hub_model_id = None,539 push_to_hub_organization = None,540 push_to_hub_token = None,541 mp_parameters = '',542 auto_find_batch_size = False,543 full_determinism = False,544 torchdynamo = None,545 ray_scope = 'last',546 ddp_timeout = 1800,547 torch_compile = False,548 torch_compile_backend = None,549 torch_compile_mode = None,550 include_tokens_per_second = False,551 include_num_input_tokens_seen = False,552 neftune_noise_alpha = None,553 optim_target_modules = None,554 batch_eval_metrics = False,555 eval_on_start = False,556 use_liger_kernel = False,557 liger_kernel_config = None,558 eval_use_gather_object = False,559 average_tokens_across_devices = True,560 max_length = 1024,561 max_prompt_length = 512,562 max_completion_length = None,563 beta = 0.1,564 label_smoothing = 0.0,565 loss_type = 'sigmoid',566 disable_dropout = True,567 cpo_alpha = 1.0,568 simpo_gamma = 0.5,569 alpha = 0.0,570 label_pad_token_id = -100,571 padding_value = None,572 truncation_mode = 'keep_end',573 generate_during_eval = False,574 is_encoder_decoder = None,575 model_init_kwargs = None,576 dataset_num_proc = None,577 vllm_sampling_params = None,578 unsloth_num_chunks = -1,579 unsloth_logit_chunk_multiplier = None,580 unsloth_grpo_mini_batch = None,581 max_seq_length = None,582 **kwargs,583 ):584 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!')585 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!')586 if num_train_epochs is None:587 num_train_epochs = 3.0 # Default to 3 epochs if None, max_steps will override588 if output_dir is None and save_strategy == 'steps' and save_steps == 500:589 output_dir = 'unsloth_training_checkpoints'590 save_strategy = 'no'591 import multiprocessing as _mp592 if dataset_num_proc is None:593 if _mp.get_start_method() != 'fork':594 dataset_num_proc = None595 else:596 import psutil597 dataset_num_proc = min(max((psutil.cpu_count() or 1)+4, 2), 64)598 memory_gb_left = psutil.virtual_memory().available / (1024**3)599 if memory_gb_left <= 2: dataset_num_proc = 1600 else: dataset_num_proc = min(dataset_num_proc, int(memory_gb_left))601 602 super().__init__(603 output_dir = output_dir,604 overwrite_output_dir = overwrite_output_dir,605 do_train = do_train,606 do_eval = do_eval,607 do_predict = do_predict,608 eval_strategy = eval_strategy,609 prediction_loss_only = prediction_loss_only,610 per_device_train_batch_size = per_device_train_batch_size,611 per_device_eval_batch_size = per_device_eval_batch_size,612 per_gpu_train_batch_size = per_gpu_train_batch_size,613 per_gpu_eval_batch_size = per_gpu_eval_batch_size,614 gradient_accumulation_steps = gradient_accumulation_steps,615 eval_accumulation_steps = eval_accumulation_steps,616 eval_delay = eval_delay,617 torch_empty_cache_steps = torch_empty_cache_steps,618 learning_rate = learning_rate,619 weight_decay = weight_decay,620 adam_beta1 = adam_beta1,621 adam_beta2 = adam_beta2,622 adam_epsilon = adam_epsilon,623 max_grad_norm = max_grad_norm,624 num_train_epochs = num_train_epochs,625 max_steps = max_steps,626 lr_scheduler_type = lr_scheduler_type,627 warmup_ratio = warmup_ratio,628 warmup_steps = warmup_steps,629 log_level = log_level,630 log_level_replica = log_level_replica,631 log_on_each_node = log_on_each_node,632 logging_dir = logging_dir,633 logging_strategy = logging_strategy,634 logging_first_step = logging_first_step,635 logging_steps = logging_steps,636 logging_nan_inf_filter = logging_nan_inf_filter,637 save_strategy = save_strategy,638 save_steps = save_steps,639 save_total_limit = save_total_limit,640 save_safetensors = save_safetensors,641 save_on_each_node = save_on_each_node,642 save_only_model = save_only_model,643 restore_callback_states_from_checkpoint = restore_callback_states_from_checkpoint,644 no_cuda = no_cuda,645 use_cpu = use_cpu,646 use_mps_device = use_mps_device,647 seed = seed,648 data_seed = data_seed,649 jit_mode_eval = jit_mode_eval,650 bf16 = bf16,651 fp16 = fp16,652 fp16_opt_level = fp16_opt_level,653 half_precision_backend = half_precision_backend,654 bf16_full_eval = bf16_full_eval,655 fp16_full_eval = fp16_full_eval,656 tf32 = tf32,657 local_rank = local_rank,658 ddp_backend = ddp_backend,659 tpu_num_cores = tpu_num_cores,660 tpu_metrics_debug = tpu_metrics_debug,661 debug = debug,662 dataloader_drop_last = dataloader_drop_last,663 eval_steps = eval_steps,664 dataloader_num_workers = dataloader_num_workers,665 dataloader_prefetch_factor = dataloader_prefetch_factor,666 past_index = past_index,667 run_name = run_name,668 disable_tqdm = disable_tqdm,669 remove_unused_columns = remove_unused_columns,670 label_names = label_names,671 load_best_model_at_end = load_best_model_at_end,672 metric_for_best_model = metric_for_best_model,673 greater_is_better = greater_is_better,674 ignore_data_skip = ignore_data_skip,675 fsdp = fsdp,676 fsdp_min_num_params = fsdp_min_num_params,677 fsdp_config = fsdp_config,678 fsdp_transformer_layer_cls_to_wrap = fsdp_transformer_layer_cls_to_wrap,679 accelerator_config = accelerator_config,680 parallelism_config = parallelism_config,681 deepspeed = deepspeed,682 label_smoothing_factor = label_smoothing_factor,683 optim = optim,684 optim_args = optim_args,685 adafactor = adafactor,686 group_by_length = group_by_length,687 length_column_name = length_column_name,688 report_to = report_to,689 project = project,690 trackio_space_id = trackio_space_id,691 ddp_find_unused_parameters = ddp_find_unused_parameters,692 ddp_bucket_cap_mb = ddp_bucket_cap_mb,693 ddp_broadcast_buffers = ddp_broadcast_buffers,694 dataloader_pin_memory = dataloader_pin_memory,695 dataloader_persistent_workers = dataloader_persistent_workers,696 skip_memory_metrics = skip_memory_metrics,697 use_legacy_prediction_loop = use_legacy_prediction_loop,698 push_to_hub = push_to_hub,699 resume_from_checkpoint = resume_from_checkpoint,700 hub_model_id = hub_model_id,701 hub_strategy = hub_strategy,702 hub_token = hub_token,703 hub_private_repo = hub_private_repo,704 hub_always_push = hub_always_push,705 hub_revision = hub_revision,706 gradient_checkpointing = gradient_checkpointing,707 gradient_checkpointing_kwargs = gradient_checkpointing_kwargs,708 include_inputs_for_metrics = include_inputs_for_metrics,709 eval_do_concat_batches = eval_do_concat_batches,710 fp16_backend = fp16_backend,711 push_to_hub_model_id = push_to_hub_model_id,712 push_to_hub_organization = push_to_hub_organization,713 push_to_hub_token = push_to_hub_token,714 mp_parameters = mp_parameters,715 auto_find_batch_size = auto_find_batch_size,716 full_determinism = full_determinism,717 torchdynamo = torchdynamo,718 ray_scope = ray_scope,719 ddp_timeout = ddp_timeout,720 torch_compile = torch_compile,721 torch_compile_backend = torch_compile_backend,722 torch_compile_mode = torch_compile_mode,723 include_tokens_per_second = include_tokens_per_second,724 include_num_input_tokens_seen = include_num_input_tokens_seen,725 neftune_noise_alpha = neftune_noise_alpha,726 optim_target_modules = optim_target_modules,727 batch_eval_metrics = batch_eval_metrics,728 eval_on_start = eval_on_start,729 use_liger_kernel = use_liger_kernel,730 liger_kernel_config = liger_kernel_config,731 eval_use_gather_object = eval_use_gather_object,732 average_tokens_across_devices = average_tokens_across_devices,733 max_length = max_length,734 max_prompt_length = max_prompt_length,735 max_completion_length = max_completion_length,736 beta = beta,737 label_smoothing = label_smoothing,738 loss_type = loss_type,739 disable_dropout = disable_dropout,740 cpo_alpha = cpo_alpha,741 simpo_gamma = simpo_gamma,742 alpha = alpha,743 label_pad_token_id = label_pad_token_id,744 padding_value = padding_value,745 truncation_mode = truncation_mode,746 generate_during_eval = generate_during_eval,747 is_encoder_decoder = is_encoder_decoder,748 model_init_kwargs = model_init_kwargs,749 dataset_num_proc = dataset_num_proc,**kwargs)750 self.vllm_sampling_params = vllm_sampling_params751 self.unsloth_num_chunks = unsloth_num_chunks752 if unsloth_grpo_mini_batch is not None:753 if self.generation_batch_size >= unsloth_grpo_mini_batch:754 self.unsloth_grpo_mini_batch = unsloth_grpo_mini_batch755 else:756 raise ValueError(757 f"Unsloth GRPO mini batch size needs to be less than or equal to the effective generation batch size, "758 f"which is self.per_device_train_batch_size * gradient_accumulation_steps."759 )760 self.unsloth_logit_chunk_multiplier = unsloth_logit_chunk_multiplier761 self.max_seq_length = max_seq_length762 763pass764 765class _UnslothCPOTrainer(Trainer):766 r""""""767 768 _tag_names = ["trl", "cpo"]769 770 def __init__(771 self,772 model: Optional[Union[PreTrainedModel, nn.Module, str]] = None,773 args: Optional[CPOConfig] = None,774 data_collator: Optional[DataCollator] = None,775 train_dataset: Optional[Dataset] = None,776 eval_dataset: Optional[Union[Dataset, dict[str, Dataset]]] = None,777 processing_class: Optional[778 Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin]779 ] = None,780 model_init: Optional[Callable[[], PreTrainedModel]] = None,781 callbacks: Optional[list[TrainerCallback]] = None,782 optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),783 preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None,784 peft_config: Optional[dict] = None,785 compute_metrics: Optional[Callable[[EvalLoopOutput], dict]] = None,786 ):787 if args.model_init_kwargs is None:788 model_init_kwargs = {}789 elif not isinstance(model, str):790 raise ValueError("You passed model_kwargs to the CPOTrainer. But your model is already instantiated.")791 else:792 model_init_kwargs = args.model_init_kwargs793 dtype = model_init_kwargs.get("dtype")794 if dtype is not None:795 # Convert to `torch.dtype` if an str is passed796 if isinstance(dtype, str) and dtype != "auto":797 dtype = getattr(torch, dtype)798 if dtype != "auto" and not isinstance(dtype, torch.dtype):799 raise ValueError(800 f"Invalid `dtype` passed to the CPOConfig. Expected a string with either `torch.dtype` or 'auto', but got {dtype}."801 )802 model_init_kwargs["dtype"] = dtype803 804 if isinstance(model, str):805 model = AutoModelForCausalLM.from_pretrained(model, **model_init_kwargs)806 807 # Initialize this variable to False. This helps tracking the case when `peft_module_casting_to_bf16`808 # has been called in order to properly call autocast if needed.809 self._peft_has_been_casted_to_bf16 = False810 811 if not is_peft_available() and peft_config is not None:812 raise ValueError(813 "PEFT is not installed and you passed a `peft_config` in the trainer's kwargs, please install it to use the PEFT models"814 )815 elif is_peft_available() and peft_config is not None:816 # if model is a peft model and we have a peft_config, we merge and unload it first817 if isinstance(model, PeftModel):818 model = model.merge_and_unload()819 820 if getattr(model, "is_loaded_in_8bit", False) or getattr(model, "is_loaded_in_4bit", False):821 _support_gc_kwargs = hasattr(822 args, "gradient_checkpointing_kwargs"823 ) and "gradient_checkpointing_kwargs" in list(824 inspect.signature(prepare_model_for_kbit_training).parameters825 )826 827 prepare_model_kwargs = {"use_gradient_checkpointing": args.gradient_checkpointing}828 829 if _support_gc_kwargs:830 prepare_model_kwargs["gradient_checkpointing_kwargs"] = args.gradient_checkpointing_kwargs831 832 model = prepare_model_for_kbit_training(model, **prepare_model_kwargs)833 elif args.gradient_checkpointing:834 # For backward compatibility with older versions of transformers835 if hasattr(model, "enable_input_require_grads"):836 model.enable_input_require_grads()837 else:838 839 def make_inputs_require_grad(module, input, output):840 output.requires_grad_(True)841 842 model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)843 844 # get peft model with the given config845 model = model846 if args.bf16 and getattr(model, "is_loaded_in_4bit", False):847 peft_module_casting_to_bf16(model)848 # If args.bf16 we need to explicitly call `generate` with torch amp autocast context manager849 self._peft_has_been_casted_to_bf16 = True850 851 # For models that use gradient_checkpointing, we need to attach a hook that enables input852 # to explicitly have `requires_grad=True`, otherwise training will either silently853 # fail or completely fail.854 elif args.gradient_checkpointing:855 # For backward compatibility with older versions of transformers856 if hasattr(model, "enable_input_require_grads"):857 model.enable_input_require_grads()858 else:859 860 def make_inputs_require_grad(module, input, output):861 output.requires_grad_(True)862 863 model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)864 865 if args.generate_during_eval and not (is_wandb_available() or is_comet_available()):866 raise ValueError(867 "`generate_during_eval=True` requires Weights and Biases or Comet to be installed."868 " Please install `wandb` or `comet-ml` to resolve."869 )870 871 if model is not None:872 self.is_encoder_decoder = model.config.is_encoder_decoder873 elif args.is_encoder_decoder is None:874 raise ValueError("When no model is provided, you need to pass the parameter is_encoder_decoder.")875 else:876 self.is_encoder_decoder = args.is_encoder_decoder877 878 if self.is_encoder_decoder:879 self.decoder_start_token_id = model.config.decoder_start_token_id880 self.pad_token_id = model.config.pad_token_id881 882 if processing_class is None:883 raise ValueError("processing_class must be specified to tokenize a CPO dataset.")884 if args.max_length is None:885 logger.warning(886 "`max_length` is not set in the CPOConfig's init"887 " it will default to `512` by default, but you should do it yourself in the future.",888 )889 max_length = 512890 else:891 max_length = args.max_length892 if args.max_prompt_length is None:893 logger.warning(894 "`max_prompt_length` is not set in the CPOConfig's init"895 " it will default to `128` by default, but you should do it yourself in the future.",896 )897 max_prompt_length = 128898 else:899 max_prompt_length = args.max_prompt_length900 901 if not max_prompt_length < max_length:902 raise ValueError(903 f"max_prompt_length ({max_prompt_length}) should be strictly less than max_length ({max_length})."904 )905 906 if args.max_completion_length is None and self.is_encoder_decoder:907 logger.warning(908 "When using an encoder decoder architecture, you should set `max_completion_length` in the CPOConfig's init"909 " it will default to `128` by default, but you should do it yourself in the future.",910 )911 max_completion_length = 128912 else:913 max_completion_length = args.max_completion_length914 915 if data_collator is None:916 data_collator = DPODataCollatorWithPadding(917 pad_token_id=processing_class.pad_token_id,918 label_pad_token_id=args.label_pad_token_id,919 is_encoder_decoder=self.is_encoder_decoder,920 )921 922 if args.remove_unused_columns:923 args.remove_unused_columns = False924 # warn users925 logger.warning(926 "When using DPODataCollatorWithPadding, you should set `remove_unused_columns=False` in your TrainingArguments"927 " we have set it for you, but you should do it yourself in the future.",928 )929 930 self.use_dpo_data_collator = True931 else:932 self.use_dpo_data_collator = False933 934 # Disable dropout in the model935 if args.disable_dropout:936 disable_dropout_in_model(model)937 938 self.max_length = max_length939 self.generate_during_eval = args.generate_during_eval940 self.label_pad_token_id = args.label_pad_token_id941 self.padding_value = args.padding_value if args.padding_value is not None else processing_class.pad_token_id942 self.max_prompt_length = max_prompt_length943 self.truncation_mode = args.truncation_mode944 self.max_completion_length = max_completion_length945 self.processing_class = processing_class946 947 if args.loss_type in ["hinge", "ipo"] and args.label_smoothing > 0:948 logger.warning(949 f"You are using the {args.loss_type} loss type that does not support label smoothing. The "950 "`label_smoothing` parameter will be ignored. Set `label_smoothing` to `0.0` to remove this warning.",951 )952 if args.loss_type == "kto_pair":953 raise ValueError("Support for kto_pair has been removed in CPOTrainer. Please use KTOTrainer.")954 955 self.beta = args.beta956 self.label_smoothing = args.label_smoothing957 self.loss_type = args.loss_type958 self.cpo_alpha = args.cpo_alpha959 self.aux_loss_enabled = getattr(model.config, "output_router_logits", False)960 self.aux_loss_coef = getattr(model.config, "router_aux_loss_coef", 0.0)961 if self.aux_loss_enabled and self.aux_loss_coef == 0.0:962 logger.warning(963 "You set `output_router_logits` to `True` in the model config, but `router_aux_loss_coef` is set to "964 "`0.0`, meaning the auxiliary loss will not be used. Either set `router_aux_loss_coef` to a value "965 "greater than `0.0`, or set `output_router_logits` to `False` if you don't want to use the auxiliary "966 "loss.",967 )968 969 if args.loss_type == "simpo":970 self.simpo_gamma = args.simpo_gamma971 972 # AlphaPO parameter for reward shaping973 self.alpha = args.alpha974 975 self._stored_metrics = defaultdict(lambda: defaultdict(list))976 977 # The trainer estimates the number of FLOPs [floating-point operations] using the number of elements in the978 # input tensor associated with the key "input_ids". However, in CPO, the sampled data does not include the979 # "input_ids" key. Instead, the available keys are "prompt_input_ids", "chosen_input_ids", and980 # "rejected_input_ids". As a result, the trainer issues the warning: "Could not estimate the number of tokens981 # of the input, floating-point operations will not be computed." To suppress this warning, we set the982 # "estimate_tokens" key in the model's "warnings_issued" dictionary to True. This acts as a flag to indicate983 # that the warning has already been issued.984 model.warnings_issued["estimate_tokens"] = True985 986 # Compute that only on the main process for faster data processing.987 # see: https://github.com/huggingface/trl/pull/1255988 with PartialState().main_process_first():989 # Extract the prompt if needed, and apply the chat template if needed990 train_dataset = train_dataset.map(maybe_extract_prompt, num_proc=args.dataset_num_proc)991 train_dataset = train_dataset.map(992 maybe_apply_chat_template, fn_kwargs={"tokenizer": processing_class}, num_proc=args.dataset_num_proc993 )994 if eval_dataset is not None:995 eval_dataset = eval_dataset.map(maybe_extract_prompt, num_proc=args.dataset_num_proc)996 eval_dataset = eval_dataset.map(997 maybe_apply_chat_template,998 fn_kwargs={"tokenizer": processing_class},999 num_proc=args.dataset_num_proc,1000 )1001 1002 # tokenize the dataset1003 train_dataset = train_dataset.map(self.tokenize_row, num_proc=args.dataset_num_proc)1004 if eval_dataset is not None:1005 eval_dataset = eval_dataset.map(self.tokenize_row, num_proc=args.dataset_num_proc)1006 1007 super().__init__(1008 model=model,1009 args=args,1010 data_collator=data_collator,1011 train_dataset=train_dataset,1012 eval_dataset=eval_dataset,1013 processing_class=processing_class,1014 model_init=model_init,1015 compute_metrics=compute_metrics,1016 callbacks=callbacks,1017 optimizers=optimizers,1018 preprocess_logits_for_metrics=preprocess_logits_for_metrics,1019 )1020 1021 # Gradient accumulation requires scaled loss. Normally, loss scaling in the parent class depends on whether the1022 # model accepts loss-related kwargs. Since we compute our own loss, this check is irrelevant. We set1023 # self.model_accepts_loss_kwargs to False to enable scaling.1024 self.model_accepts_loss_kwargs = False1025 1026 # Add tags for models that have been loaded with the correct transformers version1027 if hasattr(self.model, "add_model_tags"):1028 self.model.add_model_tags(self._tag_names)1029 1030 if not hasattr(self, "accelerator"):1031 raise AttributeError(1032 "Your `Trainer` does not have an `accelerator` object. Consider upgrading `transformers`."1033 )1034 1035 def build_tokenized_answer(self, prompt, answer):1036 """1037 Llama tokenizer does satisfy `enc(a + b) = enc(a) + enc(b)`. It does ensure `enc(a + b) = enc(a) + enc(a +1038 b)[len(enc(a)):]`. Reference:1039 https://github.com/EleutherAI/lm-evaluation-harness/pull/531#issuecomment-15955862571040 """1041 1042 full_tokenized = self.processing_class(prompt + answer, add_special_tokens=False)1043 prompt_input_ids = self.processing_class(prompt, add_special_tokens=False)["input_ids"]1044 1045 answer_input_ids = full_tokenized["input_ids"][len(prompt_input_ids) :]1046 answer_attention_mask = full_tokenized["attention_mask"][len(prompt_input_ids) :]1047 1048 # Concat tokens to form `enc(a) + enc(a + b)[len(enc(a)):]`1049 full_concat_input_ids = np.concatenate([prompt_input_ids, answer_input_ids])1050 1051 # Prepare input tokens for token by token comparison1052 full_input_ids = np.array(full_tokenized["input_ids"])1053 1054 if len(full_input_ids) != len(full_concat_input_ids):1055 raise ValueError("Prompt input ids and answer input ids should have the same length.")1056 1057 # On some tokenizers, like Llama-2 tokenizer, there are occasions where tokens1058 # can be merged together when tokenizing prompt+answer. This could result1059 # on the last token from the prompt being different when tokenized on its own1060 # vs when done as prompt+answer.1061 response_token_ids_start_idx = len(prompt_input_ids)1062 1063 # If tokenized prompt is different than both prompt+answer, then it means the1064 # last token has changed due to merging.1065 if prompt_input_ids != full_tokenized["input_ids"][:response_token_ids_start_idx]:1066 response_token_ids_start_idx -= 11067 1068 prompt_input_ids = full_tokenized["input_ids"][:response_token_ids_start_idx]1069 prompt_attention_mask = full_tokenized["attention_mask"][:response_token_ids_start_idx]1070 1071 if len(prompt_input_ids) != len(prompt_attention_mask):1072 raise ValueError("Prompt input ids and attention mask should have the same length.")1073 1074 answer_input_ids = full_tokenized["input_ids"][response_token_ids_start_idx:]1075 answer_attention_mask = full_tokenized["attention_mask"][response_token_ids_start_idx:]1076 1077 return dict(1078 prompt_input_ids=prompt_input_ids,1079 prompt_attention_mask=prompt_attention_mask,1080 input_ids=answer_input_ids,1081 attention_mask=answer_attention_mask,1082 )1083 1084 def tokenize_row(self, feature, model: Optional[Union[PreTrainedModel, nn.Module]] = None) -> dict:1085 """Tokenize a single row from a CPO specific dataset.1086 1087 At this stage, we don't convert to PyTorch tensors yet; we just handle the truncation in case the prompt +1088 chosen or prompt + rejected responses is/are too long. First we truncate the prompt; if we're still too long,1089 we truncate the chosen/rejected.1090 1091 We also create the labels for the chosen/rejected responses, which are of length equal to the sum of the length1092 of the prompt and the chosen/rejected response, with label_pad_token_id for the prompt tokens.1093 """1094 batch = {}1095 prompt = feature["prompt"]1096 chosen = feature["chosen"]1097 rejected = feature["rejected"]1098 1099 if not self.is_encoder_decoder:1100 # Check issues below for more details1101 # 1. https://github.com/huggingface/trl/issues/9071102 # 2. https://github.com/EleutherAI/lm-evaluation-harness/pull/531#issuecomment-15955862571103 # 3. https://github.com/LianjiaTech/BELLE/issues/3371104 1105 if not isinstance(prompt, str):1106 raise ValueError(f"prompt should be an str but got {type(prompt)}")1107 prompt_tokens = self.processing_class(prompt, add_special_tokens=False)1108 prompt_tokens = {f"prompt_{k}": v for k, v in prompt_tokens.items()}1109 1110 if not isinstance(chosen, str):1111 raise ValueError(f"chosen should be an str but got {type(chosen)}")1112 chosen_tokens = self.build_tokenized_answer(prompt, chosen)1113 1114 if not isinstance(rejected, str):1115 raise ValueError(f"rejected should be an str but got {type(rejected)}")1116 rejected_tokens = self.build_tokenized_answer(prompt, rejected)1117 1118 # Last prompt token might get merged by tokenizer and1119 # it should not be included for generation if that happens1120 prompt_len_input_ids = len(prompt_tokens["prompt_input_ids"])1121 1122 chosen_prompt_len_input_ids = len(chosen_tokens["prompt_input_ids"])1123 rejected_prompt_len_input_ids = len(rejected_tokens["prompt_input_ids"])1124 prompt_len_input_ids = min(chosen_prompt_len_input_ids, rejected_prompt_len_input_ids)1125 1126 for k, v in prompt_tokens.items():1127 prompt_tokens[k] = v[:prompt_len_input_ids]1128 1129 # Make sure prompts only have one different token at most an1130 # and length only differs by 1 at most1131 num_diff_tokens = sum(1132 [a != b for a, b in zip(chosen_tokens["prompt_input_ids"], rejected_tokens["prompt_input_ids"])]1133 )1134 num_diff_len = abs(chosen_prompt_len_input_ids - rejected_prompt_len_input_ids)1135 if num_diff_tokens > 1 or num_diff_len > 1:1136 raise ValueError(1137 "Chosen and rejected prompt_input_ids might only differ on the "1138 "last token due to tokenizer merge ops."1139 )1140 1141 # add BOS token to head of prompt. Avoid adding if it's already there1142 prompt_tokens, chosen_tokens, rejected_tokens = add_bos_token_if_needed(1143 self.processing_class.bos_token_id,1144 prompt_len_input_ids,1145 prompt_tokens,1146 chosen_prompt_len_input_ids,1147 chosen_tokens,1148 rejected_prompt_len_input_ids,1149 rejected_tokens,1150 )1151 1152 # add EOS token to end of answer. Avoid adding if it's already there1153 chosen_tokens, rejected_tokens = add_eos_token_if_needed(1154 self.processing_class.eos_token_id, chosen_tokens, rejected_tokens1155 )1156 1157 longer_response_length = max(len(chosen_tokens["input_ids"]), len(rejected_tokens["input_ids"]))1158 1159 # if combined sequence is too long, truncate the prompt1160 for answer_tokens in [chosen_tokens, rejected_tokens, prompt_tokens]:1161 if len(answer_tokens["prompt_input_ids"]) + longer_response_length > self.max_length:1162 if self.truncation_mode == "keep_start":1163 for k in ["prompt_input_ids", "prompt_attention_mask"]:1164 answer_tokens[k] = answer_tokens[k][: self.max_prompt_length]1165 elif self.truncation_mode == "keep_end":1166 for k in ["prompt_input_ids", "prompt_attention_mask"]:1167 answer_tokens[k] = answer_tokens[k][-self.max_prompt_length :]1168 else:1169 raise ValueError(f"Unknown truncation mode: {self.truncation_mode}")1170 1171 # if that's still too long, truncate the response1172 for answer_tokens in [chosen_tokens, rejected_tokens]:1173 if len(answer_tokens["prompt_input_ids"]) + longer_response_length > self.max_length:1174 for k in ["input_ids", "attention_mask"]:1175 answer_tokens[k] = answer_tokens[k][: self.max_length - self.max_prompt_length]1176 1177 # Create labels1178 chosen_sequence_tokens = {1179 k: chosen_tokens[f"prompt_{k}"] + chosen_tokens[k] for k in ["input_ids", "attention_mask"]1180 }1181 rejected_sequence_tokens = {1182 k: rejected_tokens[f"prompt_{k}"] + rejected_tokens[k] for k in ["input_ids", "attention_mask"]1183 }1184 chosen_sequence_tokens["labels"] = chosen_sequence_tokens["input_ids"][:]1185 chosen_sequence_tokens["labels"][: len(chosen_tokens["prompt_input_ids"])] = [1186 self.label_pad_token_id1187 ] * len(chosen_tokens["prompt_input_ids"])1188 rejected_sequence_tokens["labels"] = rejected_sequence_tokens["input_ids"][:]1189 rejected_sequence_tokens["labels"][: len(rejected_tokens["prompt_input_ids"])] = [1190 self.label_pad_token_id1191 ] * len(rejected_tokens["prompt_input_ids"])1192 1193 for k, toks in {1194 "chosen_": chosen_sequence_tokens,1195 "rejected_": rejected_sequence_tokens,1196 "": prompt_tokens,1197 }.items():1198 for type_key, tokens in toks.items():1199 if type_key == "token_type_ids":1200 continue