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.xpo_trainer import (Any, BaseImageProcessor, BasePairwiseJudge, Callable, Dataset, EvalPrediction, F, FeatureExtractionMixin, IterableDataset, OnlineDPOTrainer, OptimizerNames, Optional, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, SIMPLE_CHAT_TEMPLATE, TrainerCallback, Union, XPOConfig, XPOTrainer, empty_cache, generate_model_card, get_comet_experiment_url, get_reward, is_conversational, is_peft_available, is_wandb_available, jinja2, maybe_apply_chat_template, nn, os, selective_log_softmax, textwrap, torch, truncate_right, unwrap_model_for_generation)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 UnslothXPOConfig(XPOConfig):340 """341 342 Configuration class for the [`XPOTrainer`].343 344 Subclass of [`OnlineDPOConfig`] we can use all its arguments and add the following:345 346 Parameters:347 alpha (`float` or `list[float]`, *optional*, defaults to `1e-5`):348 Weight of the XPO loss term. If a list of floats is provided then the alpha is selected for each new epoch349 and the last alpha is used for the rest of the epochs.350 351 """352 vllm_sampling_params: Optional[Any] = field(353 default = None,354 metadata = {'help': 'vLLM SamplingParams'},355 )356 unsloth_num_chunks : Optional[int] = field(357 default = -1,358 metadata = {'help': 'Chunk size to reduce memory usage. -1 is most efficient.'},359 )360 unsloth_logit_chunk_multiplier : Optional[int] = field(361 default = None,362 metadata = {'help': 'Multiplier for chunked logit computations.'},363 )364 unsloth_grpo_mini_batch : Optional[int] = field(365 default = None,366 metadata = {'help': 'Mini batch size for GRPO hidden state accumulation. Default is None unless user defines it.'},367 )368 max_seq_length : Optional[int] = field(369 default = None,370 metadata = {'help': 'Maximum sequence length to truncate to.'},371 )372 def __init__(373 self,374 output_dir = None,375 overwrite_output_dir = None,376 do_train = False,377 do_eval = False,378 do_predict = False,379 eval_strategy = 'no',380 prediction_loss_only = False,381 per_device_train_batch_size = 4,382 per_device_eval_batch_size = 4,383 per_gpu_train_batch_size = None,384 per_gpu_eval_batch_size = None,385 gradient_accumulation_steps = 2,386 eval_accumulation_steps = 2,387 eval_delay = 0,388 torch_empty_cache_steps = 250,389 learning_rate = 5e-05,390 weight_decay = 0.01,391 adam_beta1 = 0.9,392 adam_beta2 = 0.999,393 adam_epsilon = 1e-08,394 max_grad_norm = 1.0,395 num_train_epochs = 3.0,396 max_steps = -1,397 lr_scheduler_type = 'linear',398 warmup_ratio = 0.1,399 warmup_steps = 0,400 log_level = 'passive',401 log_level_replica = 'warning',402 log_on_each_node = True,403 logging_dir = None,404 logging_strategy = 'steps',405 logging_first_step = False,406 logging_steps = 1,407 logging_nan_inf_filter = False,408 save_strategy = 'steps',409 save_steps = 500,410 save_total_limit = None,411 save_safetensors = True,412 save_on_each_node = False,413 save_only_model = False,414 restore_callback_states_from_checkpoint = False,415 no_cuda = False,416 use_cpu = False,417 use_mps_device = False,418 seed = 3407,419 data_seed = 3407,420 jit_mode_eval = False,421 bf16 = False,422 fp16 = False,423 fp16_opt_level = 'O1',424 half_precision_backend = 'auto',425 bf16_full_eval = False,426 fp16_full_eval = False,427 tf32 = None,428 local_rank = -1,429 ddp_backend = None,430 tpu_num_cores = None,431 tpu_metrics_debug = False,432 debug = '',433 dataloader_drop_last = False,434 eval_steps = None,435 dataloader_num_workers = 0,436 dataloader_prefetch_factor = None,437 past_index = -1,438 run_name = None,439 disable_tqdm = None,440 remove_unused_columns = True,441 label_names = None,442 load_best_model_at_end = False,443 metric_for_best_model = None,444 greater_is_better = None,445 ignore_data_skip = False,446 fsdp = None,447 fsdp_min_num_params = 0,448 fsdp_config = None,449 fsdp_transformer_layer_cls_to_wrap = None,450 accelerator_config = None,451 parallelism_config = None,452 deepspeed = None,453 label_smoothing_factor = 0.0,454 optim = 'adamw_8bit',455 optim_args = None,456 adafactor = False,457 group_by_length = False,458 length_column_name = 'length',459 report_to = 'none',460 project = 'huggingface',461 trackio_space_id = 'trackio',462 ddp_find_unused_parameters = None,463 ddp_bucket_cap_mb = None,464 ddp_broadcast_buffers = None,465 dataloader_pin_memory = True,466 dataloader_persistent_workers = False,467 skip_memory_metrics = True,468 use_legacy_prediction_loop = False,469 push_to_hub = False,470 resume_from_checkpoint = None,471 hub_model_id = None,472 hub_strategy = 'every_save',473 hub_token = None,474 hub_private_repo = None,475 hub_always_push = False,476 hub_revision = None,477 gradient_checkpointing = True,478 gradient_checkpointing_kwargs = None,479 include_inputs_for_metrics = False,480 eval_do_concat_batches = True,481 fp16_backend = 'auto',482 push_to_hub_model_id = None,483 push_to_hub_organization = None,484 push_to_hub_token = None,485 mp_parameters = '',486 auto_find_batch_size = False,487 full_determinism = False,488 torchdynamo = None,489 ray_scope = 'last',490 ddp_timeout = 1800,491 torch_compile = False,492 torch_compile_backend = None,493 torch_compile_mode = None,494 include_tokens_per_second = False,495 include_num_input_tokens_seen = False,496 neftune_noise_alpha = None,497 optim_target_modules = None,498 batch_eval_metrics = False,499 eval_on_start = False,500 use_liger_kernel = False,501 liger_kernel_config = None,502 eval_use_gather_object = False,503 average_tokens_across_devices = True,504 reward_model_path = None,505 judge = None,506 max_new_tokens = 64,507 max_length = 512,508 temperature = 0.9,509 top_p = 1.0,510 top_k = None,511 min_p = None,512 repetition_penalty = 1.0,513 generation_kwargs = {},514 use_transformers_paged = False,515 cache_implementation = None,516 missing_eos_penalty = None,517 loss_type = 'sigmoid',518 disable_dropout = True,519 use_vllm = False,520 vllm_model_impl = 'vllm',521 vllm_guided_decoding_regex = None,522 vllm_gpu_memory_utilization = 0.55,523 vllm_mode = 'colocate',524 vllm_server_base_url = None,525 vllm_server_host = '0.0.0.0',526 vllm_server_port = 8000,527 vllm_server_timeout = 240.0,528 vllm_tensor_parallel_size = 1,529 ds3_gather_for_generation = True,530 model_init_kwargs = None,531 reward_weights = None,532 dataset_num_proc = None,533 gpu_memory_utilization = None,534 vllm_sampling_params = None,535 unsloth_num_chunks = -1,536 unsloth_logit_chunk_multiplier = None,537 unsloth_grpo_mini_batch = None,538 max_seq_length = None,539 **kwargs,540 ):541 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!')542 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!')543 if num_train_epochs is None:544 num_train_epochs = 3.0 # Default to 3 epochs if None, max_steps will override545 if output_dir is None and save_strategy == 'steps' and save_steps == 500:546 output_dir = 'unsloth_training_checkpoints'547 save_strategy = 'no'548 import multiprocessing as _mp549 if dataset_num_proc is None:550 if _mp.get_start_method() != 'fork':551 dataset_num_proc = None552 else:553 import psutil554 dataset_num_proc = min(max((psutil.cpu_count() or 1)+4, 2), 64)555 memory_gb_left = psutil.virtual_memory().available / (1024**3)556 if memory_gb_left <= 2: dataset_num_proc = 1557 else: dataset_num_proc = min(dataset_num_proc, int(memory_gb_left))558 if temperature <= 0:559 raise ValueError('Unsloth: Please set a positive non-zero temperature since your results will be wrong.')560 elif temperature >= 10:561 raise ValueError('Unsloth: Please set a positive non-zero temperature less than 10, since sampling will be quite erratic.')562 563 564 super().__init__(565 output_dir = output_dir,566 overwrite_output_dir = overwrite_output_dir,567 do_train = do_train,568 do_eval = do_eval,569 do_predict = do_predict,570 eval_strategy = eval_strategy,571 prediction_loss_only = prediction_loss_only,572 per_device_train_batch_size = per_device_train_batch_size,573 per_device_eval_batch_size = per_device_eval_batch_size,574 per_gpu_train_batch_size = per_gpu_train_batch_size,575 per_gpu_eval_batch_size = per_gpu_eval_batch_size,576 gradient_accumulation_steps = gradient_accumulation_steps,577 eval_accumulation_steps = eval_accumulation_steps,578 eval_delay = eval_delay,579 torch_empty_cache_steps = torch_empty_cache_steps,580 learning_rate = learning_rate,581 weight_decay = weight_decay,582 adam_beta1 = adam_beta1,583 adam_beta2 = adam_beta2,584 adam_epsilon = adam_epsilon,585 max_grad_norm = max_grad_norm,586 num_train_epochs = num_train_epochs,587 max_steps = max_steps,588 lr_scheduler_type = lr_scheduler_type,589 warmup_ratio = warmup_ratio,590 warmup_steps = warmup_steps,591 log_level = log_level,592 log_level_replica = log_level_replica,593 log_on_each_node = log_on_each_node,594 logging_dir = logging_dir,595 logging_strategy = logging_strategy,596 logging_first_step = logging_first_step,597 logging_steps = logging_steps,598 logging_nan_inf_filter = logging_nan_inf_filter,599 save_strategy = save_strategy,600 save_steps = save_steps,601 save_total_limit = save_total_limit,602 save_safetensors = save_safetensors,603 save_on_each_node = save_on_each_node,604 save_only_model = save_only_model,605 restore_callback_states_from_checkpoint = restore_callback_states_from_checkpoint,606 no_cuda = no_cuda,607 use_cpu = use_cpu,608 use_mps_device = use_mps_device,609 seed = seed,610 data_seed = data_seed,611 jit_mode_eval = jit_mode_eval,612 bf16 = bf16,613 fp16 = fp16,614 fp16_opt_level = fp16_opt_level,615 half_precision_backend = half_precision_backend,616 bf16_full_eval = bf16_full_eval,617 fp16_full_eval = fp16_full_eval,618 tf32 = tf32,619 local_rank = local_rank,620 ddp_backend = ddp_backend,621 tpu_num_cores = tpu_num_cores,622 tpu_metrics_debug = tpu_metrics_debug,623 debug = debug,624 dataloader_drop_last = dataloader_drop_last,625 eval_steps = eval_steps,626 dataloader_num_workers = dataloader_num_workers,627 dataloader_prefetch_factor = dataloader_prefetch_factor,628 past_index = past_index,629 run_name = run_name,630 disable_tqdm = disable_tqdm,631 remove_unused_columns = remove_unused_columns,632 label_names = label_names,633 load_best_model_at_end = load_best_model_at_end,634 metric_for_best_model = metric_for_best_model,635 greater_is_better = greater_is_better,636 ignore_data_skip = ignore_data_skip,637 fsdp = fsdp,638 fsdp_min_num_params = fsdp_min_num_params,639 fsdp_config = fsdp_config,640 fsdp_transformer_layer_cls_to_wrap = fsdp_transformer_layer_cls_to_wrap,641 accelerator_config = accelerator_config,642 parallelism_config = parallelism_config,643 deepspeed = deepspeed,644 label_smoothing_factor = label_smoothing_factor,645 optim = optim,646 optim_args = optim_args,647 adafactor = adafactor,648 group_by_length = group_by_length,649 length_column_name = length_column_name,650 report_to = report_to,651 project = project,652 trackio_space_id = trackio_space_id,653 ddp_find_unused_parameters = ddp_find_unused_parameters,654 ddp_bucket_cap_mb = ddp_bucket_cap_mb,655 ddp_broadcast_buffers = ddp_broadcast_buffers,656 dataloader_pin_memory = dataloader_pin_memory,657 dataloader_persistent_workers = dataloader_persistent_workers,658 skip_memory_metrics = skip_memory_metrics,659 use_legacy_prediction_loop = use_legacy_prediction_loop,660 push_to_hub = push_to_hub,661 resume_from_checkpoint = resume_from_checkpoint,662 hub_model_id = hub_model_id,663 hub_strategy = hub_strategy,664 hub_token = hub_token,665 hub_private_repo = hub_private_repo,666 hub_always_push = hub_always_push,667 hub_revision = hub_revision,668 gradient_checkpointing = gradient_checkpointing,669 gradient_checkpointing_kwargs = gradient_checkpointing_kwargs,670 include_inputs_for_metrics = include_inputs_for_metrics,671 eval_do_concat_batches = eval_do_concat_batches,672 fp16_backend = fp16_backend,673 push_to_hub_model_id = push_to_hub_model_id,674 push_to_hub_organization = push_to_hub_organization,675 push_to_hub_token = push_to_hub_token,676 mp_parameters = mp_parameters,677 auto_find_batch_size = auto_find_batch_size,678 full_determinism = full_determinism,679 torchdynamo = torchdynamo,680 ray_scope = ray_scope,681 ddp_timeout = ddp_timeout,682 torch_compile = torch_compile,683 torch_compile_backend = torch_compile_backend,684 torch_compile_mode = torch_compile_mode,685 include_tokens_per_second = include_tokens_per_second,686 include_num_input_tokens_seen = include_num_input_tokens_seen,687 neftune_noise_alpha = neftune_noise_alpha,688 optim_target_modules = optim_target_modules,689 batch_eval_metrics = batch_eval_metrics,690 eval_on_start = eval_on_start,691 use_liger_kernel = use_liger_kernel,692 liger_kernel_config = liger_kernel_config,693 eval_use_gather_object = eval_use_gather_object,694 average_tokens_across_devices = average_tokens_across_devices,695 reward_model_path = reward_model_path,696 judge = judge,697 max_new_tokens = max_new_tokens,698 max_length = max_length,699 temperature = temperature,700 top_p = top_p,701 top_k = top_k,702 min_p = min_p,703 repetition_penalty = repetition_penalty,704 generation_kwargs = generation_kwargs,705 use_transformers_paged = use_transformers_paged,706 cache_implementation = cache_implementation,707 missing_eos_penalty = missing_eos_penalty,708 loss_type = loss_type,709 disable_dropout = disable_dropout,710 use_vllm = use_vllm,711 vllm_model_impl = vllm_model_impl,712 vllm_guided_decoding_regex = vllm_guided_decoding_regex,713 vllm_gpu_memory_utilization = vllm_gpu_memory_utilization,714 vllm_mode = vllm_mode,715 vllm_server_base_url = vllm_server_base_url,716 vllm_server_host = vllm_server_host,717 vllm_server_port = vllm_server_port,718 vllm_server_timeout = vllm_server_timeout,719 vllm_tensor_parallel_size = vllm_tensor_parallel_size,720 ds3_gather_for_generation = ds3_gather_for_generation,721 model_init_kwargs = model_init_kwargs,722 reward_weights = reward_weights,723 dataset_num_proc = dataset_num_proc,724 gpu_memory_utilization = gpu_memory_utilization,**kwargs)725 self.vllm_sampling_params = vllm_sampling_params726 self.unsloth_num_chunks = unsloth_num_chunks727 if unsloth_grpo_mini_batch is not None:728 if self.generation_batch_size >= unsloth_grpo_mini_batch:729 self.unsloth_grpo_mini_batch = unsloth_grpo_mini_batch730 else:731 raise ValueError(732 f"Unsloth GRPO mini batch size needs to be less than or equal to the effective generation batch size, "733 f"which is self.per_device_train_batch_size * gradient_accumulation_steps."734 )735 self.unsloth_logit_chunk_multiplier = unsloth_logit_chunk_multiplier736 self.max_seq_length = max_seq_length737 738pass739 740class _UnslothXPOTrainer(OnlineDPOTrainer):741 r""""""742 743 _tag_names = ["trl", "xpo"]744 745 def __init__(746 self,747 model: Union[PreTrainedModel, nn.Module] = None,748 ref_model: Union[PreTrainedModel, nn.Module] = None,749 reward_funcs: Optional[nn.Module] = None,750 judge: Optional[BasePairwiseJudge] = None,751 args: Optional[XPOConfig] = None,752 data_collator: Optional[Callable] = None,753 train_dataset: Optional[Union[Dataset, IterableDataset]] = None,754 eval_dataset: Optional[Union[Dataset, dict[str, Dataset]]] = None,755 processing_class: Optional[756 Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin]757 ] = None,758 reward_processing_classes: Optional[Union[PreTrainedTokenizerBase, list[PreTrainedTokenizerBase]]] = None,759 peft_config: Optional[dict] = None,760 compute_metrics: Optional[Callable[[EvalPrediction], dict]] = None,761 callbacks: Optional[list[TrainerCallback]] = None,762 optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),763 preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None,764 # Deprecated parameters765 reward_model: Optional[Union[PreTrainedModel, nn.Module]] = None,766 ) -> None:767 super().__init__(768 model=model,769 ref_model=ref_model,770 judge=judge,771 reward_funcs=reward_funcs,772 reward_model=reward_model,773 args=args,774 data_collator=data_collator,775 train_dataset=train_dataset,776 eval_dataset=eval_dataset,777 processing_class=processing_class,778 reward_processing_classes=reward_processing_classes,779 peft_config=peft_config,780 compute_metrics=compute_metrics,781 callbacks=callbacks,782 optimizers=optimizers,783 preprocess_logits_for_metrics=preprocess_logits_for_metrics,784 )785 786 self._alpha = self.args.alpha787 788 # Overwrite the stats dictionary to include XPO specific statistics789 self.stats = {790 # Remove "non_score_reward", "rlhf_reward", "scores"791 # Add "loss/dpo", "loss/xpo"792 "loss/dpo": [],793 "loss/xpo": [],794 "objective/kl": [],795 "objective/entropy": [],796 "rewards/chosen": [],797 "rewards/rejected": [],798 "rewards/accuracies": [],799 "rewards/margins": [],800 "logps/chosen": [],801 "logps/rejected": [],802 # Replace "contain_eos_token" by "model_contain_eos_token" and "ref_contain_eos_token"803 "val/model_contain_eos_token": [],804 "val/ref_contain_eos_token": [],805 "alpha": [],806 "beta": [],807 }808 if self.reward_funcs is not None:809 if len(self.reward_funcs) != 1:810 raise ValueError("XPOTrainer only supports one reward function/model.")811 self.reward_funcs = self.reward_funcs[0]812 self.stats["objective/model_scores"] = []813 self.stats["objective/ref_scores"] = []814 self.stats["objective/scores_margin"] = []815 816 @property817 def alpha(self):818 if isinstance(self._alpha, list):819 epoch = self.state.epoch820 return self._alpha[epoch] if epoch < len(self._alpha) else self._alpha[-1]821 else:822 return self._alpha823 824 def _generate_completions(self, prompts, model):825 with unwrap_model_for_generation(model, self.accelerator) as unwrapped_policy_model_for_gen:826 model_output = unwrapped_policy_model_for_gen.generate(827 input_ids=prompts["input_ids"],828 attention_mask=prompts["attention_mask"],829 generation_config=self.generation_config,830 )831 832 actual_model_for_ref_generation: torch.nn.Module833 if self.ref_model is None:834 unwrapped_main_model_for_ref_logic = self.accelerator.unwrap_model(model)835 836 if is_peft_available() and isinstance(unwrapped_main_model_for_ref_logic, PeftModel):837 actual_model_for_ref_generation = unwrapped_main_model_for_ref_logic.get_base_model()838 else:839 actual_model_for_ref_generation = unwrapped_main_model_for_ref_logic840 else:841 actual_model_for_ref_generation = self.accelerator.unwrap_model(self.ref_model)842 843 with unwrap_model_for_generation(actual_model_for_ref_generation, self.accelerator) as final_ref_model_for_gen:844 ref_output = final_ref_model_for_gen.generate(845 input_ids=prompts["input_ids"],846 attention_mask=prompts["attention_mask"],847 generation_config=self.generation_config,848 )849 850 return model_output, ref_output851 852 def _process_completions(self, model_output, ref_output, prompts):853 context_length = prompts["input_ids"].shape[1]854 855 # Process model completions856 model_completion_ids = model_output[:, context_length:]857 model_completion_ids, model_completion_mask = truncate_right(858 model_completion_ids, self.processing_class.eos_token_id, self.processing_class.pad_token_id859 )860 model_data = {861 "input_ids": torch.cat((prompts["input_ids"], model_completion_ids), dim=1),862 "attention_mask": torch.cat((prompts["attention_mask"], model_completion_mask), dim=1),863 "raw": prompts["raw"],864 }865 866 # Process reference model completions867 ref_completion_ids = ref_output[:, context_length:]868 ref_completion_ids, ref_completion_mask = truncate_right(869 ref_completion_ids, self.processing_class.eos_token_id, self.processing_class.pad_token_id870 )871 ref_data = {872 "input_ids": torch.cat((prompts["input_ids"], ref_completion_ids), dim=1),873 "attention_mask": torch.cat((prompts["attention_mask"], ref_completion_mask), dim=1),874 "raw": prompts["raw"],875 }876 877 return model_data, ref_data878 879 def _compute_rewards(self, model_data, ref_data, context_length):880 with torch.no_grad():881 _, model_scores, _ = get_reward(882 self.reward_funcs, model_data["input_ids"], self.processing_class.pad_token_id, context_length883 )884 _, ref_scores, _ = get_reward(885 self.reward_funcs, ref_data["input_ids"], self.processing_class.pad_token_id, context_length886 )887 888 # Apply EOS penalty if needed889 if self.args.missing_eos_penalty is not None:890 model_contain_eos = torch.any(model_data["input_ids"] == self.processing_class.eos_token_id, dim=-1)891 ref_contain_eos = torch.any(ref_data["input_ids"] == self.processing_class.eos_token_id, dim=-1)892 model_scores[~model_contain_eos] -= self.args.missing_eos_penalty893 ref_scores[~ref_contain_eos] -= self.args.missing_eos_penalty894 895 return model_scores, ref_scores896 897 def _compute_judge(self, model_data, ref_data, context_length):898 prompts = model_data["raw"]899 model_data_completions = self.processing_class.batch_decode(900 model_data["input_ids"][:, context_length:], skip_special_tokens=True901 )902 model_data_completions = [completion.strip() for completion in model_data_completions]903 904 ref_data_completions = self.processing_class.batch_decode(905 ref_data["input_ids"][:, context_length:], skip_special_tokens=True906 )907 ref_data_completions = [completion.strip() for completion in ref_data_completions]908 909 if is_conversational({"prompt": prompts[0]}):910 model_data_completions = [911 [{"role": "assistant", "content": completion}] for completion in model_data_completions912 ]913 environment = jinja2.Environment()914 template = environment.from_string(SIMPLE_CHAT_TEMPLATE)915 prompts = [template.render(messages=message) for message in prompts]916 model_data_completions = [template.render(messages=completion) for completion in model_data_completions]917 918 ref_data_completions = [919 [{"role": "assistant", "content": completion}] for completion in ref_data_completions920 ]921 ref_data_completions = [template.render(messages=completion) for completion in ref_data_completions]922 923 ranks_of_first_completion = self.judge.judge(924 prompts,925 list(zip(model_data_completions, ref_data_completions)),926 )927 # convert ranks to a True/False mask:928 # when rank == 0, it means the first completion is the best929 # when rank == 1, it means the second completion is the best930 return torch.tensor([rank == 0 for rank in ranks_of_first_completion], device=model_data["input_ids"].device)931 932 def _compute_logprobs(self, model, model_data, ref_data, context_length):933 def compute_logprobs_for_data(m, data):934 output = m(data["input_ids"], attention_mask=data["attention_mask"])935 logits = output.logits[:, context_length - 1 : -1]936 token_logprobs = selective_log_softmax(logits, data["input_ids"][:, context_length:])937 return token_logprobs938 939 # Compute logprobs for model completions940 model_logprobs_model_data = compute_logprobs_for_data(model, model_data)941 # Compute logprobs for model on reference completions (for XPO loss)942 model_logprobs_ref_data = compute_logprobs_for_data(model, ref_data)943 944 # Compute logprobs for reference model completions945 with torch.no_grad():946 if self.ref_model is None:947 with model.disable_adapter():948 ref_logprobs_model_data = compute_logprobs_for_data(model, model_data)949 ref_logprobs_ref_data = compute_logprobs_for_data(model, ref_data)950 else:951 ref_logprobs_model_data = compute_logprobs_for_data(self.ref_model, model_data)952 ref_logprobs_ref_data = compute_logprobs_for_data(self.ref_model, ref_data)953 954 # Mask padding tokens955 model_padding_mask = model_data["attention_mask"][:, context_length:] == 0956 ref_padding_mask = ref_data["attention_mask"][:, context_length:] == 0957 model_logprobs_model_data = model_logprobs_model_data.masked_fill(model_padding_mask, 0.0)958 model_logprobs_ref_data = model_logprobs_ref_data.masked_fill(ref_padding_mask, 0.0)959 ref_logprobs_ref_data = ref_logprobs_ref_data.masked_fill(ref_padding_mask, 0.0)960 ref_logprobs_model_data = ref_logprobs_model_data.masked_fill(model_padding_mask, 0.0)961 962 return model_logprobs_model_data, model_logprobs_ref_data, ref_logprobs_ref_data, ref_logprobs_model_data963 964 def _compute_losses(965 self,966 model_logprobs_model_data,967 model_logprobs_ref_data,968 ref_logprobs_ref_data,969 ref_logprobs_model_data,970 chosen_mask,971 ):972 # Compute log probs973 model_logprobs_model_data_sum = model_logprobs_model_data.sum(1)974 model_logprobs_ref_data_sum = model_logprobs_ref_data.sum(1)975 ref_logprobs_ref_data_sum = ref_logprobs_ref_data.sum(1)976 ref_logprobs_model_data_sum = ref_logprobs_model_data.sum(1)977 978 chosen_model_logprobs = torch.where(chosen_mask, model_logprobs_model_data_sum, model_logprobs_ref_data_sum)979 chosen_ref_logprobs = torch.where(chosen_mask, ref_logprobs_model_data_sum, ref_logprobs_ref_data_sum)980 chosen_log_ratios = chosen_model_logprobs - chosen_ref_logprobs981 982 rejected_model_logprobs = torch.where(~chosen_mask, model_logprobs_model_data_sum, model_logprobs_ref_data_sum)983 rejected_ref_logprobs = torch.where(~chosen_mask, ref_logprobs_model_data_sum, ref_logprobs_ref_data_sum)984 rejected_log_ratios = rejected_model_logprobs - rejected_ref_logprobs985 986 # Compute logits as the difference between chosen and rejected log ratios987 logits = chosen_log_ratios - rejected_log_ratios988 989 if self.args.loss_type == "sigmoid":990 dpo_losses = -F.logsigmoid(self.beta * logits)991 elif self.args.loss_type == "ipo":992 dpo_losses = (logits - 1 / (2 * self.beta)) ** 2993 else:994 raise NotImplementedError(f"invalid loss type {self.args.loss_type}")995 996 # Compute XPO specific loss997 xpo_losses = self.alpha * model_logprobs_ref_data_sum998 999 # Total loss1000 loss = (dpo_losses + xpo_losses).mean()1001 1002 return loss, dpo_losses, xpo_losses1003 1004 def _log_statistics(1005 self,1006 model_data,1007 ref_data,1008 model_logprobs_model_data,1009 model_logprobs_ref_data,1010 ref_logprobs_ref_data,1011 ref_logprobs_model_data,1012 chosen_mask,1013 dpo_losses,1014 xpo_losses,1015 context_length,1016 model_scores=None,1017 ref_scores=None,1018 ):1019 # Helper function to gather and compute mean1020 def gather_mean(tensor):1021 return self.accelerator.gather_for_metrics(tensor).mean().item()1022 1023 # Log losses1024 self.stats["loss/dpo"].append(gather_mean(dpo_losses))1025 self.stats["loss/xpo"].append(gather_mean(xpo_losses))1026 1027 # Log scores1028 if self.reward_funcs is not None:1029 self.stats["objective/model_scores"].append(gather_mean(model_scores))1030 self.stats["objective/ref_scores"].append(gather_mean(ref_scores))1031 self.stats["objective/scores_margin"].append(gather_mean(model_scores - ref_scores))1032 1033 # Log logprobs1034 model_logprobs_model_data_sum = model_logprobs_model_data.sum(1)1035 model_logprobs_ref_data_sum = model_logprobs_ref_data.sum(1)1036 ref_logprobs_ref_data_sum = ref_logprobs_ref_data.sum(1)1037 ref_logprobs_model_data_sum = ref_logprobs_model_data.sum(1)1038 1039 chosen_model_logprobs = torch.where(chosen_mask, model_logprobs_model_data_sum, model_logprobs_ref_data_sum)1040 chosen_ref_logprobs = torch.where(chosen_mask, ref_logprobs_model_data_sum, ref_logprobs_ref_data_sum)1041 chosen_log_ratios = chosen_model_logprobs - chosen_ref_logprobs1042 1043 rejected_model_logprobs = torch.where(~chosen_mask, model_logprobs_model_data_sum, model_logprobs_ref_data_sum)1044 rejected_ref_logprobs = torch.where(~chosen_mask, ref_logprobs_model_data_sum, ref_logprobs_ref_data_sum)1045 rejected_log_ratios = rejected_model_logprobs - rejected_ref_logprobs1046 1047 self.stats["logps/chosen"].append(gather_mean(chosen_model_logprobs.mean() + chosen_ref_logprobs.mean()))1048 self.stats["logps/rejected"].append(gather_mean(rejected_model_logprobs.mean() + rejected_ref_logprobs.mean()))1049 1050 # Log rewards1051 # Compute various statistics1052 chosen_rewards = chosen_log_ratios * self.beta1053 rejected_rewards = rejected_log_ratios * self.beta1054 self.stats["rewards/chosen"].append(gather_mean(chosen_rewards.mean()))1055 self.stats["rewards/rejected"].append(gather_mean(rejected_rewards.mean()))1056 1057 # Calculate KL divergence for model and ref data1058 kl_model_data = model_logprobs_model_data - ref_logprobs_model_data1059 kl_ref_data = model_logprobs_ref_data - ref_logprobs_ref_data1060 mean_kl = (kl_model_data.sum(1) + kl_ref_data.sum(1)).mean() / 21061 self.stats["objective/kl"].append(gather_mean(mean_kl))1062 1063 # Calculate entropy for model and ref data1064 entropy_model_data = -model_logprobs_model_data.sum(1)1065 entropy_ref_data = -model_logprobs_ref_data.sum(1)1066 mean_entropy = (entropy_model_data.mean() + entropy_ref_data.mean()) / 21067 self.stats["objective/entropy"].append(gather_mean(mean_entropy))1068 1069 # Calculate margins1070 margin = chosen_rewards - rejected_rewards1071 self.stats["rewards/margins"].append(gather_mean(margin.mean()))1072 1073 # Calculate accuracy1074 accuracy = (margin > 0).float()1075 self.stats["rewards/accuracies"].append(gather_mean(accuracy.mean()))1076 1077 # Log EOS token statistics1078 model_eos = (model_data["input_ids"][:, context_length:] == self.processing_class.eos_token_id).any(dim=1)1079 ref_eos = (ref_data["input_ids"][:, context_length:] == self.processing_class.eos_token_id).any(dim=1)1080 self.stats["val/model_contain_eos_token"].append(gather_mean(model_eos.float()))1081 self.stats["val/ref_contain_eos_token"].append(gather_mean(ref_eos.float()))1082 1083 # Log alpha and beta1084 self.stats["alpha"].append(self.alpha)1085 self.stats["beta"].append(self.beta)1086 1087 def training_step(1088 self, model: nn.Module, inputs: dict[str, Union[torch.Tensor, Any]], num_items_in_batch: Optional[int] = None1089 ) -> torch.Tensor:1090 model.train()1091 1092 # Apply chat template and tokenize the input1093 batch_size = len(next(iter(inputs.values())))1094 prompts = inputs["prompt"]1095 inputs = [{k: v[i] for k, v in inputs.items()} for i in range(batch_size)]1096 inputs = [maybe_apply_chat_template(x, self.processing_class) for x in inputs]1097 inputs = [self.tokenize_row(x, self.model.config.is_encoder_decoder, self.processing_class) for x in inputs]1098 inputs = self.data_collator(inputs)1099 1100 # need the prompt_ only1101 inputs = self._prepare_inputs(inputs)1102 context_length = inputs["prompt_input_ids"].shape[1]1103 prompts = {1104 "input_ids": inputs["prompt_input_ids"],1105 "attention_mask": inputs["prompt_attention_mask"],1106 "raw": prompts,1107 }1108 del inputs1109 1110 # Sample completions from both the model and the reference model1111 model_output, ref_output = self._generate_completions(prompts, model)1112 1113 # Process model completions1114 model_data, ref_data = self._process_completions(model_output, ref_output, prompts)1115 1116 # Compute rewards1117 if self.reward_funcs is not None:1118 model_scores, ref_scores = self._compute_rewards(model_data, ref_data, context_length)1119 chosen_mask = model_scores >= ref_scores1120 else:1121 model_scores, ref_scores = None, None1122 chosen_mask = self._compute_judge(model_data, ref_data, context_length)1123 1124 # Compute logprobs1125 model_logprobs_model_data, model_logprobs_ref_data, ref_logprobs_ref_data, ref_logprobs_model_data = (1126 self._compute_logprobs(model, model_data, ref_data, context_length)1127 )1128 1129 # Compute loss1130 loss, dpo_losses, xpo_losses = self._compute_losses(1131 model_logprobs_model_data,1132 model_logprobs_ref_data,1133 ref_logprobs_ref_data,1134 ref_logprobs_model_data,1135 chosen_mask,1136 )1137 1138 # Log everything1139 self._log_statistics(1140 model_data,1141 ref_data,1142 model_logprobs_model_data.detach(),1143 model_logprobs_ref_data.detach(),1144 ref_logprobs_ref_data,1145 ref_logprobs_model_data,1146 chosen_mask,1147 dpo_losses.detach(),1148 xpo_losses.detach(),1149 context_length,1150 model_scores,1151 ref_scores,1152 )1153 1154 if (1155 self.args.torch_empty_cache_steps is not None1156 and self.state.global_step % self.args.torch_empty_cache_steps == 01157 ):1158 empty_cache()1159 1160 kwargs = {}1161 # For LOMO optimizers you need to explicitly use the learning rate1162 if self.args.optim in [OptimizerNames.LOMO, OptimizerNames.ADALOMO]:1163 kwargs["learning_rate"] = self._get_learning_rate()1164 1165 if self.args.n_gpu > 1:1166 loss = loss.mean() # mean() to average on multi-gpu parallel training1167 1168 if self.use_apex:1169 with amp.scale_loss(loss, self.optimizer) as scaled_loss:1170 scaled_loss.backward()1171 else:1172 self.accelerator.backward(loss, **kwargs)1173 1174 return loss.detach() / self.args.gradient_accumulation_steps1175 1176 def create_model_card(1177 self,1178 model_name: Optional[str] = None,1179 dataset_name: Optional[str] = None,1180 tags: Union[str, list[str], None] = None,1181 ):1182 """1183 Creates a draft of a model card using the information available to the `Trainer`.1184 1185 Args:1186 model_name (`str` or `None`, *optional*, defaults to `None`):1187 Name of the model.1188 dataset_name (`str` or `None`, *optional*, defaults to `None`):1189 Name of the dataset used for training.1190 tags (`str`, `list[str]` or `None`, *optional*, defaults to `None`):1191 Tags to be associated with the model card.1192 """1193 if not self.is_world_process_zero():1194 return1195 1196 if hasattr(self.model.config, "_name_or_path") and not os.path.isdir(self.model.config._name_or_path):1197 base_model = self.model.config._name_or_path1198 else:1199 base_model = None1200