| """ |
| 2025.10.1 |
| 2025.10.1 |
| 4.56.2 |
| 0.22.2 |
| __UNSLOTH_VERSIONING__ |
| """ |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| from torch import Tensor |
| import torch |
| import torch.nn as nn |
| from torch.nn import functional as F |
| from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable |
| from trl.trainer.ppo_trainer import (Accelerator, BaseImageProcessor, CallbackHandler, DEFAULT_CALLBACKS, DEFAULT_PROGRESS_CALLBACK, DataCollatorWithPadding, DataLoader, Dataset, ExportableState, FeatureExtractionMixin, GenerationConfig, INVALID_LOGPROB, OnlineTrainerState, Optional, PPOConfig, PPOTrainer, Path, PeftConfig, PeftModel, PolicyAndValueWrapper, PreTrainedTokenizerBase, PrinterCallback, ProcessorMixin, Trainer, TrainerCallback, TrainerControl, Union, batch_generation, broadcast, contextmanager, create_reference_model, defaultdict, disable_dropout_in_model, empty_cache, exact_div, first_true_indices, forward, gather_object, gc, generate_model_card, get_comet_experiment_url, get_peft_model, get_reporting_integration_callbacks, get_reward, is_peft_available, is_rich_available, is_wandb_available, log_table_to_comet_experiment, masked_mean, masked_whiten, math, nn, np, nullcontext, os, pd, peft_module_casting_to_bf16, prepare_deepspeed, print_rich_table, selective_log_softmax, textwrap, time, torch, truncate_response, unwrap_model_for_generation, wandb, Optional, PeftModel, Trainer, is_peft_available, os, torch) |
|
|
|
|
| import os |
| from typing import * |
| from dataclasses import dataclass, field |
| from packaging.version import Version |
| import torch |
| import numpy as np |
| from contextlib import nullcontext |
| from torch.nn import functional as F |
| from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling as TransformersDataCollatorForLanguageModeling |
| from transformers.training_args import ParallelMode |
|
|
| |
| import functools |
| from types import MethodType |
| def prepare_for_training_mode(f): |
| @functools.wraps(f) |
| def wrapper(self, *args, **kwargs): |
| |
| if hasattr(self, 'model') and hasattr(self.model, "for_training"): |
| self.model.for_training() |
| output = f(self, *args, **kwargs) |
| |
| if hasattr(self, 'model') and hasattr(self.model, "for_inference"): |
| self.model.for_inference() |
| return output |
| return wrapper |
| pass |
|
|
| torch_compile_options = { |
| "epilogue_fusion" : True, |
| "max_autotune" : False, |
| "shape_padding" : True, |
| "trace.enabled" : False, |
| "triton.cudagraphs" : False, |
| } |
|
|
| @torch.compile(dynamic = True, fullgraph = True, options = torch_compile_options,) |
| def chunked_selective_log_softmax(logits, index): |
| |
| chunked_logits = torch.chunk(logits.reshape(-1, logits.shape[-1]), chunks = 4, dim = 0) |
| chunked_index = torch.chunk(index.reshape(-1), chunks = 4, dim = 0) |
| all_per_token_logps = [] |
| |
| for chunk_logits, chunk_index in zip(chunked_logits, chunked_index): |
| chunk_logits = chunk_logits.to(torch.float32) |
| selected_logits = torch.gather(chunk_logits, dim = -1, index = chunk_index.unsqueeze(-1)).squeeze(-1) |
| logsumexp_values = torch.logsumexp(chunk_logits, dim = -1) |
| per_token_logps = selected_logits - logsumexp_values |
| all_per_token_logps.append(per_token_logps) |
| pass |
| all_per_token_logps = torch.concat(all_per_token_logps) |
| all_per_token_logps = all_per_token_logps.reshape((logits.shape[0], logits.shape[1])) |
| return all_per_token_logps |
|
|
| def calculate_pad_tokens_in_prompt( |
| input_ids: torch.Tensor, |
| logits_to_keep: int, |
| pad_token_id: int |
| ) -> torch.Tensor: |
| """ |
| Given prompt tensor, it returns all the left padded tokens in that sequence. so [pad, pad, pad, cat] = 3 tokens |
| """ |
| if logits_to_keep >= input_ids.shape[1]: |
| raise ValueError("logits_to_keep must be smaller than the sequence length.") |
|
|
| prompt_section = input_ids[:, :-logits_to_keep] |
|
|
| padding_mask = (prompt_section == pad_token_id) |
|
|
| pad_token_counts = padding_mask.sum(dim=1) |
|
|
| return pad_token_counts |
|
|
| def create_completion_attention_mask( |
| completion_input_ids: torch.Tensor, |
| left_pad_tokens_per_prompt: torch.Tensor, |
| max_left_pad: int, |
| pad_token_id: int |
| ) -> torch.Tensor: |
| """ |
| Given that we have a sequence, [p,p,p,c,c,c,pad,pad,pad] |
| |
| Where p are extra prompt tokens we got from slicing the torch tensor, c is completion tokens |
| and pad are pad tokens, this function would make a completion mask that would 0 out the pad |
| and p tokens. so in this example [0,0,0,1,1,1,0,0,0] |
| """ |
| batch_size, completion_len = completion_input_ids.shape |
| device = completion_input_ids.device |
|
|
| num_tokens_to_mask = max_left_pad - left_pad_tokens_per_prompt |
|
|
| indices = torch.arange(completion_len, device=device).unsqueeze(0) |
| shift_mask = indices >= num_tokens_to_mask.unsqueeze(1) |
|
|
| non_padding_mask = (completion_input_ids != pad_token_id) |
|
|
| final_mask = shift_mask & non_padding_mask |
|
|
| return final_mask |
|
|
| def left_pack_padding(tensor: torch.Tensor, pad_id: int) -> torch.Tensor: |
| """ |
| Moves all padding tokens in each sequence of a batch to the right. |
| """ |
| mask = (tensor != pad_id) |
| |
| sorted_indices = torch.argsort(mask, dim=1, descending=True, stable=True) |
| packed_tensor = torch.gather(tensor, 1, sorted_indices) |
| return packed_tensor |
| @dataclass |
| class UnslothPPOConfig(PPOConfig): |
| """ |
| |
| Configuration class for the [`PPOTrainer`]. |
| |
| This class includes only the parameters that are specific to PPO training. For a full list of training arguments, |
| please refer to the [`~transformers.TrainingArguments`] and [`OnPolicyConfig`] documentation. Note that default |
| values in this class may differ from those in [`~transformers.TrainingArguments`]. |
| |
| Using [`~transformers.HfArgumentParser`] we can turn this class into |
| [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the |
| command line. |
| |
| Parameters: |
| exp_name (`str`, *optional*, defaults to `os.path.basename(__file__)[:-3]`): |
| Name of this experiment. |
| reward_model_path (`str`, *optional*, defaults to `"EleutherAI/pythia-160m"`): |
| Path to the reward model. |
| model_adapter_name (`str` or `None`, *optional*, defaults to `None`): |
| Name of the train target PEFT adapter, when using LoRA with multiple adapters. |
| ref_adapter_name (`str` or `None`, *optional*, defaults to `None`): |
| Name of the reference PEFT adapter, when using LoRA with multiple adapters. |
| num_ppo_epochs (`int`, *optional*, defaults to `4`): |
| Number of epochs to train. |
| whiten_rewards (`bool`, *optional*, defaults to `False`): |
| Whether to whiten the rewards. |
| kl_coef (`float`, *optional*, defaults to `0.05`): |
| KL coefficient. |
| kl_estimator (`Literal["k1", "k3"]`, *optional*, defaults to `"k1"`): |
| Which estimator for KL-Divergence to use from [Approximating KL |
| Divergence](http://joschu.net/blog/kl-approx.html). Defaults to "k1", a straightforward, unbiased |
| estimator. Can be set to "k3", an unbiased estimator with lower variance which "appears to be a strictly |
| better estimator". Cannot be set to "k2", as it is used for logging purposes. |
| cliprange (`float`, *optional*, defaults to `0.2`): |
| Clip range. |
| vf_coef (`float`, *optional*, defaults to `0.1`): |
| Value function coefficient. |
| cliprange_value (`float`, *optional*, defaults to `0.2`): |
| Clip range for the value function. |
| gamma (`float`, *optional*, defaults to `1.0`): |
| Discount factor. |
| lam (`float`, *optional*, defaults to `0.95`): |
| Lambda value for GAE. |
| ds3_gather_for_generation (`bool`, *optional*, defaults to `True`): |
| This setting applies to DeepSpeed ZeRO-3. If enabled, the policy model weights are gathered for generation, |
| improving generation speed. However, disabling this option allows training models that exceed the VRAM |
| capacity of a single GPU, albeit at the cost of slower generation. |
| |
| """ |
| vllm_sampling_params: Optional[Any] = field( |
| default = None, |
| metadata = {'help': 'vLLM SamplingParams'}, |
| ) |
| unsloth_num_chunks : Optional[int] = field( |
| default = -1, |
| metadata = {'help': 'Chunk size to reduce memory usage. -1 is most efficient.'}, |
| ) |
| |
| def __init__( |
| self, |
| output_dir = None, |
| overwrite_output_dir = None, |
| do_train = False, |
| do_eval = False, |
| do_predict = False, |
| eval_strategy = 'no', |
| prediction_loss_only = False, |
| per_device_train_batch_size = 4, |
| per_device_eval_batch_size = 4, |
| per_gpu_train_batch_size = None, |
| per_gpu_eval_batch_size = None, |
| gradient_accumulation_steps = 2, |
| eval_accumulation_steps = 2, |
| eval_delay = 0, |
| torch_empty_cache_steps = 250, |
| learning_rate = 5e-05, |
| weight_decay = 0.01, |
| adam_beta1 = 0.9, |
| adam_beta2 = 0.999, |
| adam_epsilon = 1e-08, |
| max_grad_norm = 1.0, |
| num_train_epochs = 3.0, |
| max_steps = -1, |
| lr_scheduler_type = 'linear', |
| warmup_ratio = 0.1, |
| warmup_steps = 0, |
| log_level = 'passive', |
| log_level_replica = 'warning', |
| log_on_each_node = True, |
| logging_dir = None, |
| logging_strategy = 'steps', |
| logging_first_step = False, |
| logging_steps = 1, |
| logging_nan_inf_filter = False, |
| save_strategy = 'steps', |
| save_steps = 500, |
| save_total_limit = None, |
| save_safetensors = True, |
| save_on_each_node = False, |
| save_only_model = False, |
| restore_callback_states_from_checkpoint = False, |
| no_cuda = False, |
| use_cpu = False, |
| use_mps_device = False, |
| seed = 3407, |
| data_seed = 3407, |
| jit_mode_eval = False, |
| use_ipex = False, |
| bf16 = False, |
| fp16 = False, |
| fp16_opt_level = 'O1', |
| half_precision_backend = 'auto', |
| bf16_full_eval = False, |
| fp16_full_eval = False, |
| tf32 = None, |
| local_rank = -1, |
| ddp_backend = None, |
| tpu_num_cores = None, |
| tpu_metrics_debug = False, |
| debug = '', |
| dataloader_drop_last = False, |
| eval_steps = None, |
| dataloader_num_workers = 0, |
| dataloader_prefetch_factor = None, |
| past_index = -1, |
| run_name = None, |
| disable_tqdm = None, |
| remove_unused_columns = True, |
| label_names = None, |
| load_best_model_at_end = False, |
| metric_for_best_model = None, |
| greater_is_better = None, |
| ignore_data_skip = False, |
| fsdp = '', |
| fsdp_min_num_params = 0, |
| fsdp_config = None, |
| fsdp_transformer_layer_cls_to_wrap = None, |
| accelerator_config = None, |
| parallelism_config = None, |
| deepspeed = None, |
| label_smoothing_factor = 0.0, |
| optim = 'adamw_8bit', |
| optim_args = None, |
| adafactor = False, |
| group_by_length = False, |
| length_column_name = 'length', |
| report_to = None, |
| ddp_find_unused_parameters = None, |
| ddp_bucket_cap_mb = None, |
| ddp_broadcast_buffers = None, |
| dataloader_pin_memory = True, |
| dataloader_persistent_workers = False, |
| skip_memory_metrics = True, |
| use_legacy_prediction_loop = False, |
| push_to_hub = False, |
| resume_from_checkpoint = None, |
| hub_model_id = None, |
| hub_strategy = 'every_save', |
| hub_token = None, |
| hub_private_repo = None, |
| hub_always_push = False, |
| hub_revision = None, |
| gradient_checkpointing = True, |
| gradient_checkpointing_kwargs = None, |
| include_inputs_for_metrics = False, |
| eval_do_concat_batches = True, |
| fp16_backend = 'auto', |
| push_to_hub_model_id = None, |
| push_to_hub_organization = None, |
| push_to_hub_token = None, |
| mp_parameters = '', |
| auto_find_batch_size = False, |
| full_determinism = False, |
| torchdynamo = None, |
| ray_scope = 'last', |
| ddp_timeout = 1800, |
| torch_compile = False, |
| torch_compile_backend = None, |
| torch_compile_mode = None, |
| include_tokens_per_second = False, |
| include_num_input_tokens_seen = False, |
| neftune_noise_alpha = None, |
| optim_target_modules = None, |
| batch_eval_metrics = False, |
| eval_on_start = False, |
| use_liger_kernel = False, |
| liger_kernel_config = None, |
| eval_use_gather_object = False, |
| average_tokens_across_devices = True, |
| dataset_num_proc = None, |
| num_mini_batches = 1, |
| total_episodes = None, |
| local_rollout_forward_batch_size = 64, |
| num_sample_generations = 10, |
| response_length = 53, |
| stop_token = None, |
| stop_token_id = None, |
| temperature = 0.7, |
| missing_eos_penalty = None, |
| sft_model_path = 'EleutherAI/pythia-160m', |
| world_size = None, |
| num_total_batches = None, |
| micro_batch_size = None, |
| local_batch_size = None, |
| batch_size = None, |
| local_mini_batch_size = None, |
| mini_batch_size = None, |
| exp_name = 'ppo_config', |
| reward_model_path = 'EleutherAI/pythia-160m', |
| model_adapter_name = None, |
| ref_adapter_name = None, |
| num_ppo_epochs = 4, |
| whiten_rewards = False, |
| kl_coef = 0.05, |
| kl_estimator = 'k1', |
| cliprange = 0.2, |
| vf_coef = 0.1, |
| cliprange_value = 0.2, |
| gamma = 1.0, |
| lam = 0.95, |
| ds3_gather_for_generation = True, |
| vllm_sampling_params = None, |
| unsloth_num_chunks = -1, |
| |
| **kwargs, |
| ): |
| if learning_rate < 1e-7: print(f'Unsloth: Your learning rate of `{learning_rate}` is too small and less than 1e-7! Consider increasing it, otherwise gradient updates will be close to 0!') |
| if learning_rate > 1: print(f'Unsloth: Your learning rate of `{learning_rate}` is way too larger > 1! Consider decreasing it to 1e-1, otherwise gradient updates will explode!') |
| if output_dir is None and save_strategy == 'steps' and save_steps == 500: |
| output_dir = 'unsloth_training_checkpoints' |
| save_strategy = 'no' |
| if dataset_num_proc is None: |
| from multiprocessing import cpu_count |
| dataset_num_proc = max(cpu_count()+4, 2) |
| if temperature <= 0: |
| raise MathError('Unsloth: Please set a positive non-zero temperature since your results will be wrong.') |
| elif temperature >= 10: |
| raise MathError('Unsloth: Please set a positive non-zero temperature less than 10, since sampling will be quite erratic.') |
| |
| |
| super().__init__( |
| output_dir = output_dir, |
| overwrite_output_dir = overwrite_output_dir, |
| do_train = do_train, |
| do_eval = do_eval, |
| do_predict = do_predict, |
| eval_strategy = eval_strategy, |
| prediction_loss_only = prediction_loss_only, |
| per_device_train_batch_size = per_device_train_batch_size, |
| per_device_eval_batch_size = per_device_eval_batch_size, |
| per_gpu_train_batch_size = per_gpu_train_batch_size, |
| per_gpu_eval_batch_size = per_gpu_eval_batch_size, |
| gradient_accumulation_steps = gradient_accumulation_steps, |
| eval_accumulation_steps = eval_accumulation_steps, |
| eval_delay = eval_delay, |
| torch_empty_cache_steps = torch_empty_cache_steps, |
| learning_rate = learning_rate, |
| weight_decay = weight_decay, |
| adam_beta1 = adam_beta1, |
| adam_beta2 = adam_beta2, |
| adam_epsilon = adam_epsilon, |
| max_grad_norm = max_grad_norm, |
| num_train_epochs = num_train_epochs, |
| max_steps = max_steps, |
| lr_scheduler_type = lr_scheduler_type, |
| warmup_ratio = warmup_ratio, |
| warmup_steps = warmup_steps, |
| log_level = log_level, |
| log_level_replica = log_level_replica, |
| log_on_each_node = log_on_each_node, |
| logging_dir = logging_dir, |
| logging_strategy = logging_strategy, |
| logging_first_step = logging_first_step, |
| logging_steps = logging_steps, |
| logging_nan_inf_filter = logging_nan_inf_filter, |
| save_strategy = save_strategy, |
| save_steps = save_steps, |
| save_total_limit = save_total_limit, |
| save_safetensors = save_safetensors, |
| save_on_each_node = save_on_each_node, |
| save_only_model = save_only_model, |
| restore_callback_states_from_checkpoint = restore_callback_states_from_checkpoint, |
| no_cuda = no_cuda, |
| use_cpu = use_cpu, |
| use_mps_device = use_mps_device, |
| seed = seed, |
| data_seed = data_seed, |
| jit_mode_eval = jit_mode_eval, |
| use_ipex = use_ipex, |
| bf16 = bf16, |
| fp16 = fp16, |
| fp16_opt_level = fp16_opt_level, |
| half_precision_backend = half_precision_backend, |
| bf16_full_eval = bf16_full_eval, |
| fp16_full_eval = fp16_full_eval, |
| tf32 = tf32, |
| local_rank = local_rank, |
| ddp_backend = ddp_backend, |
| tpu_num_cores = tpu_num_cores, |
| tpu_metrics_debug = tpu_metrics_debug, |
| debug = debug, |
| dataloader_drop_last = dataloader_drop_last, |
| eval_steps = eval_steps, |
| dataloader_num_workers = dataloader_num_workers, |
| dataloader_prefetch_factor = dataloader_prefetch_factor, |
| past_index = past_index, |
| run_name = run_name, |
| disable_tqdm = disable_tqdm, |
| remove_unused_columns = remove_unused_columns, |
| label_names = label_names, |
| load_best_model_at_end = load_best_model_at_end, |
| metric_for_best_model = metric_for_best_model, |
| greater_is_better = greater_is_better, |
| ignore_data_skip = ignore_data_skip, |
| fsdp = fsdp, |
| fsdp_min_num_params = fsdp_min_num_params, |
| fsdp_config = fsdp_config, |
| fsdp_transformer_layer_cls_to_wrap = fsdp_transformer_layer_cls_to_wrap, |
| accelerator_config = accelerator_config, |
| parallelism_config = parallelism_config, |
| deepspeed = deepspeed, |
| label_smoothing_factor = label_smoothing_factor, |
| optim = optim, |
| optim_args = optim_args, |
| adafactor = adafactor, |
| group_by_length = group_by_length, |
| length_column_name = length_column_name, |
| report_to = report_to, |
| ddp_find_unused_parameters = ddp_find_unused_parameters, |
| ddp_bucket_cap_mb = ddp_bucket_cap_mb, |
| ddp_broadcast_buffers = ddp_broadcast_buffers, |
| dataloader_pin_memory = dataloader_pin_memory, |
| dataloader_persistent_workers = dataloader_persistent_workers, |
| skip_memory_metrics = skip_memory_metrics, |
| use_legacy_prediction_loop = use_legacy_prediction_loop, |
| push_to_hub = push_to_hub, |
| resume_from_checkpoint = resume_from_checkpoint, |
| hub_model_id = hub_model_id, |
| hub_strategy = hub_strategy, |
| hub_token = hub_token, |
| hub_private_repo = hub_private_repo, |
| hub_always_push = hub_always_push, |
| hub_revision = hub_revision, |
| gradient_checkpointing = gradient_checkpointing, |
| gradient_checkpointing_kwargs = gradient_checkpointing_kwargs, |
| include_inputs_for_metrics = include_inputs_for_metrics, |
| eval_do_concat_batches = eval_do_concat_batches, |
| fp16_backend = fp16_backend, |
| push_to_hub_model_id = push_to_hub_model_id, |
| push_to_hub_organization = push_to_hub_organization, |
| push_to_hub_token = push_to_hub_token, |
| mp_parameters = mp_parameters, |
| auto_find_batch_size = auto_find_batch_size, |
| full_determinism = full_determinism, |
| torchdynamo = torchdynamo, |
| ray_scope = ray_scope, |
| ddp_timeout = ddp_timeout, |
| torch_compile = torch_compile, |
| torch_compile_backend = torch_compile_backend, |
| torch_compile_mode = torch_compile_mode, |
| include_tokens_per_second = include_tokens_per_second, |
| include_num_input_tokens_seen = include_num_input_tokens_seen, |
| neftune_noise_alpha = neftune_noise_alpha, |
| optim_target_modules = optim_target_modules, |
| batch_eval_metrics = batch_eval_metrics, |
| eval_on_start = eval_on_start, |
| use_liger_kernel = use_liger_kernel, |
| liger_kernel_config = liger_kernel_config, |
| eval_use_gather_object = eval_use_gather_object, |
| average_tokens_across_devices = average_tokens_across_devices, |
| dataset_num_proc = dataset_num_proc, |
| num_mini_batches = num_mini_batches, |
| total_episodes = total_episodes, |
| local_rollout_forward_batch_size = local_rollout_forward_batch_size, |
| num_sample_generations = num_sample_generations, |
| response_length = response_length, |
| stop_token = stop_token, |
| stop_token_id = stop_token_id, |
| temperature = temperature, |
| missing_eos_penalty = missing_eos_penalty, |
| sft_model_path = sft_model_path, |
| world_size = world_size, |
| num_total_batches = num_total_batches, |
| micro_batch_size = micro_batch_size, |
| local_batch_size = local_batch_size, |
| batch_size = batch_size, |
| local_mini_batch_size = local_mini_batch_size, |
| mini_batch_size = mini_batch_size, |
| exp_name = exp_name, |
| reward_model_path = reward_model_path, |
| model_adapter_name = model_adapter_name, |
| ref_adapter_name = ref_adapter_name, |
| num_ppo_epochs = num_ppo_epochs, |
| whiten_rewards = whiten_rewards, |
| kl_coef = kl_coef, |
| kl_estimator = kl_estimator, |
| cliprange = cliprange, |
| vf_coef = vf_coef, |
| cliprange_value = cliprange_value, |
| gamma = gamma, |
| lam = lam, |
| ds3_gather_for_generation = ds3_gather_for_generation,**kwargs) |
| self.vllm_sampling_params = vllm_sampling_params |
| self.unsloth_num_chunks = unsloth_num_chunks |
| |
| pass |
|
|
| class _UnslothPPOTrainer(Trainer): |
| _tag_names = ["trl", "ppo"] |
|
|
| def __init__( |
| self, |
| args: PPOConfig, |
| processing_class: Optional[ |
| Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin] |
| ], |
| model: nn.Module, |
| ref_model: Optional[nn.Module], |
| reward_model: nn.Module, |
| train_dataset: Dataset, |
| value_model: nn.Module, |
| data_collator: Optional[DataCollatorWithPadding] = None, |
| eval_dataset: Optional[Union[Dataset, dict[str, Dataset]]] = None, |
| |
| optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), |
| callbacks: Optional[list[TrainerCallback]] = None, |
| peft_config: Optional["PeftConfig"] = None, |
| ) -> None: |
| if ref_model is model: |
| raise ValueError( |
| "`model` and `ref_model` cannot be the same object. If you want `ref_model` to be the " |
| "same as `model`, you must make a copy of it, or `None` if you use peft." |
| ) |
|
|
| self.args = args |
| self.processing_class = processing_class |
| self.policy_model = model |
|
|
| |
| if data_collator is None: |
| data_collator = DataCollatorWithPadding(self.processing_class) |
|
|
| |
| if args.stop_token and args.stop_token_id: |
| raise ValueError("You cannot set both `stop_token` and `stop_token_id`.") |
| elif args.stop_token: |
| if args.stop_token == "eos": |
| self.policy_model.generation_config.eos_token_id = self.stop_token_id = processing_class.eos_token_id |
| else: |
| raise ValueError( |
| f"Unknown `stop_token` {args.stop_token}. Allowed values are: `'eos'` and `None` (no stop token)." |
| ) |
| else: |
| self.policy_model.generation_config.eos_token_id = self.stop_token_id = args.stop_token_id |
|
|
| |
| if self.args.kl_estimator not in {"k1", "k3"}: |
| raise ValueError( |
| "kl_estimator must be either 'k1' (straightforward, unbiased) or 'k3' (lower variance, unbiased, " |
| "appears to be a strictly better estimator). See " |
| "[Approximating KL Divergence](http://joschu.net/blog/kl-approx.html) for details." |
| ) |
|
|
| |
| if not is_peft_available() and peft_config is not None: |
| raise ImportError( |
| "PEFT is not installed and you passed a `peft_config` in the trainer's kwargs, please install it to use the PEFT models" |
| ) |
| elif is_peft_available() and peft_config is not None: |
| |
| if isinstance(self.policy_model, PeftModel): |
| self.policy_model = self.policy_model.merge_and_unload() |
|
|
| |
| self.policy_model = get_peft_model(self.policy_model, peft_config) |
| if args.bf16 and getattr(self.policy_model, "is_loaded_in_4bit", False): |
| peft_module_casting_to_bf16(self.policy_model) |
|
|
| self.is_peft_model = is_peft_available() and isinstance(self.policy_model, PeftModel) |
| self.model_adapter_name = args.model_adapter_name |
| self.ref_adapter_name = args.ref_adapter_name |
|
|
| if ref_model: |
| self.ref_model = ref_model |
| elif self.is_peft_model: |
| self.ref_model = None |
| else: |
| self.ref_model = create_reference_model(self.policy_model) |
|
|
| self.reward_model = reward_model |
| self.train_dataset = train_dataset |
| self.train_dataset_len = len(train_dataset) |
| self.value_model = value_model |
| self.data_collator = data_collator |
| self.eval_dataset = eval_dataset |
| self.optimizer, self.lr_scheduler = optimizers |
| self.optimizer_cls_and_kwargs = None |
|
|
| |
| |
| |
| if args.total_episodes is None: |
| args.total_episodes = int(args.num_train_epochs * self.train_dataset_len) |
| accelerator = Accelerator(gradient_accumulation_steps=args.gradient_accumulation_steps) |
| self.accelerator = accelerator |
| args.world_size = accelerator.num_processes |
| args.local_batch_size = args.per_device_train_batch_size * args.gradient_accumulation_steps |
| args.micro_batch_size = int(args.per_device_train_batch_size * args.world_size) |
| args.batch_size = int(args.local_batch_size * args.world_size) |
| args.mini_batch_size = exact_div( |
| args.batch_size, args.num_mini_batches, "`batch_size` must be a multiple of `num_mini_batches`" |
| ) |
| args.local_mini_batch_size = exact_div( |
| args.local_batch_size, args.num_mini_batches, "`local_batch_size` must be a multiple of `num_mini_batches`" |
| ) |
| if args.whiten_rewards: |
| assert args.local_mini_batch_size >= 8, ( |
| f"Per-rank minibatch size {args.local_mini_batch_size} is insufficient for whitening" |
| ) |
| |
| |
| args.num_total_batches = math.ceil( |
| args.total_episodes / args.batch_size |
| ) |
| time_tensor = torch.tensor(int(time.time()), device=accelerator.device) |
| time_int = broadcast(time_tensor, 0).item() |
| args.run_name = f"{args.exp_name}__{args.seed}__{time_int}" |
| self.local_seed = args.seed + accelerator.process_index * 100003 |
| if args.num_sample_generations > 0: |
| self.sample_generations_freq = max(1, args.num_total_batches // args.num_sample_generations) |
| self.local_dataloader_batch_size = args.local_batch_size |
|
|
| |
| |
| |
| for module in [self.policy_model, self.ref_model, self.value_model, self.reward_model]: |
| if module is not None: |
| disable_dropout_in_model(module) |
| self.model = PolicyAndValueWrapper(self.policy_model, self.value_model) |
| self.model.config = self.policy_model.config |
| self.create_optimizer_and_scheduler( |
| num_training_steps=args.num_total_batches |
| ) |
|
|
| |
| |
| |
| default_callbacks = DEFAULT_CALLBACKS + get_reporting_integration_callbacks(self.args.report_to) |
| self.callbacks = default_callbacks if callbacks is None else default_callbacks + callbacks |
| self.callback_handler = CallbackHandler( |
| self.callbacks, self.model, self.processing_class, self.optimizer, self.lr_scheduler |
| ) |
| self.add_callback(PrinterCallback if self.args.disable_tqdm else DEFAULT_PROGRESS_CALLBACK) |
| self.control = TrainerControl() |
| self.state = OnlineTrainerState( |
| is_local_process_zero=self.is_local_process_zero(), |
| is_world_process_zero=self.is_world_process_zero(), |
| stateful_callbacks=[ |
| cb for cb in self.callback_handler.callbacks + [self.control] if isinstance(cb, ExportableState) |
| ], |
| ) |
| self.current_flos = 0 |
| self.hp_search_backend = None |
| self.is_deepspeed_enabled = getattr(self.accelerator.state, "deepspeed_plugin", None) is not None |
| self.is_fsdp_enabled = getattr(self.accelerator.state, "fsdp_plugin", None) is not None |
| |
| self.hub_model_id = None |
| if self.args.push_to_hub: |
| self.init_hf_repo() |
| if self.args.should_save: |
| os.makedirs(self.args.output_dir, exist_ok=True) |
|
|
| |
| if hasattr(self.model, "add_model_tags"): |
| self.model.add_model_tags(self._tag_names) |
|
|
| |
| |
| |
| self.dataloader = DataLoader( |
| self.train_dataset, |
| batch_size=self.local_dataloader_batch_size, |
| shuffle=True, |
| collate_fn=self.data_collator, |
| drop_last=True, |
| ) |
| |
| |
| torch.manual_seed(args.seed) |
| self.model, self.optimizer, self.dataloader = accelerator.prepare(self.model, self.optimizer, self.dataloader) |
| torch.manual_seed(self.local_seed) |
|
|
| self.eval_dataloader = DataLoader( |
| self.eval_dataset, |
| batch_size=args.per_device_eval_batch_size, |
| collate_fn=self.data_collator, |
| drop_last=True, |
| ) |
| self.eval_dataloader = accelerator.prepare(self.eval_dataloader) |
|
|
| if self.is_deepspeed_enabled: |
| self.reward_model = prepare_deepspeed( |
| self.reward_model, args.per_device_train_batch_size, args.fp16, args.bf16 |
| ) |
|
|
| if self.ref_model is None: |
| if not self.is_peft_model: |
| raise ValueError("No reference model and model is not a Peft model.") |
| else: |
| self.ref_model = prepare_deepspeed( |
| self.ref_model, args.per_device_train_batch_size, args.fp16, args.bf16 |
| ) |
| else: |
| if self.ref_model is None: |
| if not self.is_peft_model: |
| raise ValueError("No reference model and model is not a Peft model.") |
| else: |
| self.ref_model = self.ref_model.to(self.accelerator.device) |
| self.reward_model = self.reward_model.to(self.accelerator.device) |
|
|
| def get_train_dataloader(self) -> DataLoader: |
| return self.dataloader |
|
|
| def get_eval_dataloader(self) -> DataLoader: |
| return self.eval_dataloader |
|
|
| @contextmanager |
| def null_ref_context(self): |
| """Context manager for handling null reference model (that is, peft adapter manipulation).""" |
| with ( |
| self.accelerator.unwrap_model(self.model.policy).disable_adapter() |
| if self.is_peft_model and not self.ref_adapter_name |
| else nullcontext() |
| ): |
| if self.ref_adapter_name: |
| self.model.policy.set_adapter(self.ref_adapter_name) |
| yield |
| if self.ref_adapter_name: |
| self.model.policy.set_adapter(self.model_adapter_name or "default") |
|
|
| def save_model(self, output_dir: Optional[str] = None, _internal_call: bool = False): |
| backup_model = self.model |
| self.model = self.model.policy |
|
|
| if self.is_deepspeed_enabled: |
| backup_deepspeed = self.deepspeed |
| self.deepspeed = self.model |
|
|
| super().save_model(output_dir, _internal_call) |
|
|
| self.model = backup_model |
|
|
| if self.is_deepspeed_enabled: |
| self.deepspeed = backup_deepspeed |
|
|
| def train(self): |
| args = self.args |
| accelerator = self.accelerator |
| optimizer = self.optimizer |
| model = self.model |
| ref_policy = self.ref_model |
| reward_model = self.reward_model |
| processing_class = self.processing_class |
| dataloader = self.dataloader |
| device = accelerator.device |
|
|
| def repeat_generator(): |
| while True: |
| yield from dataloader |
|
|
| iter_dataloader = iter(repeat_generator()) |
| generation_config = GenerationConfig( |
| max_new_tokens=args.response_length, |
| temperature=(args.temperature + 1e-7), |
| top_k=0.0, |
| top_p=1.0, |
| do_sample=True, |
| ) |
|
|
| accelerator.print("===training policy===") |
| start_time = time.time() |
| stats_shape = (args.num_ppo_epochs, args.num_mini_batches, args.gradient_accumulation_steps) |
| approxkl_stats = torch.zeros(stats_shape, device=device) |
| pg_clipfrac_stats = torch.zeros(stats_shape, device=device) |
| pg_loss_stats = torch.zeros(stats_shape, device=device) |
| vf_loss_stats = torch.zeros(stats_shape, device=device) |
| vf_clipfrac_stats = torch.zeros(stats_shape, device=device) |
| entropy_stats = torch.zeros(stats_shape, device=device) |
| ratio_stats = torch.zeros(stats_shape, device=device) |
| model.train() |
|
|
| |
| self.state.global_step = 0 |
| self.state.episode = 0 |
| self.state.max_steps = args.num_total_batches |
| self.state.num_train_epochs = args.total_episodes / self.train_dataset_len |
| |
| if args.logging_steps is not None: |
| if args.logging_steps < 1: |
| self.state.logging_steps = math.ceil(self.state.max_steps * args.logging_steps) |
| else: |
| self.state.logging_steps = args.logging_steps |
| if args.eval_steps is not None: |
| if args.eval_steps < 1: |
| self.state.eval_steps = math.ceil(self.state.max_steps * args.eval_steps) |
| else: |
| self.state.eval_steps = args.eval_steps |
| if args.save_steps is not None: |
| if args.save_steps < 1: |
| self.state.save_steps = math.ceil(self.state.max_steps * args.save_steps) |
| else: |
| self.state.save_steps = args.save_steps |
| self.control = self.callback_handler.on_train_begin(args, self.state, self.control) |
|
|
| |
| if self.is_deepspeed_enabled: |
| self.deepspeed = self.model |
| self.model_wrapped = self.model |
|
|
| for update in range(1, args.num_total_batches + 1): |
| self.state.episode += 1 * args.batch_size |
| data = next(iter_dataloader) |
| with torch.no_grad(): |
| queries = data["input_ids"].to(device) |
| context_length = queries.shape[1] |
| responses = [] |
| postprocessed_responses = [] |
| logprobs = [] |
| ref_logprobs = [] |
| scores = [] |
| sequence_lengths = [] |
| values = [] |
| with unwrap_model_for_generation( |
| self.model, self.accelerator, gather_deepspeed3_params=self.args.ds3_gather_for_generation |
| ) as unwrapped_model: |
| query_responses, logitss = batch_generation( |
| unwrapped_model.policy, |
| queries, |
| args.local_rollout_forward_batch_size, |
| processing_class.pad_token_id, |
| generation_config, |
| ) |
|
|
| for i in range(0, queries.shape[0], args.local_rollout_forward_batch_size): |
| query = queries[i : i + args.local_rollout_forward_batch_size] |
| query_response = query_responses[i : i + args.local_rollout_forward_batch_size] |
| response = query_response[:, context_length:] |
| logits = logitss[i : i + args.local_rollout_forward_batch_size] |
| logprob = selective_log_softmax(logits, response) |
| del logits |
| empty_cache() |
|
|
| if ref_policy is None: |
| with self.null_ref_context(): |
| ref_output = forward(model.policy, query_response, processing_class.pad_token_id) |
| else: |
| ref_output = forward(ref_policy, query_response, processing_class.pad_token_id) |
| ref_logits = ref_output.logits[:, context_length - 1 : -1] |
| ref_logits /= args.temperature + 1e-7 |
| ref_logprob = selective_log_softmax(ref_logits, response) |
| del ref_output, ref_logits |
| empty_cache() |
|
|
| |
| postprocessed_response = response |
| if self.stop_token_id is not None: |
| postprocessed_response = truncate_response( |
| self.stop_token_id, processing_class.pad_token_id, response |
| ) |
|
|
| |
| postprocessed_query_response = torch.cat((query, postprocessed_response), 1) |
| sequence_length = first_true_indices(postprocessed_response == processing_class.pad_token_id) - 1 |
| unwrapped_value_model = accelerator.unwrap_model(model).value_model |
| full_value, _, _ = get_reward( |
| unwrapped_value_model, query_response, processing_class.pad_token_id, context_length |
| ) |
| value = full_value[:, context_length - 1 : -1].squeeze(-1) |
| _, score, _ = get_reward( |
| reward_model, postprocessed_query_response, processing_class.pad_token_id, context_length |
| ) |
|
|
| responses.append(response) |
| postprocessed_responses.append(postprocessed_response) |
| logprobs.append(logprob) |
| ref_logprobs.append(ref_logprob) |
| sequence_lengths.append(sequence_length) |
| scores.append(score) |
| values.append(value) |
| responses = torch.cat(responses, 0) |
| postprocessed_responses = torch.cat(postprocessed_responses, 0) |
| logprobs = torch.cat(logprobs, 0) |
| ref_logprobs = torch.cat(ref_logprobs, 0) |
| sequence_lengths = torch.cat(sequence_lengths, 0) |
| scores = torch.cat(scores, 0) |
| values = torch.cat(values, 0) |
| del (logprob, ref_logprob, full_value, value, score, unwrapped_model) |
| empty_cache() |
| gc.collect() |
|
|
| |
| |
| contain_eos_token = torch.any(postprocessed_responses == self.processing_class.eos_token_id, dim=-1) |
| if self.args.missing_eos_penalty is not None: |
| scores[~contain_eos_token] -= self.args.missing_eos_penalty |
| |
|
|
| |
| response_idxs = torch.arange(responses.shape[1], device=responses.device).repeat(responses.shape[0], 1) |
| padding_mask = response_idxs > sequence_lengths.unsqueeze(1) |
| logprobs = torch.masked_fill(logprobs, padding_mask, INVALID_LOGPROB) |
| ref_logprobs = torch.masked_fill(ref_logprobs, padding_mask, INVALID_LOGPROB) |
| sequence_lengths_p1 = sequence_lengths + 1 |
| padding_mask_p1 = response_idxs > (sequence_lengths_p1.unsqueeze(1)) |
| values = torch.masked_fill(values, padding_mask_p1, 0) |
|
|
| |
| |
| logr = ref_logprobs - logprobs |
| kl = -logr if args.kl_estimator == "k1" else (logr.exp() - 1) - logr |
| non_score_reward = -args.kl_coef * kl |
| rewards = non_score_reward.clone() |
| actual_start = torch.arange(rewards.size(0), device=rewards.device) |
| actual_end = torch.where(sequence_lengths_p1 < rewards.size(1), sequence_lengths_p1, sequence_lengths) |
| rewards[[actual_start, actual_end]] += scores |
|
|
| |
| if args.whiten_rewards: |
| rewards = masked_whiten(rewards, mask=~padding_mask_p1, shift_mean=False) |
| rewards = torch.masked_fill(rewards, padding_mask_p1, 0) |
|
|
| |
| lastgaelam = 0 |
| advantages_reversed = [] |
| gen_length = responses.shape[1] |
| for t in reversed(range(gen_length)): |
| nextvalues = values[:, t + 1] if t < gen_length - 1 else 0.0 |
| delta = rewards[:, t] + args.gamma * nextvalues - values[:, t] |
| lastgaelam = delta + args.gamma * args.lam * lastgaelam |
| advantages_reversed.append(lastgaelam) |
| advantages = torch.stack(advantages_reversed[::-1], axis=1) |
| returns = advantages + values |
| advantages = masked_whiten(advantages, ~padding_mask) |
| advantages = torch.masked_fill(advantages, padding_mask, 0) |
| empty_cache() |
|
|
| |
| for ppo_epoch_idx in range(args.num_ppo_epochs): |
| b_inds = np.random.permutation(args.local_batch_size) |
| minibatch_idx = 0 |
| for mini_batch_start in range(0, args.local_batch_size, args.local_mini_batch_size): |
| mini_batch_end = mini_batch_start + args.local_mini_batch_size |
| mini_batch_inds = b_inds[mini_batch_start:mini_batch_end] |
| gradient_accumulation_idx = 0 |
| for micro_batch_start in range(0, args.local_mini_batch_size, args.per_device_train_batch_size): |
| with accelerator.accumulate(model): |
| micro_batch_end = micro_batch_start + args.per_device_train_batch_size |
| micro_batch_inds = mini_batch_inds[micro_batch_start:micro_batch_end] |
| mb_advantage = advantages[micro_batch_inds] |
| mb_responses = responses[micro_batch_inds] |
| mb_query_responses = query_responses[micro_batch_inds] |
| mb_logprobs = logprobs[micro_batch_inds] |
| mb_return = returns[micro_batch_inds] |
| mb_values = values[micro_batch_inds] |
|
|
| output, vpred_temp = forward(model, mb_query_responses, processing_class.pad_token_id) |
| logits = output.logits[:, context_length - 1 : -1] |
| logits /= args.temperature + 1e-7 |
| new_logprobs = selective_log_softmax(logits, mb_responses) |
| new_logprobs = torch.masked_fill( |
| new_logprobs, padding_mask[micro_batch_inds], INVALID_LOGPROB |
| ) |
| vpred = vpred_temp[:, context_length - 1 : -1].squeeze(-1) |
| vpred = torch.masked_fill(vpred, padding_mask_p1[micro_batch_inds], 0) |
| vpredclipped = torch.clamp( |
| vpred, |
| mb_values - args.cliprange_value, |
| mb_values + args.cliprange_value, |
| ) |
| vf_losses1 = torch.square(vpred - mb_return) |
| vf_losses2 = torch.square(vpredclipped - mb_return) |
| vf_loss_max = torch.max(vf_losses1, vf_losses2) |
| vf_loss = 0.5 * masked_mean(vf_loss_max, ~padding_mask_p1[micro_batch_inds]) |
| vf_clipfrac = masked_mean( |
| (vf_losses2 > vf_losses1).float(), ~padding_mask_p1[micro_batch_inds] |
| ) |
| logprobs_diff = new_logprobs - mb_logprobs |
| ratio = torch.exp(logprobs_diff) |
| pg_losses = -mb_advantage * ratio |
| pg_losses2 = -mb_advantage * torch.clamp(ratio, 1.0 - args.cliprange, 1.0 + args.cliprange) |
| pg_loss_max = torch.max(pg_losses, pg_losses2) |
| pg_loss = masked_mean(pg_loss_max, ~padding_mask[micro_batch_inds]) |
| loss = pg_loss + args.vf_coef * vf_loss |
| accelerator.backward(loss) |
| optimizer.step() |
| optimizer.zero_grad() |
| with torch.no_grad(): |
| pg_clipfrac = masked_mean( |
| (pg_losses2 > pg_losses).float(), ~padding_mask[micro_batch_inds] |
| ) |
| prob_dist = torch.nn.functional.softmax(logits, dim=-1, dtype = torch.float32).to(logits.dtype) |
| entropy = torch.logsumexp(logits, dim=-1) - torch.sum(prob_dist * logits, dim=-1) |
| approxkl = 0.5 * (logprobs_diff**2).mean() |
| approxkl_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = approxkl |
| pg_clipfrac_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = ( |
| pg_clipfrac |
| ) |
| pg_loss_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = pg_loss |
| vf_loss_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = vf_loss |
| vf_clipfrac_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = ( |
| vf_clipfrac |
| ) |
| entropy_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = entropy.mean() |
| ratio_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = ratio.mean() |
| gradient_accumulation_idx += 1 |
| minibatch_idx += 1 |
| |
| |
| del ( |
| output, vpred_temp, logits, new_logprobs, vpred, vpredclipped, |
| vf_losses1, vf_losses2, vf_loss, vf_clipfrac, logprobs_diff, ratio, pg_losses, pg_losses2, pg_loss_max, |
| pg_loss, loss, pg_clipfrac, prob_dist, entropy, approxkl, mb_return, |
| mb_advantage, mb_values, mb_responses, mb_query_responses, mb_logprobs, |
| ) |
| |
| empty_cache() |
| with torch.no_grad(): |
| mean_kl = kl.sum(1).mean() |
| mean_entropy = (-logprobs).sum(1).mean() |
| mean_non_score_reward = non_score_reward.sum(1).mean() |
| rlhf_reward = mean_non_score_reward + scores.mean() |
| eps = int(self.state.episode / (time.time() - start_time)) |
| metrics = {} |
| metrics["eps"] = eps |
| metrics["objective/kl"] = self.accelerator.gather_for_metrics(mean_kl).mean().item() |
| metrics["objective/entropy"] = self.accelerator.gather_for_metrics(mean_entropy).mean().item() |
| metrics["objective/non_score_reward"] = ( |
| self.accelerator.gather_for_metrics(mean_non_score_reward).mean().item() |
| ) |
| metrics["objective/rlhf_reward"] = self.accelerator.gather_for_metrics(rlhf_reward).mean().item() |
| metrics["objective/scores"] = self.accelerator.gather_for_metrics(scores.mean()).mean().item() |
| metrics["policy/approxkl_avg"] = self.accelerator.gather_for_metrics(approxkl_stats).mean().item() |
| metrics["policy/clipfrac_avg"] = self.accelerator.gather_for_metrics(pg_clipfrac_stats).mean().item() |
| metrics["loss/policy_avg"] = self.accelerator.gather_for_metrics(pg_loss_stats).mean().item() |
| metrics["loss/value_avg"] = self.accelerator.gather_for_metrics(vf_loss_stats).mean().item() |
| metrics["val/clipfrac_avg"] = self.accelerator.gather_for_metrics(vf_clipfrac_stats).mean().item() |
| metrics["policy/entropy_avg"] = self.accelerator.gather_for_metrics(entropy_stats).mean().item() |
| metrics["val/ratio"] = self.accelerator.gather_for_metrics(ratio_stats).mean().item() |
| metrics["val/ratio_var"] = self.accelerator.gather_for_metrics(ratio_stats).var().item() |
| metrics["val/num_eos_tokens"] = (responses == processing_class.eos_token_id).sum().item() |
| metrics["lr"] = self.lr_scheduler.get_last_lr()[0] |
| metrics["episode"] = self.state.episode |
| self.state.epoch = self.state.episode / self.train_dataset_len |
| self.state.global_step += 1 |
| self.log(metrics) |
|
|
| self.lr_scheduler.step() |
| self.control = self.callback_handler.on_step_end(args, self.state, self.control) |
| if self.control.should_save: |
| self._save_checkpoint(model, trial=None) |
| self.control = self.callback_handler.on_save(self.args, self.state, self.control) |
| del kl, mean_kl, mean_entropy, mean_non_score_reward, scores, metrics, non_score_reward |
| empty_cache() |
| gc.collect() |
|
|
| if args.num_sample_generations > 0 and (update - 1) % self.sample_generations_freq == 0: |
| self.generate_completions(sampling=True) |
| empty_cache() |
| del ( |
| query_responses, |
| responses, |
| postprocessed_responses, |
| logprobs, |
| ref_logprobs, |
| values, |
| sequence_lengths, |
| contain_eos_token, |
| sequence_lengths_p1, |
| response_idxs, |
| padding_mask, |
| padding_mask_p1, |
| rewards, |
| actual_start, |
| actual_end, |
| advantages, |
| returns, |
| ) |
| empty_cache() |
|
|
| |
| self.control = self.callback_handler.on_train_end(args, self.state, self.control) |
| if self.control.should_save: |
| self._save_checkpoint(model, trial=None, metrics=None) |
| self.control = self.callback_handler.on_save(self.args, self.state, self.control) |
|
|
| def generate_completions(self, sampling: bool = False): |
| args = self.args |
| processing_class = self.processing_class |
| generation_config = GenerationConfig( |
| max_new_tokens=self.args.response_length, |
| temperature=(0.01 + 1e-7), |
| top_k=0.0, |
| top_p=1.0, |
| do_sample=True, |
| ) |
|
|
| table = defaultdict(list) |
| with unwrap_model_for_generation( |
| self.model, self.accelerator, gather_deepspeed3_params=self.args.ds3_gather_for_generation |
| ) as unwrapped_model: |
| for batch in self.eval_dataloader: |
| query = batch["input_ids"] |
| with torch.no_grad(): |
| context_length = query.shape[1] |
| query_response, _ = batch_generation( |
| unwrapped_model.policy, |
| query, |
| query.shape[0], |
| processing_class.pad_token_id, |
| generation_config, |
| ) |
| response = query_response[:, context_length:] |
| postprocessed_response = response |
| if self.stop_token_id is not None: |
| postprocessed_response = truncate_response( |
| self.stop_token_id, processing_class.pad_token_id, response |
| ) |
| table["query"].extend( |
| gather_object(processing_class.batch_decode(query, skip_special_tokens=True)) |
| ) |
| table["model response"].extend( |
| gather_object(processing_class.batch_decode(postprocessed_response)) |
| ) |
|
|
| postprocessed_query_response = torch.cat((query, postprocessed_response), 1) |
| _, score, _ = get_reward( |
| self.reward_model, postprocessed_query_response, processing_class.pad_token_id, context_length |
| ) |
| table["score"].extend(self.accelerator.gather_for_metrics(score).float().cpu().numpy()) |
|
|
| if sampling: |
| break |
| df = pd.DataFrame(table) |
|
|
| if self.accelerator.is_main_process: |
| if is_rich_available(): |
| print_rich_table(df.iloc[0 : 0 + 5]) |
| if "wandb" in args.report_to: |
| import wandb |
|
|
| if wandb.run is not None: |
| wandb.log({"completions": wandb.Table(dataframe=df)}) |
|
|
| if "comet_ml" in args.report_to: |
| log_table_to_comet_experiment( |
| name="completions.csv", |
| table=df, |
| ) |
|
|
| |
| def _save_checkpoint(self, model, trial): |
| if self.args.hub_model_id is None: |
| model_name = Path(self.args.output_dir).name |
| else: |
| model_name = self.args.hub_model_id.split("/")[-1] |
| self.create_model_card(model_name=model_name) |
| super()._save_checkpoint(model, trial) |
|
|
| def create_model_card( |
| self, |
| model_name: Optional[str] = None, |
| dataset_name: Optional[str] = None, |
| tags: Union[str, list[str], None] = None, |
| ): |
| """ |
| Creates a draft of a model card using the information available to the `Trainer`. |
| |
| Args: |
| model_name (`str` or `None`, *optional*, defaults to `None`): |
| Name of the model. |
| dataset_name (`str` or `None`, *optional*, defaults to `None`): |
| Name of the dataset used for training. |
| tags (`str`, `list[str]` or `None`, *optional*, defaults to `None`): |
| Tags to be associated with the model card. |
| """ |
| if not self.is_world_process_zero(): |
| return |
|
|
| if hasattr(self.model.config, "_name_or_path") and not os.path.isdir(self.model.config._name_or_path): |
| base_model = self.model.config._name_or_path |
| else: |
| base_model = None |
|
|
| |
| if tags is None: |
| tags = set() |
| elif isinstance(tags, str): |
| tags = {tags} |
| else: |
| tags = set(tags) |
|
|
| if hasattr(self.model.config, "unsloth_version"): |
| tags.add("unsloth") |
|
|
| if "JOB_ID" in os.environ: |
| tags.add("hf_jobs") |
|
|
| tags.update(self._tag_names) |
|
|
| |
| citation = textwrap.dedent("""\ |
| @article{mziegler2019fine-tuning, |
| title = {{Fine-Tuning Language Models from Human Preferences}}, |
| author = {Daniel M. Ziegler and Nisan Stiennon and Jeffrey Wu and Tom B. Brown and Alec Radford and Dario Amodei and Paul F. Christiano and Geoffrey Irving}, |
| year = 2019, |
| eprint = {arXiv:1909.08593} |
| }""") |
|
|
| model_card = generate_model_card( |
| base_model=base_model, |
| model_name=model_name, |
| hub_model_id=self.hub_model_id, |
| dataset_name=dataset_name, |
| tags=tags, |
| wandb_url=wandb.run.url if is_wandb_available() and wandb.run is not None else None, |
| comet_url=get_comet_experiment_url(), |
| trainer_name="PPO", |
| trainer_citation=citation, |
| paper_title="Fine-Tuning Language Models from Human Preferences", |
| paper_id="1909.08593", |
| ) |
|
|
| model_card.save(os.path.join(self.args.output_dir, "README.md")) |
| class UnslothPPOTrainer(_UnslothPPOTrainer): |
| """ |
| |
| """ |
| def __init__( |
| self, |
| args, |
| processing_class, |
| model, |
| ref_model, |
| reward_model, |
| train_dataset, |
| value_model, |
| data_collator = None, |
| eval_dataset = None, |
| callbacks = None, |
| peft_config = None, |
| **kwargs |
| ): |
| if args is None: args = UnslothPPOConfig() |
| use_bf16 = getattr(args, 'bf16', False) |
| if type(use_bf16) is not bool: use_bf16 = False |
| use_fp16 = getattr(args, 'fp16', False) |
| if type(use_fp16) is not bool: use_fp16 = False |
| force_float32 = False |
| full_finetuning = os.environ.get('UNSLOTH_ENABLE_FULL_FINETUNING', '0') == '1' |
| if not full_finetuning and (os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1'): |
| print('Unsloth: Switching to float32 training since model cannot work with float16') |
| force_float32 = True |
| mixed_precision_dtype = os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') |
| dtype = getattr(model.config, 'dtype', None) or getattr(model.config, 'torch_dtype', None) |
| if dtype is None: dtype = model.get_input_embeddings().dtype |
| from unsloth_zoo.utils import _get_dtype |
| dtype = _get_dtype(dtype) |
| float16 = dtype == torch.float16 |
| if not force_float32 and (float16 and use_bf16): raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`') |
| if not force_float32 and (not float16 and use_fp16): raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`') |
| if force_float32: |
| |
| args.fp16 = False |
| args.bf16 = False |
| os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' |
| elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32': |
| |
| args.fp16 = float16 |
| args.bf16 = not float16 |
| os.environ['ACCELERATE_MIXED_PRECISION'] = 'fp16' if float16 else 'bf16' |
| if getattr(args, 'eval_dataset', None) is not None and getattr(args, 'eval_strategy', 'no') == 'no': |
| args.eval_strategy = 'steps' |
| if getattr(args, 'eval_steps', None) is None: args.eval_steps = 0.1 |
| ga_steps = getattr(args, 'gradient_accumulation_steps', None) |
| if ga_steps is not None and ga_steps > 1: |
| from transformers import __version__ as transformers_version |
| if Version(transformers_version) <= Version('4.45.2'): |
| print('**** Unsloth: Please use our fixed gradient_accumulation_steps by updating transformers, TRL and Unsloth!\n' |
| '`pip install --upgrade --no-cache-dir --force-reinstall --no-deps unsloth transformers trl unsloth_zoo`') |
| if getattr(args, 'eval_strategy', 'no') != 'no': |
| eval_bsz = getattr(args, 'per_device_eval_batch_size', 8) |
| if eval_bsz == 8 and args.per_device_train_batch_size < eval_bsz: args.per_device_eval_batch_size = args.per_device_train_batch_size |
| if getattr(args, 'eval_accumulation_steps', None) is None and ga_steps is not None: args.eval_accumulation_steps = ga_steps |
| fp16_full_eval = getattr(args, 'fp16_full_eval', False) |
| if type(fp16_full_eval) is not bool: fp16_full_eval = False |
| bf16_full_eval = getattr(args, 'bf16_full_eval', False) |
| if type(bf16_full_eval) is not bool: bf16_full_eval = False |
| if args.fp16 and bf16_full_eval: args.bf16_full_eval = False; args.fp16_full_eval = True |
| if args.bf16 and fp16_full_eval: args.bf16_full_eval = True; args.fp16_full_eval = False |
| if force_float32: |
| args.bf16_full_eval = False |
| args.fp16_full_eval = False |
| elif os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') == 'bfloat16': |
| args.bf16_full_eval = True |
| args.fp16_full_eval = False |
| elif not bf16_full_eval and not fp16_full_eval: |
| args.bf16_full_eval = args.bf16 |
| args.fp16_full_eval = args.fp16 |
| _output_logits = False |
| if locals().get('compute_metrics', None) is not None: _output_logits = True |
| if locals().get('preprocess_logits_for_metrics', None) is not None: _output_logits = True |
| if _output_logits: |
| os.environ['UNSLOTH_RETURN_LOGITS'] = '1' |
| if 'max_seq_length' not in locals() and not hasattr(args, 'max_seq_length'): |
| pass |
| else: |
| model_max_seq_length = getattr(model, 'max_seq_length', None) |
| args_max_seq_length = getattr(args, 'max_seq_length', None) |
| if args_max_seq_length is None and model_max_seq_length is not None: |
| max_seq_length = model.max_seq_length |
| if hasattr(args, 'max_seq_length'): args.max_seq_length = max_seq_length |
| if model is not None and hasattr(model, 'for_training'): |
| model.for_training() |
| if 'tokenizer' in locals() and hasattr(tokenizer, 'padding_side'): tokenizer.padding_side = 'right' |
| if 'processing_class' in locals(): |
| if hasattr(processing_class, 'padding_side'): processing_class.padding_side = 'right' |
| if hasattr(processing_class, 'tokenizer') and hasattr(processing_class.tokenizer, 'padding_side'): processing_class.tokenizer.padding_side = 'right' |
| __tokenizer = processing_class if 'processing_class' in locals() else tokenizer |
| from unsloth_zoo.vision_utils import UnslothVisionDataCollator |
| if not isinstance(data_collator, UnslothVisionDataCollator): |
| if isinstance(data_collator, DataCollatorForSeq2Seq) and 'labels' not in train_dataset.column_names: |
| data_collator = TransformersDataCollatorForLanguageModeling( |
| __tokenizer, |
| mlm = False, |
| mlm_probability = 0.0, |
| pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), |
| ) |
| elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling) and 'labels' in train_dataset.column_names: |
| data_collator = DataCollatorForSeq2Seq( |
| __tokenizer, |
| pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), |
| ) |
| else: |
| if hasattr(args, 'remove_unused_columns'): args.remove_unused_columns = False |
| if hasattr(args, 'dataset_text_field'): args.dataset_text_field = '' |
| if hasattr(args, 'dataset_kwargs'): args.dataset_kwargs = {'skip_prepare_dataset': True} |
| if not isinstance(data_collator, UnslothVisionDataCollator): |
| if not hasattr(__tokenizer, 'pad') and hasattr(__tokenizer, 'tokenizer'): |
| if isinstance(data_collator, DataCollatorForSeq2Seq): |
| data_collator = DataCollatorForSeq2Seq( |
| __tokenizer.tokenizer, |
| pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), |
| ) |
| else: |
| data_collator = TransformersDataCollatorForLanguageModeling( |
| __tokenizer.tokenizer, |
| mlm = False, |
| mlm_probability = 0.0, |
| pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), |
| ) |
| other_metrics = [] |
| |
| from unsloth_zoo.logging_utils import PatchRLStatistics |
| PatchRLStatistics('ppo_trainer', other_metrics) |
| |
| |
| |
| if getattr(args, "parallel_mode", None) == ParallelMode.NOT_DISTRIBUTED and args.n_gpu > 1: |
| if getattr(args, "_n_gpu", 1) != 1: |
| args._n_gpu = 1 |
| if "model" in locals() and hasattr(model, "for_training"): |
| model.for_training() |
| super().__init__( |
| args = args, |
| processing_class = processing_class, |
| model = model, |
| ref_model = ref_model, |
| reward_model = reward_model, |
| train_dataset = train_dataset, |
| value_model = value_model, |
| data_collator = data_collator, |
| eval_dataset = eval_dataset, |
| callbacks = callbacks, |
| peft_config = peft_config,**kwargs) |
| if "model" in locals() and hasattr(model, "for_inference"): |
| model.for_inference() |
| if hasattr(self, 'neftune_hook_handle'): |
| self.neftune_hook_handle.remove() |
| if hasattr(self, 'neftune_hook_handle'): del self.neftune_hook_handle |
| if getattr(args, 'neftune_noise_alpha', None) is not None: |
| model.get_input_embeddings().neftune_noise_alpha = self.neftune_noise_alpha |
| pass |
| if hasattr(self, 'accelerator'): |
| scaler = self.accelerator.scaler |
| current_model = model |
| while hasattr(current_model, 'model'): |
| current_model.accelerator_scaler = scaler |
| current_model = current_model.model |
| current_model.accelerator_scaler = scaler |
| pass |
| if hasattr(self, 'train'): |
| self.train = MethodType(prepare_for_training_mode(self.__class__.train), self) |
| pass |
| |
| pass |
|
|
