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.nash_md_trainer import (Any, BaseImageProcessor, BasePairwiseJudge, Callable, Dataset, EvalPrediction, F, FeatureExtractionMixin, GeometricMixtureWrapper, IterableDataset, NashMDConfig, NashMDTrainer, OnlineDPOTrainer, OptimizerNames, Optional, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, SIMPLE_CHAT_TEMPLATE, TrainerCallback, Union, empty_cache, generate_model_card, get_comet_experiment_url, get_reward, is_conversational, is_peft_available, is_wandb_available, jinja2, maybe_apply_chat_template, nn, os, 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 UnslothNashMDConfig(NashMDConfig):340 """341 342 Configuration class for the [`NashMDTrainer`].343 344 Subclass of [`OnlineDPOConfig`] we can use all its arguments and add the following:345 346 Parameters:347 mixture_coef (`float` or `list[float]`, *optional*, defaults to `0.5`):348 Logit mixture coefficient for the model and reference model. If a list of floats is provided then the349 mixture coefficient is selected for each new epoch and the last coefficient is used for the rest of the350 epochs.351 352 """353 vllm_sampling_params: Optional[Any] = field(354 default = None,355 metadata = {'help': 'vLLM SamplingParams'},356 )357 unsloth_num_chunks : Optional[int] = field(358 default = -1,359 metadata = {'help': 'Chunk size to reduce memory usage. -1 is most efficient.'},360 )361 unsloth_logit_chunk_multiplier : Optional[int] = field(362 default = None,363 metadata = {'help': 'Multiplier for chunked logit computations.'},364 )365 unsloth_grpo_mini_batch : Optional[int] = field(366 default = None,367 metadata = {'help': 'Mini batch size for GRPO hidden state accumulation. Default is None unless user defines it.'},368 )369 max_seq_length : Optional[int] = field(370 default = None,371 metadata = {'help': 'Maximum sequence length to truncate to.'},372 )373 def __init__(374 self,375 output_dir = None,376 overwrite_output_dir = None,377 do_train = False,378 do_eval = False,379 do_predict = False,380 eval_strategy = 'no',381 prediction_loss_only = False,382 per_device_train_batch_size = 4,383 per_device_eval_batch_size = 4,384 per_gpu_train_batch_size = None,385 per_gpu_eval_batch_size = None,386 gradient_accumulation_steps = 2,387 eval_accumulation_steps = 2,388 eval_delay = 0,389 torch_empty_cache_steps = 250,390 learning_rate = 5e-05,391 weight_decay = 0.01,392 adam_beta1 = 0.9,393 adam_beta2 = 0.999,394 adam_epsilon = 1e-08,395 max_grad_norm = 1.0,396 num_train_epochs = 3.0,397 max_steps = -1,398 lr_scheduler_type = 'linear',399 warmup_ratio = 0.1,400 warmup_steps = 0,401 log_level = 'passive',402 log_level_replica = 'warning',403 log_on_each_node = True,404 logging_dir = None,405 logging_strategy = 'steps',406 logging_first_step = False,407 logging_steps = 1,408 logging_nan_inf_filter = False,409 save_strategy = 'steps',410 save_steps = 500,411 save_total_limit = None,412 save_safetensors = True,413 save_on_each_node = False,414 save_only_model = False,415 restore_callback_states_from_checkpoint = False,416 no_cuda = False,417 use_cpu = False,418 use_mps_device = False,419 seed = 3407,420 data_seed = 3407,421 jit_mode_eval = False,422 bf16 = False,423 fp16 = False,424 fp16_opt_level = 'O1',425 half_precision_backend = 'auto',426 bf16_full_eval = False,427 fp16_full_eval = False,428 tf32 = None,429 local_rank = -1,430 ddp_backend = None,431 tpu_num_cores = None,432 tpu_metrics_debug = False,433 debug = '',434 dataloader_drop_last = False,435 eval_steps = None,436 dataloader_num_workers = 0,437 dataloader_prefetch_factor = None,438 past_index = -1,439 run_name = None,440 disable_tqdm = None,441 remove_unused_columns = True,442 label_names = None,443 load_best_model_at_end = False,444 metric_for_best_model = None,445 greater_is_better = None,446 ignore_data_skip = False,447 fsdp = None,448 fsdp_min_num_params = 0,449 fsdp_config = None,450 fsdp_transformer_layer_cls_to_wrap = None,451 accelerator_config = None,452 parallelism_config = None,453 deepspeed = None,454 label_smoothing_factor = 0.0,455 optim = 'adamw_8bit',456 optim_args = None,457 adafactor = False,458 group_by_length = False,459 length_column_name = 'length',460 report_to = 'none',461 project = 'huggingface',462 trackio_space_id = 'trackio',463 ddp_find_unused_parameters = None,464 ddp_bucket_cap_mb = None,465 ddp_broadcast_buffers = None,466 dataloader_pin_memory = True,467 dataloader_persistent_workers = False,468 skip_memory_metrics = True,469 use_legacy_prediction_loop = False,470 push_to_hub = False,471 resume_from_checkpoint = None,472 hub_model_id = None,473 hub_strategy = 'every_save',474 hub_token = None,475 hub_private_repo = None,476 hub_always_push = False,477 hub_revision = None,478 gradient_checkpointing = True,479 gradient_checkpointing_kwargs = None,480 include_inputs_for_metrics = False,481 eval_do_concat_batches = True,482 fp16_backend = 'auto',483 push_to_hub_model_id = None,484 push_to_hub_organization = None,485 push_to_hub_token = None,486 mp_parameters = '',487 auto_find_batch_size = False,488 full_determinism = False,489 torchdynamo = None,490 ray_scope = 'last',491 ddp_timeout = 1800,492 torch_compile = False,493 torch_compile_backend = None,494 torch_compile_mode = None,495 include_tokens_per_second = False,496 include_num_input_tokens_seen = False,497 neftune_noise_alpha = None,498 optim_target_modules = None,499 batch_eval_metrics = False,500 eval_on_start = False,501 use_liger_kernel = False,502 liger_kernel_config = None,503 eval_use_gather_object = False,504 average_tokens_across_devices = True,505 reward_model_path = None,506 judge = None,507 max_new_tokens = 64,508 max_length = 512,509 temperature = 0.9,510 top_p = 1.0,511 top_k = None,512 min_p = None,513 repetition_penalty = 1.0,514 generation_kwargs = {},515 use_transformers_paged = False,516 cache_implementation = None,517 missing_eos_penalty = None,518 loss_type = 'sigmoid',519 disable_dropout = True,520 use_vllm = False,521 vllm_model_impl = 'vllm',522 vllm_guided_decoding_regex = None,523 vllm_gpu_memory_utilization = 0.55,524 vllm_mode = 'colocate',525 vllm_server_base_url = None,526 vllm_server_host = '0.0.0.0',527 vllm_server_port = 8000,528 vllm_server_timeout = 240.0,529 vllm_tensor_parallel_size = 1,530 ds3_gather_for_generation = True,531 model_init_kwargs = None,532 reward_weights = None,533 dataset_num_proc = None,534 gpu_memory_utilization = None,535 vllm_sampling_params = None,536 unsloth_num_chunks = -1,537 unsloth_logit_chunk_multiplier = None,538 unsloth_grpo_mini_batch = None,539 max_seq_length = None,540 **kwargs,541 ):542 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!')543 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!')544 if num_train_epochs is None:545 num_train_epochs = 3.0 # Default to 3 epochs if None, max_steps will override546 if output_dir is None and save_strategy == 'steps' and save_steps == 500:547 output_dir = 'unsloth_training_checkpoints'548 save_strategy = 'no'549 import multiprocessing as _mp550 if dataset_num_proc is None:551 if _mp.get_start_method() != 'fork':552 dataset_num_proc = None553 else:554 import psutil555 dataset_num_proc = min(max((psutil.cpu_count() or 1)+4, 2), 64)556 memory_gb_left = psutil.virtual_memory().available / (1024**3)557 if memory_gb_left <= 2: dataset_num_proc = 1558 else: dataset_num_proc = min(dataset_num_proc, int(memory_gb_left))559 if temperature <= 0:560 raise ValueError('Unsloth: Please set a positive non-zero temperature since your results will be wrong.')561 elif temperature >= 10:562 raise ValueError('Unsloth: Please set a positive non-zero temperature less than 10, since sampling will be quite erratic.')563 564 565 super().__init__(566 output_dir = output_dir,567 overwrite_output_dir = overwrite_output_dir,568 do_train = do_train,569 do_eval = do_eval,570 do_predict = do_predict,571 eval_strategy = eval_strategy,572 prediction_loss_only = prediction_loss_only,573 per_device_train_batch_size = per_device_train_batch_size,574 per_device_eval_batch_size = per_device_eval_batch_size,575 per_gpu_train_batch_size = per_gpu_train_batch_size,576 per_gpu_eval_batch_size = per_gpu_eval_batch_size,577 gradient_accumulation_steps = gradient_accumulation_steps,578 eval_accumulation_steps = eval_accumulation_steps,579 eval_delay = eval_delay,580 torch_empty_cache_steps = torch_empty_cache_steps,581 learning_rate = learning_rate,582 weight_decay = weight_decay,583 adam_beta1 = adam_beta1,584 adam_beta2 = adam_beta2,585 adam_epsilon = adam_epsilon,586 max_grad_norm = max_grad_norm,587 num_train_epochs = num_train_epochs,588 max_steps = max_steps,589 lr_scheduler_type = lr_scheduler_type,590 warmup_ratio = warmup_ratio,591 warmup_steps = warmup_steps,592 log_level = log_level,593 log_level_replica = log_level_replica,594 log_on_each_node = log_on_each_node,595 logging_dir = logging_dir,596 logging_strategy = logging_strategy,597 logging_first_step = logging_first_step,598 logging_steps = logging_steps,599 logging_nan_inf_filter = logging_nan_inf_filter,600 save_strategy = save_strategy,601 save_steps = save_steps,602 save_total_limit = save_total_limit,603 save_safetensors = save_safetensors,604 save_on_each_node = save_on_each_node,605 save_only_model = save_only_model,606 restore_callback_states_from_checkpoint = restore_callback_states_from_checkpoint,607 no_cuda = no_cuda,608 use_cpu = use_cpu,609 use_mps_device = use_mps_device,610 seed = seed,611 data_seed = data_seed,612 jit_mode_eval = jit_mode_eval,613 bf16 = bf16,614 fp16 = fp16,615 fp16_opt_level = fp16_opt_level,616 half_precision_backend = half_precision_backend,617 bf16_full_eval = bf16_full_eval,618 fp16_full_eval = fp16_full_eval,619 tf32 = tf32,620 local_rank = local_rank,621 ddp_backend = ddp_backend,622 tpu_num_cores = tpu_num_cores,623 tpu_metrics_debug = tpu_metrics_debug,624 debug = debug,625 dataloader_drop_last = dataloader_drop_last,626 eval_steps = eval_steps,627 dataloader_num_workers = dataloader_num_workers,628 dataloader_prefetch_factor = dataloader_prefetch_factor,629 past_index = past_index,630 run_name = run_name,631 disable_tqdm = disable_tqdm,632 remove_unused_columns = remove_unused_columns,633 label_names = label_names,634 load_best_model_at_end = load_best_model_at_end,635 metric_for_best_model = metric_for_best_model,636 greater_is_better = greater_is_better,637 ignore_data_skip = ignore_data_skip,638 fsdp = fsdp,639 fsdp_min_num_params = fsdp_min_num_params,640 fsdp_config = fsdp_config,641 fsdp_transformer_layer_cls_to_wrap = fsdp_transformer_layer_cls_to_wrap,642 accelerator_config = accelerator_config,643 parallelism_config = parallelism_config,644 deepspeed = deepspeed,645 label_smoothing_factor = label_smoothing_factor,646 optim = optim,647 optim_args = optim_args,648 adafactor = adafactor,649 group_by_length = group_by_length,650 length_column_name = length_column_name,651 report_to = report_to,652 project = project,653 trackio_space_id = trackio_space_id,654 ddp_find_unused_parameters = ddp_find_unused_parameters,655 ddp_bucket_cap_mb = ddp_bucket_cap_mb,656 ddp_broadcast_buffers = ddp_broadcast_buffers,657 dataloader_pin_memory = dataloader_pin_memory,658 dataloader_persistent_workers = dataloader_persistent_workers,659 skip_memory_metrics = skip_memory_metrics,660 use_legacy_prediction_loop = use_legacy_prediction_loop,661 push_to_hub = push_to_hub,662 resume_from_checkpoint = resume_from_checkpoint,663 hub_model_id = hub_model_id,664 hub_strategy = hub_strategy,665 hub_token = hub_token,666 hub_private_repo = hub_private_repo,667 hub_always_push = hub_always_push,668 hub_revision = hub_revision,669 gradient_checkpointing = gradient_checkpointing,670 gradient_checkpointing_kwargs = gradient_checkpointing_kwargs,671 include_inputs_for_metrics = include_inputs_for_metrics,672 eval_do_concat_batches = eval_do_concat_batches,673 fp16_backend = fp16_backend,674 push_to_hub_model_id = push_to_hub_model_id,675 push_to_hub_organization = push_to_hub_organization,676 push_to_hub_token = push_to_hub_token,677 mp_parameters = mp_parameters,678 auto_find_batch_size = auto_find_batch_size,679 full_determinism = full_determinism,680 torchdynamo = torchdynamo,681 ray_scope = ray_scope,682 ddp_timeout = ddp_timeout,683 torch_compile = torch_compile,684 torch_compile_backend = torch_compile_backend,685 torch_compile_mode = torch_compile_mode,686 include_tokens_per_second = include_tokens_per_second,687 include_num_input_tokens_seen = include_num_input_tokens_seen,688 neftune_noise_alpha = neftune_noise_alpha,689 optim_target_modules = optim_target_modules,690 batch_eval_metrics = batch_eval_metrics,691 eval_on_start = eval_on_start,692 use_liger_kernel = use_liger_kernel,693 liger_kernel_config = liger_kernel_config,694 eval_use_gather_object = eval_use_gather_object,695 average_tokens_across_devices = average_tokens_across_devices,696 reward_model_path = reward_model_path,697 judge = judge,698 max_new_tokens = max_new_tokens,699 max_length = max_length,700 temperature = temperature,701 top_p = top_p,702 top_k = top_k,703 min_p = min_p,704 repetition_penalty = repetition_penalty,705 generation_kwargs = generation_kwargs,706 use_transformers_paged = use_transformers_paged,707 cache_implementation = cache_implementation,708 missing_eos_penalty = missing_eos_penalty,709 loss_type = loss_type,710 disable_dropout = disable_dropout,711 use_vllm = use_vllm,712 vllm_model_impl = vllm_model_impl,713 vllm_guided_decoding_regex = vllm_guided_decoding_regex,714 vllm_gpu_memory_utilization = vllm_gpu_memory_utilization,715 vllm_mode = vllm_mode,716 vllm_server_base_url = vllm_server_base_url,717 vllm_server_host = vllm_server_host,718 vllm_server_port = vllm_server_port,719 vllm_server_timeout = vllm_server_timeout,720 vllm_tensor_parallel_size = vllm_tensor_parallel_size,721 ds3_gather_for_generation = ds3_gather_for_generation,722 model_init_kwargs = model_init_kwargs,723 reward_weights = reward_weights,724 dataset_num_proc = dataset_num_proc,725 gpu_memory_utilization = gpu_memory_utilization,**kwargs)726 self.vllm_sampling_params = vllm_sampling_params727 self.unsloth_num_chunks = unsloth_num_chunks728 if unsloth_grpo_mini_batch is not None:729 if self.generation_batch_size >= unsloth_grpo_mini_batch:730 self.unsloth_grpo_mini_batch = unsloth_grpo_mini_batch731 else:732 raise ValueError(733 f"Unsloth GRPO mini batch size needs to be less than or equal to the effective generation batch size, "734 f"which is self.per_device_train_batch_size * gradient_accumulation_steps."735 )736 self.unsloth_logit_chunk_multiplier = unsloth_logit_chunk_multiplier737 self.max_seq_length = max_seq_length738 739pass740 741class _UnslothNashMDTrainer(OnlineDPOTrainer):742 r""""""743 744 _tag_names = ["trl", "nash-md"]745 746 def __init__(747 self,748 model: Union[PreTrainedModel, nn.Module] = None,749 ref_model: Union[PreTrainedModel, nn.Module] = None,750 reward_funcs: Union[PreTrainedModel, nn.Module, None] = None,751 judge: Optional[BasePairwiseJudge] = None,752 args: Optional[NashMDConfig] = None,753 data_collator: Optional[Callable] = None,754 train_dataset: Optional[Union[Dataset, IterableDataset]] = None,755 eval_dataset: Optional[Union[Dataset, dict[str, Dataset]]] = None,756 processing_class: Optional[757 Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin]758 ] = 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 reward_funcs=reward_funcs,771 judge=judge,772 args=args,773 data_collator=data_collator,774 train_dataset=train_dataset,775 eval_dataset=eval_dataset,776 processing_class=processing_class,777 reward_processing_classes=processing_class,778 peft_config=peft_config,779 compute_metrics=compute_metrics,780 callbacks=callbacks,781 optimizers=optimizers,782 preprocess_logits_for_metrics=preprocess_logits_for_metrics,783 reward_model=reward_model,784 )785 786 self._mixture_coef = self.args.mixture_coef787 788 # Overwrite the stats dictionary to include NashMD specific statistics789 self.stats = {790 # Remove "non_score_reward", "rlhf_reward", "scores_margin"791 # Add "mixture_coef"792 "loss/kl": [],793 "objective/entropy": [],794 "loss/score": [],795 "rewards/probabilities": [],796 "rewards/accuracies": [],797 "rewards/margins": [],798 "logps/chosen": [],799 "logps/rejected": [],800 "val/model_contain_eos_token": [],801 "val/ref_contain_eos_token": [],802 "beta": [],803 "mixture_coef": [],804 }805 if self.reward_funcs is not None:806 if len(self.reward_funcs) != 1:807 raise ValueError("NashMDTrainer only supports one reward function/model.")808 self.reward_funcs = self.reward_funcs[0]809 self.stats["rewards/chosen"] = []810 self.stats["rewards/rejected"] = []811 812 @property813 def mixture_coef(self):814 if isinstance(self._mixture_coef, list):815 epoch = self.state.epoch816 return self._mixture_coef[epoch] if epoch < len(self._mixture_coef) else self._mixture_coef[-1]817 else:818 return self._mixture_coef819 820 def _generate_completions(self, model, prompts):821 # Generate completions from the policy model.822 with unwrap_model_for_generation(model, self.accelerator) as unwrapped_policy_for_gen_ctx:823 model_output = unwrapped_policy_for_gen_ctx.generate(824 input_ids=prompts["input_ids"],825 attention_mask=prompts["attention_mask"],826 generation_config=self.generation_config,827 )828 829 # Get the DDP/FSDP unwrapped version of the main model.830 # This will be the policy model for GeometricMixtureWrapper (PEFT adapters active if PEFT is used).831 policy_model_for_gmw = self.accelerator.unwrap_model(model)832 833 # Determine the correct reference model for GeometricMixtureWrapper.834 # This also needs to be DDP/FSDP unwrapped.835 ref_model_for_gmw: torch.nn.Module836 if self.ref_model is None:837 # No explicit ref_model is provided.838 # Use the base of the main `model` if it's a PEFT model.839 # policy_model_for_gmw is already DDP-unwrapped.840 if is_peft_available() and isinstance(policy_model_for_gmw, PeftModel):841 ref_model_for_gmw = policy_model_for_gmw.get_base_model()842 else:843 # Not a PEFT model (or PEFT not available), or already a base model.844 # Use the DDP-unwrapped policy model itself as the reference.845 ref_model_for_gmw = policy_model_for_gmw846 else:847 # An explicit ref_model is provided. Unwrap it for DDP/FSDP.848 ref_model_for_gmw = self.accelerator.unwrap_model(self.ref_model)849 850 # Both models given to GeometricMixtureWrapper (policy_model_for_gmw and ref_model_for_gmw) are DDP-unwrapped.851 with torch.no_grad(): # Ensure no_grad context for mixture model generation852 mixture_model = GeometricMixtureWrapper(853 model=policy_model_for_gmw,854 ref_model=ref_model_for_gmw,855 generation_config=self.generation_config,856 mixture_coef=self.mixture_coef,857 device=self.accelerator.device,858 )859 860 mixture_output = mixture_model.generate(861 input_ids=prompts["input_ids"],862 attention_mask=prompts["attention_mask"],863 generation_config=self.generation_config,864 )865 866 return model_output, mixture_output867 868 def _process_completions(self, model_output, mixture_output, prompts):869 context_length = prompts["input_ids"].shape[1]870 871 # Process model completions872 model_completion_ids = model_output[:, context_length:]873 model_completion_ids, model_completion_mask = truncate_right(874 model_completion_ids, self.processing_class.eos_token_id, self.processing_class.pad_token_id875 )876 model_data = {877 "input_ids": torch.cat((prompts["input_ids"], model_completion_ids), dim=1),878 "attention_mask": torch.cat((prompts["attention_mask"], model_completion_mask), dim=1),879 "raw": prompts["raw"],880 }881 882 # Process reference model completions883 mixture_completion_ids = mixture_output[:, context_length:]884 mixture_completion_ids, mixture_completion_mask = truncate_right(885 mixture_completion_ids, self.processing_class.eos_token_id, self.processing_class.pad_token_id886 )887 mixture_data = {888 "input_ids": torch.cat((prompts["input_ids"], mixture_completion_ids), dim=1),889 "attention_mask": torch.cat((prompts["attention_mask"], mixture_completion_mask), dim=1),890 "raw": prompts["raw"],891 }892 893 return model_data, mixture_data894 895 def _compute_rewards(self, model_data, mixture_data, context_length):896 with torch.no_grad():897 _, model_scores, _ = get_reward(898 self.reward_funcs, model_data["input_ids"], self.processing_class.pad_token_id, context_length899 )900 _, mixture_scores, _ = get_reward(901 self.reward_funcs, mixture_data["input_ids"], self.processing_class.pad_token_id, context_length902 )903 904 # Apply EOS penalty if needed905 if self.args.missing_eos_penalty is not None:906 model_contain_eos = torch.any(model_data["input_ids"] == self.processing_class.eos_token_id, dim=-1)907 mixture_contain_eos = torch.any(mixture_data["input_ids"] == self.processing_class.eos_token_id, dim=-1)908 model_scores[~model_contain_eos] -= self.args.missing_eos_penalty909 mixture_scores[~mixture_contain_eos] -= self.args.missing_eos_penalty910 911 return model_scores, mixture_scores912 913 def _compute_judge(self, model_data, mixture_data, context_length):914 prompts = model_data["raw"]915 model_data_completions = self.processing_class.batch_decode(916 model_data["input_ids"][:, context_length:], skip_special_tokens=True917 )918 model_data_completions = [completion.strip() for completion in model_data_completions]919 920 mixture_data_completions = self.processing_class.batch_decode(921 mixture_data["input_ids"][:, context_length:], skip_special_tokens=True922 )923 mixture_data_completions = [completion.strip() for completion in mixture_data_completions]924 if is_conversational({"prompt": prompts[0]}):925 model_data_completions = [926 [{"role": "assistant", "content": completion}] for completion in model_data_completions927 ]928 environment = jinja2.Environment()929 template = environment.from_string(SIMPLE_CHAT_TEMPLATE)930 prompts = [template.render(messages=message) for message in prompts]931 model_data_completions = [template.render(messages=completion) for completion in model_data_completions]932 933 mixture_data_completions = [934 [{"role": "assistant", "content": completion}] for completion in mixture_data_completions935 ]936 mixture_data_completions = [937 template.render(messages=completion) for completion in mixture_data_completions938 ]939 940 probability = self.judge.judge(941 prompts,942 list(zip(model_data_completions, mixture_data_completions)),943 return_scores=True,944 )945 return torch.tensor(probability, device=model_data["input_ids"].device)946 947 def _compute_logprobs(self, model, model_data, context_length):948 def compute_logprobs_for_data(m, data):949 output = m(data["input_ids"], attention_mask=data["attention_mask"])950 logits = output.logits[:, context_length - 1 : -1]951 token_logprobs = selective_log_softmax(logits, data["input_ids"][:, context_length:])952 return token_logprobs953 954 # Compute logprobs for model completions under the model955 model_logprobs_model_data = compute_logprobs_for_data(model, model_data)956 957 # Compute logprobs of model completions under the reference model958 with torch.no_grad():959 if self.ref_model is None:960 with model.disable_adapter():961 ref_logprobs_model_data = compute_logprobs_for_data(model, model_data)962 else:963 ref_logprobs_model_data = compute_logprobs_for_data(self.ref_model, model_data)964 965 # Mask padding tokens966 model_padding_mask = model_data["attention_mask"][:, context_length:] == 0967 model_logprobs_model_data = model_logprobs_model_data.masked_fill(model_padding_mask, 0.0)968 ref_logprobs_model_data = ref_logprobs_model_data.masked_fill(model_padding_mask, 0.0)969 970 return (model_logprobs_model_data, ref_logprobs_model_data)971 972 def _compute_losses(973 self,974 model_logprobs_model_data,975 ref_logprobs_model_data,976 probability,977 ):978 # reinforce score where 0.5 is a control variate979 score = (probability - 0.5) * model_logprobs_model_data.sum(1)980 981 # kl divergence via reinforce982 with torch.no_grad():983 log_ratio = model_logprobs_model_data - ref_logprobs_model_data984 kl_div_log = log_ratio.sum(1)985 kl_div_loss = (log_ratio * model_logprobs_model_data).sum(1)986 987 # final loss988 loss = self.beta * kl_div_loss - score989 990 return loss.mean(), score, kl_div_log991 992 def _log_statistics(993 self,994 model_data,995 mixture_data,996 model_logprobs_model_data,997 ref_logprobs_model_data,998 probability,999 score,1000 kl_div,1001 context_length,1002 model_scores=None,1003 mixture_scores=None,1004 ):1005 # Helper function to gather and compute mean1006 def gather_mean(tensor):1007 return self.accelerator.gather_for_metrics(tensor).mean().item()1008 1009 # Log score1010 self.stats["loss/score"].append(gather_mean(score))1011 # Log KL divergence1012 self.stats["loss/kl"].append(gather_mean(kl_div))1013 1014 # Log logprobs1015 model_logprobs_model_data_sum = model_logprobs_model_data.sum(1)1016 ref_logprobs_model_data_sum = ref_logprobs_model_data.sum(1)1017 1018 self.stats["logps/chosen"].append(gather_mean(model_logprobs_model_data_sum))1019 self.stats["logps/rejected"].append(gather_mean(ref_logprobs_model_data_sum))1020 1021 # Log rewards1022 if self.reward_funcs is not None:1023 self.stats["rewards/chosen"].append(gather_mean(model_scores))1024 self.stats["rewards/rejected"].append(gather_mean(mixture_scores))1025 1026 # Log probabilities1027 self.stats["rewards/probabilities"].append(gather_mean(probability))1028 1029 # Calculate entropy for model data1030 entropy_model_data = -model_logprobs_model_data.sum(1)1031 self.stats["objective/entropy"].append(gather_mean(entropy_model_data))1032 1033 # Calculate margins1034 margin = model_logprobs_model_data_sum - ref_logprobs_model_data_sum1035 self.stats["rewards/margins"].append(gather_mean(margin))1036 1037 # Calculate accuracy1038 accuracy = (margin > 0).float()1039 self.stats["rewards/accuracies"].append(gather_mean(accuracy))1040 1041 # Log EOS token statistics1042 model_eos = (model_data["input_ids"][:, context_length:] == self.processing_class.eos_token_id).any(dim=1)1043 mixture_eos = (mixture_data["input_ids"][:, context_length:] == self.processing_class.eos_token_id).any(dim=1)1044 self.stats["val/model_contain_eos_token"].append(gather_mean(model_eos.float()))1045 self.stats["val/ref_contain_eos_token"].append(gather_mean(mixture_eos.float()))1046 1047 # Log beta and mixture coef1048 self.stats["beta"].append(self.beta)1049 self.stats["mixture_coef"].append(self.mixture_coef)1050 1051 def training_step(1052 self, model: nn.Module, inputs: dict[str, Union[torch.Tensor, Any]], num_items_in_batch: Optional[int] = None1053 ) -> torch.Tensor:1054 model.train()1055 1056 # Apply chat template and tokenize the input1057 batch_size = len(next(iter(inputs.values())))1058 prompts = inputs["prompt"]1059 inputs = [{k: v[i] for k, v in inputs.items()} for i in range(batch_size)]1060 inputs = [maybe_apply_chat_template(x, self.processing_class) for x in inputs]1061 inputs = [self.tokenize_row(x, self.model.config.is_encoder_decoder, self.processing_class) for x in inputs]1062 inputs = self.data_collator(inputs)1063 1064 # need the prompt_ only1065 inputs = self._prepare_inputs(inputs)1066 context_length = inputs["prompt_input_ids"].shape[1]1067 prompts = {1068 "input_ids": inputs["prompt_input_ids"],1069 "attention_mask": inputs["prompt_attention_mask"],1070 "raw": prompts,1071 }1072 del inputs1073 1074 # Sample completions from both the model and the reference model1075 model_output, mixture_output = self._generate_completions(model, prompts)1076 1077 # Process model completions1078 model_data, mixture_data = self._process_completions(model_output, mixture_output, prompts)1079 1080 # Compute rewards1081 if self.reward_funcs is not None:1082 model_scores, mixture_scores = self._compute_rewards(model_data, mixture_data, context_length)1083 # probability of the model data vs the mixture data1084 probability = F.sigmoid(model_scores - mixture_scores)1085 else:1086 model_scores, mixture_scores = None, None1087 probability = self._compute_judge(model_data, mixture_data, context_length)1088 1089 # Compute logprobs1090 model_logprobs_model_data, ref_logprobs_model_data = self._compute_logprobs(model, model_data, context_length)1091 1092 # Compute loss1093 loss, score, kl_div = self._compute_losses(model_logprobs_model_data, ref_logprobs_model_data, probability)1094 1095 # Log everything1096 self._log_statistics(1097 model_data,1098 mixture_data,1099 model_logprobs_model_data.detach(),1100 ref_logprobs_model_data,1101 probability,1102 score.detach(),1103 kl_div.detach(),1104 context_length,1105 model_scores,1106 mixture_scores,1107 )1108 1109 if (1110 self.args.torch_empty_cache_steps is not None1111 and self.state.global_step % self.args.torch_empty_cache_steps == 01112 ):1113 empty_cache()1114 1115 kwargs = {}1116 # For LOMO optimizers you need to explicitly use the learning rate1117 if self.args.optim in [OptimizerNames.LOMO, OptimizerNames.ADALOMO]:1118 kwargs["learning_rate"] = self._get_learning_rate()1119 1120 if self.args.n_gpu > 1:1121 loss = loss.mean() # mean() to average on multi-gpu parallel training1122 1123 if self.use_apex:1124 with amp.scale_loss(loss, self.optimizer) as scaled_loss:1125 scaled_loss.backward()1126 else:1127 self.accelerator.backward(loss, **kwargs)1128 1129 return loss.detach() / self.args.gradient_accumulation_steps1130 1131 def create_model_card(1132 self,1133 model_name: Optional[str] = None,1134 dataset_name: Optional[str] = None,1135 tags: Union[str, list[str], None] = None,1136 ):1137 """1138 Creates a draft of a model card using the information available to the `Trainer`.1139 1140 Args:1141 model_name (`str` or `None`, *optional*, defaults to `None`):1142 Name of the model.1143 dataset_name (`str` or `None`, *optional*, defaults to `None`):1144 Name of the dataset used for training.1145 tags (`str`, `list[str]` or `None`, *optional*, defaults to `None`):1146 Tags to be associated with the model card.1147 """1148 if not self.is_world_process_zero():1149 return1150 1151 if hasattr(self.model.config, "_name_or_path") and not os.path.isdir(self.model.config._name_or_path):1152 base_model = self.model.config._name_or_path1153 else:1154 base_model = None1155 1156 # normalize `tags` to a mutable set1157 if tags is None:1158 tags = set()1159 elif isinstance(tags, str):1160 tags = {tags}1161 else:1162 tags = set(tags)1163 1164 if hasattr(self.model.config, "unsloth_version"):1165 tags.add("unsloth")1166 1167 if "JOB_ID" in os.environ:1168 tags.add("hf_jobs")1169 1170 tags.update(self._tag_names)1171 1172 # docstyle-ignore1173 citation = textwrap.dedent("""\1174 @inproceedings{munos2024nash,1175 title = {{Nash Learning from Human Feedback}},1176 author = {R{\'{e}}mi Munos and Michal Valko and Daniele Calandriello and Mohammad Gheshlaghi Azar and Mark Rowland and Zhaohan Daniel Guo and Yunhao Tang and Matthieu Geist and Thomas Mesnard and C{\\^{o}}me Fiegel and Andrea Michi and Marco Selvi and Sertan Girgin and Nikola Momchev and Olivier Bachem and Daniel J. Mankowitz and Doina Precup and Bilal Piot},1177 year = 2024,1178 booktitle = {Forty-first International Conference on Machine Learning, {ICML} 2024, Vienna, Austria, July 21-27, 2024},1179 publisher = {OpenReview.net},1180 url = {https://openreview.net/forum?id=Y5AmNYiyCQ}1181 }""")1182 1183 model_card = generate_model_card(1184 base_model=base_model,1185 model_name=model_name,1186 hub_model_id=self.hub_model_id,1187 dataset_name=dataset_name,1188 tags=tags,1189 wandb_url=wandb.run.url if is_wandb_available() and wandb.run is not None else None,1190 comet_url=get_comet_experiment_url(),1191 trainer_name="Nash-MD",1192 trainer_citation=citation,1193 paper_title="Nash Learning from Human Feedback",1194 paper_id="2312.00886",1195 )1196 1197 model_card.save(os.path.join(self.args.output_dir, "README.md"))1198class UnslothNashMDTrainer(_UnslothNashMDTrainer):1199 """1200