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.orpo_trainer import (Any, AutoModelForCausalLM, BaseImageProcessor, Callable, DPODataCollatorWithPadding, DataCollator, DataLoader, Dataset, EvalLoopOutput, F, FeatureExtractionMixin, Literal, ORPOConfig, ORPOTrainer, Optional, PartialState, Path, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, Trainer, TrainerCallback, Union, add_bos_token_if_needed, add_eos_token_if_needed, autocast, defaultdict, disable_dropout_in_model, generate_model_card, get_comet_experiment_url, inspect, is_comet_available, is_peft_available, is_torch_fx_proxy, is_torch_xla_available, is_wandb_available, log_table_to_comet_experiment, 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, Callable, DPODataCollatorWithPadding, DataCollator, Dataset, EvalLoopOutput, F, FeatureExtractionMixin, ORPOConfig, ORPOTrainer, 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 UnslothORPOConfig(ORPOConfig):340 """341 342 Configuration class for the [`ORPOTrainer`].343 344 This class includes only the parameters that are specific to ORPO 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 relative ratio loss weight in the ORPO loss. In the363 [paper](https://huggingface.co/papers/2403.07691), it is denoted by λ. In the364 [code](https://github.com/xfactlab/orpo), it is denoted by `alpha`.365 disable_dropout (`bool`, *optional*, defaults to `True`):366 Whether to disable dropout in the model.367 label_pad_token_id (`int`, *optional*, defaults to `-100`):368 Label pad token id. This argument is required if you want to use the default data collator.369 padding_value (`int` or `None`, *optional*, defaults to `None`):370 Padding value to use. If `None`, the padding value of the tokenizer is used.371 truncation_mode (`str`, *optional*, defaults to `"keep_end"`):372 Truncation mode to use when the prompt is too long. Possible values are `"keep_end"` or `"keep_start"`.373 This argument is required if you want to use the default data collator.374 generate_during_eval (`bool`, *optional*, defaults to `False`):375 If `True`, generates and logs completions from the model to W&B or Comet during evaluation.376 is_encoder_decoder (`bool` or `None`, *optional*, defaults to `None`):377 When using the `model_init` argument (callable) to instantiate the model instead of the `model` argument,378 you need to specify if the model returned by the callable is an encoder-decoder model.379 model_init_kwargs (`dict[str, Any]` or `None`, *optional*, defaults to `None`):380 Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model from a381 string.382 dataset_num_proc (`int` or `None`, *optional*, defaults to `None`):383 Number of processes to use for processing the dataset.384 385 """386 vllm_sampling_params: Optional[Any] = field(387 default = None,388 metadata = {'help': 'vLLM SamplingParams'},389 )390 unsloth_num_chunks : Optional[int] = field(391 default = -1,392 metadata = {'help': 'Chunk size to reduce memory usage. -1 is most efficient.'},393 )394 unsloth_logit_chunk_multiplier : Optional[int] = field(395 default = None,396 metadata = {'help': 'Multiplier for chunked logit computations.'},397 )398 unsloth_grpo_mini_batch : Optional[int] = field(399 default = None,400 metadata = {'help': 'Mini batch size for GRPO hidden state accumulation. Default is None unless user defines it.'},401 )402 max_seq_length : Optional[int] = field(403 default = None,404 metadata = {'help': 'Maximum sequence length to truncate to.'},405 )406 def __init__(407 self,408 output_dir = None,409 overwrite_output_dir = None,410 do_train = False,411 do_eval = False,412 do_predict = False,413 eval_strategy = 'no',414 prediction_loss_only = False,415 per_device_train_batch_size = 4,416 per_device_eval_batch_size = 4,417 per_gpu_train_batch_size = None,418 per_gpu_eval_batch_size = None,419 gradient_accumulation_steps = 2,420 eval_accumulation_steps = 2,421 eval_delay = 0,422 torch_empty_cache_steps = 250,423 learning_rate = 5e-05,424 weight_decay = 0.01,425 adam_beta1 = 0.9,426 adam_beta2 = 0.999,427 adam_epsilon = 1e-08,428 max_grad_norm = 1.0,429 num_train_epochs = 3.0,430 max_steps = -1,431 lr_scheduler_type = 'linear',432 warmup_ratio = 0.1,433 warmup_steps = 0,434 log_level = 'passive',435 log_level_replica = 'warning',436 log_on_each_node = True,437 logging_dir = None,438 logging_strategy = 'steps',439 logging_first_step = False,440 logging_steps = 1,441 logging_nan_inf_filter = False,442 save_strategy = 'steps',443 save_steps = 500,444 save_total_limit = None,445 save_safetensors = True,446 save_on_each_node = False,447 save_only_model = False,448 restore_callback_states_from_checkpoint = False,449 no_cuda = False,450 use_cpu = False,451 use_mps_device = False,452 seed = 3407,453 data_seed = 3407,454 jit_mode_eval = False,455 bf16 = False,456 fp16 = False,457 fp16_opt_level = 'O1',458 half_precision_backend = 'auto',459 bf16_full_eval = False,460 fp16_full_eval = False,461 tf32 = None,462 local_rank = -1,463 ddp_backend = None,464 tpu_num_cores = None,465 tpu_metrics_debug = False,466 debug = '',467 dataloader_drop_last = False,468 eval_steps = None,469 dataloader_num_workers = 0,470 dataloader_prefetch_factor = None,471 past_index = -1,472 run_name = None,473 disable_tqdm = None,474 remove_unused_columns = True,475 label_names = None,476 load_best_model_at_end = False,477 metric_for_best_model = None,478 greater_is_better = None,479 ignore_data_skip = False,480 fsdp = None,481 fsdp_min_num_params = 0,482 fsdp_config = None,483 fsdp_transformer_layer_cls_to_wrap = None,484 accelerator_config = None,485 parallelism_config = None,486 deepspeed = None,487 label_smoothing_factor = 0.0,488 optim = 'adamw_8bit',489 optim_args = None,490 adafactor = False,491 group_by_length = False,492 length_column_name = 'length',493 report_to = 'none',494 project = 'huggingface',495 trackio_space_id = 'trackio',496 ddp_find_unused_parameters = None,497 ddp_bucket_cap_mb = None,498 ddp_broadcast_buffers = None,499 dataloader_pin_memory = True,500 dataloader_persistent_workers = False,501 skip_memory_metrics = True,502 use_legacy_prediction_loop = False,503 push_to_hub = False,504 resume_from_checkpoint = None,505 hub_model_id = None,506 hub_strategy = 'every_save',507 hub_token = None,508 hub_private_repo = None,509 hub_always_push = False,510 hub_revision = None,511 gradient_checkpointing = True,512 gradient_checkpointing_kwargs = None,513 include_inputs_for_metrics = False,514 eval_do_concat_batches = True,515 fp16_backend = 'auto',516 push_to_hub_model_id = None,517 push_to_hub_organization = None,518 push_to_hub_token = None,519 mp_parameters = '',520 auto_find_batch_size = False,521 full_determinism = False,522 torchdynamo = None,523 ray_scope = 'last',524 ddp_timeout = 1800,525 torch_compile = False,526 torch_compile_backend = None,527 torch_compile_mode = None,528 include_tokens_per_second = False,529 include_num_input_tokens_seen = False,530 neftune_noise_alpha = None,531 optim_target_modules = None,532 batch_eval_metrics = False,533 eval_on_start = False,534 use_liger_kernel = False,535 liger_kernel_config = None,536 eval_use_gather_object = False,537 average_tokens_across_devices = True,538 max_length = 1024,539 max_prompt_length = 512,540 max_completion_length = None,541 beta = 0.1,542 disable_dropout = True,543 label_pad_token_id = -100,544 padding_value = None,545 truncation_mode = 'keep_end',546 generate_during_eval = False,547 is_encoder_decoder = None,548 model_init_kwargs = None,549 dataset_num_proc = None,550 vllm_sampling_params = None,551 unsloth_num_chunks = -1,552 unsloth_logit_chunk_multiplier = None,553 unsloth_grpo_mini_batch = None,554 max_seq_length = None,555 **kwargs,556 ):557 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!')558 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!')559 if num_train_epochs is None:560 num_train_epochs = 3.0 # Default to 3 epochs if None, max_steps will override561 if output_dir is None and save_strategy == 'steps' and save_steps == 500:562 output_dir = 'unsloth_training_checkpoints'563 save_strategy = 'no'564 import multiprocessing as _mp565 if dataset_num_proc is None:566 if _mp.get_start_method() != 'fork':567 dataset_num_proc = None568 else:569 import psutil570 dataset_num_proc = min(max((psutil.cpu_count() or 1)+4, 2), 64)571 memory_gb_left = psutil.virtual_memory().available / (1024**3)572 if memory_gb_left <= 2: dataset_num_proc = 1573 else: dataset_num_proc = min(dataset_num_proc, int(memory_gb_left))574 575 super().__init__(576 output_dir = output_dir,577 overwrite_output_dir = overwrite_output_dir,578 do_train = do_train,579 do_eval = do_eval,580 do_predict = do_predict,581 eval_strategy = eval_strategy,582 prediction_loss_only = prediction_loss_only,583 per_device_train_batch_size = per_device_train_batch_size,584 per_device_eval_batch_size = per_device_eval_batch_size,585 per_gpu_train_batch_size = per_gpu_train_batch_size,586 per_gpu_eval_batch_size = per_gpu_eval_batch_size,587 gradient_accumulation_steps = gradient_accumulation_steps,588 eval_accumulation_steps = eval_accumulation_steps,589 eval_delay = eval_delay,590 torch_empty_cache_steps = torch_empty_cache_steps,591 learning_rate = learning_rate,592 weight_decay = weight_decay,593 adam_beta1 = adam_beta1,594 adam_beta2 = adam_beta2,595 adam_epsilon = adam_epsilon,596 max_grad_norm = max_grad_norm,597 num_train_epochs = num_train_epochs,598 max_steps = max_steps,599 lr_scheduler_type = lr_scheduler_type,600 warmup_ratio = warmup_ratio,601 warmup_steps = warmup_steps,602 log_level = log_level,603 log_level_replica = log_level_replica,604 log_on_each_node = log_on_each_node,605 logging_dir = logging_dir,606 logging_strategy = logging_strategy,607 logging_first_step = logging_first_step,608 logging_steps = logging_steps,609 logging_nan_inf_filter = logging_nan_inf_filter,610 save_strategy = save_strategy,611 save_steps = save_steps,612 save_total_limit = save_total_limit,613 save_safetensors = save_safetensors,614 save_on_each_node = save_on_each_node,615 save_only_model = save_only_model,616 restore_callback_states_from_checkpoint = restore_callback_states_from_checkpoint,617 no_cuda = no_cuda,618 use_cpu = use_cpu,619 use_mps_device = use_mps_device,620 seed = seed,621 data_seed = data_seed,622 jit_mode_eval = jit_mode_eval,623 bf16 = bf16,624 fp16 = fp16,625 fp16_opt_level = fp16_opt_level,626 half_precision_backend = half_precision_backend,627 bf16_full_eval = bf16_full_eval,628 fp16_full_eval = fp16_full_eval,629 tf32 = tf32,630 local_rank = local_rank,631 ddp_backend = ddp_backend,632 tpu_num_cores = tpu_num_cores,633 tpu_metrics_debug = tpu_metrics_debug,634 debug = debug,635 dataloader_drop_last = dataloader_drop_last,636 eval_steps = eval_steps,637 dataloader_num_workers = dataloader_num_workers,638 dataloader_prefetch_factor = dataloader_prefetch_factor,639 past_index = past_index,640 run_name = run_name,641 disable_tqdm = disable_tqdm,642 remove_unused_columns = remove_unused_columns,643 label_names = label_names,644 load_best_model_at_end = load_best_model_at_end,645 metric_for_best_model = metric_for_best_model,646 greater_is_better = greater_is_better,647 ignore_data_skip = ignore_data_skip,648 fsdp = fsdp,649 fsdp_min_num_params = fsdp_min_num_params,650 fsdp_config = fsdp_config,651 fsdp_transformer_layer_cls_to_wrap = fsdp_transformer_layer_cls_to_wrap,652 accelerator_config = accelerator_config,653 parallelism_config = parallelism_config,654 deepspeed = deepspeed,655 label_smoothing_factor = label_smoothing_factor,656 optim = optim,657 optim_args = optim_args,658 adafactor = adafactor,659 group_by_length = group_by_length,660 length_column_name = length_column_name,661 report_to = report_to,662 project = project,663 trackio_space_id = trackio_space_id,664 ddp_find_unused_parameters = ddp_find_unused_parameters,665 ddp_bucket_cap_mb = ddp_bucket_cap_mb,666 ddp_broadcast_buffers = ddp_broadcast_buffers,667 dataloader_pin_memory = dataloader_pin_memory,668 dataloader_persistent_workers = dataloader_persistent_workers,669 skip_memory_metrics = skip_memory_metrics,670 use_legacy_prediction_loop = use_legacy_prediction_loop,671 push_to_hub = push_to_hub,672 resume_from_checkpoint = resume_from_checkpoint,673 hub_model_id = hub_model_id,674 hub_strategy = hub_strategy,675 hub_token = hub_token,676 hub_private_repo = hub_private_repo,677 hub_always_push = hub_always_push,678 hub_revision = hub_revision,679 gradient_checkpointing = gradient_checkpointing,680 gradient_checkpointing_kwargs = gradient_checkpointing_kwargs,681 include_inputs_for_metrics = include_inputs_for_metrics,682 eval_do_concat_batches = eval_do_concat_batches,683 fp16_backend = fp16_backend,684 push_to_hub_model_id = push_to_hub_model_id,685 push_to_hub_organization = push_to_hub_organization,686 push_to_hub_token = push_to_hub_token,687 mp_parameters = mp_parameters,688 auto_find_batch_size = auto_find_batch_size,689 full_determinism = full_determinism,690 torchdynamo = torchdynamo,691 ray_scope = ray_scope,692 ddp_timeout = ddp_timeout,693 torch_compile = torch_compile,694 torch_compile_backend = torch_compile_backend,695 torch_compile_mode = torch_compile_mode,696 include_tokens_per_second = include_tokens_per_second,697 include_num_input_tokens_seen = include_num_input_tokens_seen,698 neftune_noise_alpha = neftune_noise_alpha,699 optim_target_modules = optim_target_modules,700 batch_eval_metrics = batch_eval_metrics,701 eval_on_start = eval_on_start,702 use_liger_kernel = use_liger_kernel,703 liger_kernel_config = liger_kernel_config,704 eval_use_gather_object = eval_use_gather_object,705 average_tokens_across_devices = average_tokens_across_devices,706 max_length = max_length,707 max_prompt_length = max_prompt_length,708 max_completion_length = max_completion_length,709 beta = beta,710 disable_dropout = disable_dropout,711 label_pad_token_id = label_pad_token_id,712 padding_value = padding_value,713 truncation_mode = truncation_mode,714 generate_during_eval = generate_during_eval,715 is_encoder_decoder = is_encoder_decoder,716 model_init_kwargs = model_init_kwargs,717 dataset_num_proc = dataset_num_proc,**kwargs)718 self.vllm_sampling_params = vllm_sampling_params719 self.unsloth_num_chunks = unsloth_num_chunks720 if unsloth_grpo_mini_batch is not None:721 if self.generation_batch_size >= unsloth_grpo_mini_batch:722 self.unsloth_grpo_mini_batch = unsloth_grpo_mini_batch723 else:724 raise ValueError(725 f"Unsloth GRPO mini batch size needs to be less than or equal to the effective generation batch size, "726 f"which is self.per_device_train_batch_size * gradient_accumulation_steps."727 )728 self.unsloth_logit_chunk_multiplier = unsloth_logit_chunk_multiplier729 self.max_seq_length = max_seq_length730 731pass732 733class _UnslothORPOTrainer(Trainer):734 r""""""735 736 _tag_names = ["trl", "orpo"]737 738 def __init__(739 self,740 model: Optional[Union[PreTrainedModel, nn.Module, str]] = None,741 args: Optional[ORPOConfig] = None,742 data_collator: Optional[DataCollator] = None,743 train_dataset: Optional[Dataset] = None,744 eval_dataset: Optional[Union[Dataset, dict[str, Dataset]]] = None,745 processing_class: Optional[746 Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin]747 ] = None,748 model_init: Optional[Callable[[], PreTrainedModel]] = None,749 callbacks: Optional[list[TrainerCallback]] = None,750 optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),751 preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None,752 peft_config: Optional[dict] = None,753 compute_metrics: Optional[Callable[[EvalLoopOutput], dict]] = None,754 ):755 if args.model_init_kwargs is None:756 model_init_kwargs = {}757 elif not isinstance(model, str):758 raise ValueError("You passed model_kwargs to the ORPOTrainer. But your model is already instantiated.")759 else:760 model_init_kwargs = args.model_init_kwargs761 dtype = model_init_kwargs.get("dtype")762 if dtype is not None:763 # Convert to `torch.dtype` if an str is passed764 if isinstance(dtype, str) and dtype != "auto":765 dtype = getattr(torch, dtype)766 if dtype != "auto" and not isinstance(dtype, torch.dtype):767 raise ValueError(768 f"Invalid `dtype` passed to the ORPOConfig. Expected a string with either `torch.dtype` or 'auto', but got {dtype}."769 )770 model_init_kwargs["dtype"] = dtype771 772 if isinstance(model, str):773 model = AutoModelForCausalLM.from_pretrained(model, **model_init_kwargs)774 775 # Initialize this variable to False. This helps tracking the case when `peft_module_casting_to_bf16`776 # has been called in order to properly call autocast if needed.777 self._peft_has_been_casted_to_bf16 = False778 779 if not is_peft_available() and peft_config is not None:780 raise ValueError(781 "PEFT is not installed and you passed a `peft_config` in the trainer's kwargs, please install it to use the PEFT models"782 )783 elif is_peft_available() and peft_config is not None:784 # if model is a peft model and we have a peft_config, we merge and unload it first785 if isinstance(model, PeftModel):786 model = model.merge_and_unload()787 788 if getattr(model, "is_loaded_in_8bit", False) or getattr(model, "is_loaded_in_4bit", False):789 _support_gc_kwargs = hasattr(790 args, "gradient_checkpointing_kwargs"791 ) and "gradient_checkpointing_kwargs" in list(792 inspect.signature(prepare_model_for_kbit_training).parameters793 )794 795 prepare_model_kwargs = {"use_gradient_checkpointing": args.gradient_checkpointing}796 797 if _support_gc_kwargs:798 prepare_model_kwargs["gradient_checkpointing_kwargs"] = args.gradient_checkpointing_kwargs799 800 model = prepare_model_for_kbit_training(model, **prepare_model_kwargs)801 elif args.gradient_checkpointing:802 # For backward compatibility with older versions of transformers803 if hasattr(model, "enable_input_require_grads"):804 model.enable_input_require_grads()805 else:806 807 def make_inputs_require_grad(module, input, output):808 output.requires_grad_(True)809 810 model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)811 812 # get peft model with the given config813 model = model814 if args.bf16 and getattr(model, "is_loaded_in_4bit", False):815 peft_module_casting_to_bf16(model)816 # If args.bf16 we need to explicitly call `generate` with torch amp autocast context manager817 self._peft_has_been_casted_to_bf16 = True818 819 # For models that use gradient_checkpointing, we need to attach a hook that enables input820 # to explicitly have `requires_grad=True`, otherwise training will either silently821 # fail or completely fail.822 elif args.gradient_checkpointing:823 # For backward compatibility with older versions of transformers824 if hasattr(model, "enable_input_require_grads"):825 model.enable_input_require_grads()826 else:827 828 def make_inputs_require_grad(module, input, output):829 output.requires_grad_(True)830 831 model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)832 833 if args.generate_during_eval and not (is_wandb_available() or is_comet_available()):834 raise ValueError(835 "`generate_during_eval=True` requires Weights and Biases or Comet to be installed."836 " Please install `wandb` or `comet-ml` to resolve."837 )838 839 if model is not None:840 self.is_encoder_decoder = model.config.is_encoder_decoder841 elif args.is_encoder_decoder is None:842 raise ValueError("When no model is provided, you need to pass the parameter is_encoder_decoder.")843 else:844 self.is_encoder_decoder = args.is_encoder_decoder845 846 if self.is_encoder_decoder:847 self.decoder_start_token_id = model.config.decoder_start_token_id848 self.pad_token_id = model.config.pad_token_id849 850 if processing_class is None:851 raise ValueError("processing_class must be specified to tokenize a ORPO dataset.")852 if args.max_length is None:853 logger.warning(854 "`max_length` is not set in the ORPOConfig's init"855 " it will default to `512` by default, but you should do it yourself in the future.",856 )857 max_length = 512858 else:859 max_length = args.max_length860 if args.max_prompt_length is None:861 logger.warning(862 "`max_prompt_length` is not set in the ORPOConfig's init"863 " it will default to `128` by default, but you should do it yourself in the future.",864 )865 max_prompt_length = 128866 else:867 max_prompt_length = args.max_prompt_length868 869 if args.max_completion_length is None and self.is_encoder_decoder:870 logger.warning(871 "When using an encoder decoder architecture, you should set `max_completion_length` in the ORPOConfig's init"872 " it will default to `128` by default, but you should do it yourself in the future.",873 )874 self.max_completion_length = 128875 else:876 self.max_completion_length = args.max_completion_length877 878 if data_collator is None:879 data_collator = DPODataCollatorWithPadding(880 pad_token_id=processing_class.pad_token_id,881 label_pad_token_id=args.label_pad_token_id,882 is_encoder_decoder=self.is_encoder_decoder,883 )884 885 if args.remove_unused_columns:886 args.remove_unused_columns = False887 # warn users888 logger.warning(889 "When using DPODataCollatorWithPadding, you should set `remove_unused_columns=False` in your TrainingArguments"890 " we have set it for you, but you should do it yourself in the future.",891 )892 893 self.use_dpo_data_collator = True894 else:895 self.use_dpo_data_collator = False896 897 # Disable dropout in the model and reference model898 if args.disable_dropout:899 disable_dropout_in_model(model)900 901 self.max_length = max_length902 self.generate_during_eval = args.generate_during_eval903 self.label_pad_token_id = args.label_pad_token_id904 self.padding_value = args.padding_value if args.padding_value is not None else processing_class.pad_token_id905 self.max_prompt_length = max_prompt_length906 self.truncation_mode = args.truncation_mode907 self.processing_class = processing_class908 909 self.beta = args.beta910 self.aux_loss_enabled = getattr(model.config, "output_router_logits", False)911 self.aux_loss_coef = getattr(model.config, "router_aux_loss_coef", 0.0)912 if self.aux_loss_enabled and self.aux_loss_coef == 0.0:913 logger.warning(914 "You set `output_router_logits` to `True` in the model config, but `router_aux_loss_coef` is set to "915 "`0.0`, meaning the auxiliary loss will not be used. Either set `router_aux_loss_coef` to a value "916 "greater than `0.0`, or set `output_router_logits` to `False` if you don't want to use the auxiliary "917 "loss.",918 )919 920 self._stored_metrics = defaultdict(lambda: defaultdict(list))921 922 # The trainer estimates the number of FLOPs [floating-point operations] using the number of elements in the923 # input tensor associated with the key "input_ids". However, in ORPO, the sampled data does not include the924 # "input_ids" key. Instead, the available keys are "prompt_input_ids", "chosen_input_ids", and925 # "rejected_input_ids". As a result, the trainer issues the warning: "Could not estimate the number of tokens926 # of the input, floating-point operations will not be computed." To suppress this warning, we set the927 # "estimate_tokens" key in the model's "warnings_issued" dictionary to True. This acts as a flag to indicate928 # that the warning has already been issued.929 model.warnings_issued["estimate_tokens"] = True930 931 # Compute that only on the main process for faster data processing.932 # see: https://github.com/huggingface/trl/pull/1255933 with PartialState().main_process_first():934 # Extract the prompt if needed, and apply the chat template if needed935 train_dataset = train_dataset.map(maybe_extract_prompt, num_proc=args.dataset_num_proc)936 train_dataset = train_dataset.map(937 maybe_apply_chat_template, fn_kwargs={"tokenizer": processing_class}, num_proc=args.dataset_num_proc938 )939 train_dataset = train_dataset.map(self.tokenize_row, num_proc=args.dataset_num_proc)940 if eval_dataset is not None:941 eval_dataset = eval_dataset.map(maybe_extract_prompt, num_proc=args.dataset_num_proc)942 eval_dataset = eval_dataset.map(943 maybe_apply_chat_template,944 fn_kwargs={"tokenizer": processing_class},945 num_proc=args.dataset_num_proc,946 )947 eval_dataset = eval_dataset.map(self.tokenize_row, num_proc=args.dataset_num_proc)948 949 super().__init__(950 model=model,951 args=args,952 data_collator=data_collator,953 train_dataset=train_dataset,954 eval_dataset=eval_dataset,955 processing_class=processing_class,956 model_init=model_init,957 compute_metrics=compute_metrics,958 callbacks=callbacks,959 optimizers=optimizers,960 preprocess_logits_for_metrics=preprocess_logits_for_metrics,961 )962 963 # Gradient accumulation requires scaled loss. Normally, loss scaling in the parent class depends on whether the964 # model accepts loss-related kwargs. Since we compute our own loss, this check is irrelevant. We set965 # self.model_accepts_loss_kwargs to False to enable scaling.966 self.model_accepts_loss_kwargs = False967 968 # Add tags for models that have been loaded with the correct transformers version969 if hasattr(self.model, "add_model_tags"):970 self.model.add_model_tags(self._tag_names)971 972 if not hasattr(self, "accelerator"):973 raise AttributeError(974 "Your `Trainer` does not have an `accelerator` object. Consider upgrading `transformers`."975 )976 977 def build_tokenized_answer(self, prompt, answer):978 """979 Llama tokenizer does satisfy `enc(a + b) = enc(a) + enc(b)`. It does ensure `enc(a + b) = enc(a) + enc(a +980 b)[len(enc(a)):]`. Reference:981 https://github.com/EleutherAI/lm-evaluation-harness/pull/531#issuecomment-1595586257982 """983 984 full_tokenized = self.processing_class(prompt + answer, add_special_tokens=False)985 prompt_input_ids = self.processing_class(prompt, add_special_tokens=False)["input_ids"]986 987 answer_input_ids = full_tokenized["input_ids"][len(prompt_input_ids) :]988 answer_attention_mask = full_tokenized["attention_mask"][len(prompt_input_ids) :]989 990 # Concat tokens to form `enc(a) + enc(a + b)[len(enc(a)):]`991 full_concat_input_ids = np.concatenate([prompt_input_ids, answer_input_ids])992 993 # Prepare input tokens for token by token comparison994 full_input_ids = np.array(full_tokenized["input_ids"])995 996 if len(full_input_ids) != len(full_concat_input_ids):997 raise ValueError("Prompt input ids and answer input ids should have the same length.")998 999 # On some tokenizers, like Llama-2 tokenizer, there are occasions where tokens1000 # can be merged together when tokenizing prompt+answer. This could result1001 # on the last token from the prompt being different when tokenized on its own1002 # vs when done as prompt+answer.1003 response_token_ids_start_idx = len(prompt_input_ids)1004 1005 # If tokenized prompt is different than both prompt+answer, then it means the1006 # last token has changed due to merging.1007 if prompt_input_ids != full_tokenized["input_ids"][:response_token_ids_start_idx]:1008 response_token_ids_start_idx -= 11009 1010 prompt_input_ids = full_tokenized["input_ids"][:response_token_ids_start_idx]1011 prompt_attention_mask = full_tokenized["attention_mask"][:response_token_ids_start_idx]1012 1013 if len(prompt_input_ids) != len(prompt_attention_mask):1014 raise ValueError("Prompt input ids and attention mask should have the same length.")1015 1016 answer_input_ids = full_tokenized["input_ids"][response_token_ids_start_idx:]1017 answer_attention_mask = full_tokenized["attention_mask"][response_token_ids_start_idx:]1018 1019 return dict(1020 prompt_input_ids=prompt_input_ids,1021 prompt_attention_mask=prompt_attention_mask,1022 input_ids=answer_input_ids,1023 attention_mask=answer_attention_mask,1024 )1025 1026 def tokenize_row(self, feature, model: Optional[Union[PreTrainedModel, nn.Module]] = None) -> dict:1027 """Tokenize a single row from a ORPO specific dataset.1028 1029 At this stage, we don't convert to PyTorch tensors yet; we just handle the truncation in case the prompt +1030 chosen or prompt + rejected responses is/are too long. First we truncate the prompt; if we're still too long,1031 we truncate the chosen/rejected.1032 1033 We also create the labels for the chosen/rejected responses, which are of length equal to the sum of the length1034 of the prompt and the chosen/rejected response, with label_pad_token_id for the prompt tokens.1035 """1036 batch = {}1037 prompt = feature["prompt"]1038 chosen = feature["chosen"]1039 rejected = feature["rejected"]1040 1041 if not self.is_encoder_decoder:1042 # Check issues below for more details1043 # 1. https://github.com/huggingface/trl/issues/9071044 # 2. https://github.com/EleutherAI/lm-evaluation-harness/pull/531#issuecomment-15955862571045 # 3. https://github.com/LianjiaTech/BELLE/issues/3371046 1047 if not isinstance(prompt, str):1048 raise ValueError(f"prompt should be an str but got {type(prompt)}")1049 prompt_tokens = self.processing_class(prompt, add_special_tokens=False)1050 prompt_tokens = {f"prompt_{k}": v for k, v in prompt_tokens.items()}1051 1052 if not isinstance(chosen, str):1053 raise ValueError(f"chosen should be an str but got {type(chosen)}")1054 chosen_tokens = self.build_tokenized_answer(prompt, chosen)1055 1056 if not isinstance(rejected, str):1057 raise ValueError(f"rejected should be an str but got {type(rejected)}")1058 rejected_tokens = self.build_tokenized_answer(prompt, rejected)1059 1060 # Last prompt token might get merged by tokenizer and1061 # it should not be included for generation if that happens1062 prompt_len_input_ids = len(prompt_tokens["prompt_input_ids"])1063 1064 chosen_prompt_len_input_ids = len(chosen_tokens["prompt_input_ids"])1065 rejected_prompt_len_input_ids = len(rejected_tokens["prompt_input_ids"])1066 prompt_len_input_ids = min(chosen_prompt_len_input_ids, rejected_prompt_len_input_ids)1067 1068 for k, v in prompt_tokens.items():1069 prompt_tokens[k] = v[:prompt_len_input_ids]1070 1071 # Make sure prompts only have one different token at most an1072 # and length only differs by 1 at most1073 num_diff_tokens = sum(1074 [a != b for a, b in zip(chosen_tokens["prompt_input_ids"], rejected_tokens["prompt_input_ids"])]1075 )1076 num_diff_len = abs(chosen_prompt_len_input_ids - rejected_prompt_len_input_ids)1077 if num_diff_tokens > 1 or num_diff_len > 1:1078 raise ValueError(1079 "Chosen and rejected prompt_input_ids might only differ on the "1080 "last token due to tokenizer merge ops."1081 )1082 1083 # add BOS token to head of prompt. Avoid adding if it's already there1084 prompt_tokens, chosen_tokens, rejected_tokens = add_bos_token_if_needed(1085 self.processing_class.bos_token_id,1086 prompt_len_input_ids,1087 prompt_tokens,1088 chosen_prompt_len_input_ids,1089 chosen_tokens,1090 rejected_prompt_len_input_ids,1091 rejected_tokens,1092 )1093 1094 # add EOS token to end of answer. Avoid adding if it's already there1095 chosen_tokens, rejected_tokens = add_eos_token_if_needed(1096 self.processing_class.eos_token_id, chosen_tokens, rejected_tokens1097 )1098 1099 longer_response_length = max(len(chosen_tokens["input_ids"]), len(rejected_tokens["input_ids"]))1100 1101 # if combined sequence is too long, truncate the prompt1102 for answer_tokens in [chosen_tokens, rejected_tokens, prompt_tokens]:1103 if len(answer_tokens["prompt_input_ids"]) + longer_response_length > self.max_length:1104 if self.truncation_mode == "keep_start":1105 for k in ["prompt_input_ids", "prompt_attention_mask"]:1106 answer_tokens[k] = answer_tokens[k][: self.max_prompt_length]1107 elif self.truncation_mode == "keep_end":1108 for k in ["prompt_input_ids", "prompt_attention_mask"]:1109 answer_tokens[k] = answer_tokens[k][-self.max_prompt_length :]1110 else:1111 raise ValueError(f"Unknown truncation mode: {self.truncation_mode}")1112 1113 # if that's still too long, truncate the response1114 for answer_tokens in [chosen_tokens, rejected_tokens]:1115 if len(answer_tokens["prompt_input_ids"]) + longer_response_length > self.max_length:1116 for k in ["input_ids", "attention_mask"]:1117 answer_tokens[k] = answer_tokens[k][: self.max_length - self.max_prompt_length]1118 1119 # Create labels1120 chosen_sequence_tokens = {1121 k: chosen_tokens[f"prompt_{k}"] + chosen_tokens[k] for k in ["input_ids", "attention_mask"]1122 }1123 rejected_sequence_tokens = {1124 k: rejected_tokens[f"prompt_{k}"] + rejected_tokens[k] for k in ["input_ids", "attention_mask"]1125 }1126 chosen_sequence_tokens["labels"] = chosen_sequence_tokens["input_ids"][:]1127 chosen_sequence_tokens["labels"][: len(chosen_tokens["prompt_input_ids"])] = [1128 self.label_pad_token_id1129 ] * len(chosen_tokens["prompt_input_ids"])1130 rejected_sequence_tokens["labels"] = rejected_sequence_tokens["input_ids"][:]1131 rejected_sequence_tokens["labels"][: len(rejected_tokens["prompt_input_ids"])] = [1132 self.label_pad_token_id1133 ] * len(rejected_tokens["prompt_input_ids"])1134 1135 for k, toks in {1136 "chosen_": chosen_sequence_tokens,1137 "rejected_": rejected_sequence_tokens,1138 "": prompt_tokens,1139 }.items():1140 for type_key, tokens in toks.items():1141 if type_key == "token_type_ids":1142 continue1143 batch[f"{k}{type_key}"] = tokens1144 1145 else:1146 chosen_tokens = self.processing_class(1147 chosen, truncation=True, max_length=self.max_completion_length, add_special_tokens=True1148 )1149 rejected_tokens = self.processing_class(1150 rejected, truncation=True, max_length=self.max_completion_length, add_special_tokens=True1151 )1152 prompt_tokens = self.processing_class(1153 prompt, truncation=True, max_length=self.max_prompt_length, add_special_tokens=True1154 )1155 1156 batch["chosen_labels"] = chosen_tokens["input_ids"]1157 batch["rejected_labels"] = rejected_tokens["input_ids"]1158 batch["prompt_input_ids"] = prompt_tokens["input_ids"]1159 batch["prompt_attention_mask"] = prompt_tokens["attention_mask"]1160 1161 if model is not None and hasattr(model, "prepare_decoder_input_ids_from_labels"):1162 batch["rejected_decoder_input_ids"] = model.prepare_decoder_input_ids_from_labels(1163 labels=torch.tensor(batch["rejected_labels"])1164 )1165 batch["chosen_decoder_input_ids"] = model.prepare_decoder_input_ids_from_labels(1166 labels=torch.tensor(batch["chosen_labels"])1167 )1168 1169 if is_torch_xla_available():1170 # Pad the sequences to global max_length to avoid TorchXLA recompilation1171 for k in batch:1172 if "labels" in k or self.is_encoder_decoder:1173 pad_value = self.label_pad_token_id1174 elif k.endswith("_input_ids"):1175 pad_value = self.padding_value1176 elif k.endswith("_attention_mask"):1177 pad_value = 01178 batch[k] = batch[k] + [pad_value] * (self.max_length - len(batch[k]))1179 return batch1180 1181 @staticmethod1182 def concatenated_inputs(1183 batch: dict[str, Union[list, torch.LongTensor]],1184 is_encoder_decoder: bool = False,1185 label_pad_token_id: int = -100,1186 padding_value: int = 0,1187 device: Optional[torch.device] = None,1188 ) -> dict[str, torch.LongTensor]:1189 """Concatenate the chosen and rejected inputs into a single tensor.1190 1191 Args:1192 batch:1193 A batch of data. Must contain the keys 'chosen_input_ids' and 'rejected_input_ids', which are tensors1194 of shape (batch_size, sequence_length).1195 is_encoder_decoder:1196 Whether the model is an encoder-decoder model.1197 label_pad_token_id:1198 The label pad token id.1199 padding_value:1200 The padding value to use for the concatenated inputs_ids.